From e41efcddcdfb06eb546a97641c650607286b08ed Mon Sep 17 00:00:00 2001 From: chflame163 Date: Sun, 27 Oct 2024 10:17:15 +0800 Subject: [PATCH] fix sam2 model load --- py/sam_2_ultra.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/py/sam_2_ultra.py b/py/sam_2_ultra.py index b014f99..1deff2b 100644 --- a/py/sam_2_ultra.py +++ b/py/sam_2_ultra.py @@ -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")