Spaces:

hujiecpp
/

PE3R

Running on Zero

hujiecpp commited on Feb 19

Commit

224be48

1 Parent(s): 67b7ca3

init project

Files changed (1) hide show

modules/pe3r/models.py CHANGED Viewed

@@ -36,7 +36,7 @@ class Models:
         SAM1_DECODER_CKP = './checkpoints/Prompt_guided_Mask_Decoder.pt'
         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", device=device)
         image_encoder = sam1.image_encoder
         prompt_encoder, mask_decoder = sam_model_registry['prompt_guided_decoder'](SAM1_DECODER_CKP)

         SAM1_DECODER_CKP = './checkpoints/Prompt_guided_Mask_Decoder.pt'
         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)