nodes/chflame163-ComfyUI_LayerStyle
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix sam2 model load

Browse Source
This commit is contained in:
chflame163
2024-10-27 10:17:15 +08:00
parent 2df6cf23ed
commit e41efcddcd
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/sam_2_ultra.py
+3 -3
View File
@@ -192,16 +192,16 @@ def load_model(model_path, model_cfg_path, segmentor, dtype, device):
if segmentor == 'single_image':
model_class = SAM2Base
model = initialize_model(model_class, model_config, segmentor, image_encoder, memory_attention, memory_encoder, sam_mask_decoder_extra_args, dtype, device)
model.load_state_dict(sd)
model.load_state_dict(sd, strict=False)
model = SAM2ImagePredictor(model)
elif segmentor == 'video':
model_class = SAM2VideoPredictor
model = initialize_model(model_class, model_config, segmentor, image_encoder, memory_attention, memory_encoder, sam_mask_decoder_extra_args, dtype, device)
model.load_state_dict(sd)
model.load_state_dict(sd, strict=False)
elif segmentor == 'automaskgenerator':
model_class = SAM2Base
model = initialize_model(model_class, model_config, segmentor, image_encoder, memory_attention, memory_encoder, sam_mask_decoder_extra_args, dtype, device)
model.load_state_dict(sd)
model.load_state_dict(sd, strict=False)
model = SAM2AutomaticMaskGenerator(model)
else:
raise ValueError(f"Segmentor {segmentor} not supported")