Spaces:

hujiecpp
/

PE3R

Running on Zero

hujiecpp commited on Feb 19

Commit

4addd4f

1 Parent(s): 224be48

init project

Files changed (1) hide show

modules/pe3r/models.py CHANGED Viewed

@@ -37,7 +37,7 @@ class Models:
         self.mobilesamv2 = sam_model_registry['sam_vit_h'](None)
         # image_encoder=sam_model_registry['sam_vit_h_encoder'](SAM1_ENCODER_CKP)
         sam1 = SamModel.from_pretrained('facebook/sam-vit-huge')
-        image_encoder = sam1.image_encoder
         prompt_encoder, mask_decoder = sam_model_registry['prompt_guided_decoder'](SAM1_DECODER_CKP)
         self.mobilesamv2.prompt_encoder = prompt_encoder

         self.mobilesamv2 = sam_model_registry['sam_vit_h'](None)
         # image_encoder=sam_model_registry['sam_vit_h_encoder'](SAM1_ENCODER_CKP)
         sam1 = SamModel.from_pretrained('facebook/sam-vit-huge')
+        image_encoder = sam1.vision_encoder
         prompt_encoder, mask_decoder = sam_model_registry['prompt_guided_decoder'](SAM1_DECODER_CKP)
         self.mobilesamv2.prompt_encoder = prompt_encoder