Initial commit

mvdiffusion/models/unet_mv2d_condition.py

@@ -35,7 +35,7 @@ from diffusers.models.embeddings import (
     TimestepEmbedding,
     Timesteps,
 )
-from diffusers.models.modeling_utils import ModelMixin, load_state_dict, _load_state_dict_into_model
+from diffusers.models.modeling_utils import ModelMixin, load_state_dict
 from diffusers.models.unets.unet_2d_blocks import (
     CrossAttnDownBlock2D,
     CrossAttnUpBlock2D,
@@ -1506,6 +1506,7 @@ class UNetMV2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixi
                         del state_dict[checkpoint_key]
             return mismatched_keys
 
+        error_msgs = []
         if state_dict is not None:
             # Whole checkpoint
             mismatched_keys = _find_mismatched_keys(
@@ -1514,7 +1515,11 @@ class UNetMV2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixi
                 original_loaded_keys,
                 ignore_mismatched_sizes,
             )
-            error_msgs = _load_state_dict_into_model(model_to_load, state_dict)
+            # Use PyTorch's load_state_dict with strict=False to handle mismatched keys
+            incompatible_keys = model_to_load.load_state_dict(state_dict, strict=False)
+            if incompatible_keys.missing_keys:
+                error_msgs.append(f"Missing keys: {incompatible_keys.missing_keys}")
+            # Note: unexpected_keys are already tracked separately, so we don't include them here
 
         if len(error_msgs) > 0:
             error_msg = "\n\t".join(error_msgs)