pop control_noise_refiner from 2.0 state_dict

Author: hlky

src/diffusers/loaders/single_file_model.py

@@ -49,6 +49,7 @@
     convert_stable_cascade_unet_single_file_to_diffusers,
     convert_wan_transformer_to_diffusers,
     convert_wan_vae_to_diffusers,
+    convert_z_image_controlnet_checkpoint_to_diffusers,
     convert_z_image_transformer_checkpoint_to_diffusers,
     create_controlnet_diffusers_config_from_ldm,
     create_unet_diffusers_config_from_ldm,
@@ -174,14 +175,18 @@
         "default_subfolder": "transformer",
     },
     "ZImageControlNetModel": {
-        "checkpoint_mapping_fn": lambda x: x,
+        "checkpoint_mapping_fn": convert_z_image_controlnet_checkpoint_to_diffusers,
         "config_create_fn": create_z_image_controlnet_config,
     },
 }
 
 
 def _should_convert_state_dict_to_diffusers(model_state_dict, checkpoint_state_dict):
-    return not set(model_state_dict.keys()).issubset(set(checkpoint_state_dict.keys()))
+    model_state_dict_keys = set(model_state_dict.keys())
+    checkpoint_state_dict_keys = set(checkpoint_state_dict.keys())
+    is_subset = model_state_dict_keys.issubset(checkpoint_state_dict_keys)
+    is_match = model_state_dict_keys == checkpoint_state_dict_keys
+    return not (is_subset and is_match)
 
 
 def _get_single_file_loadable_mapping_class(cls):

src/diffusers/loaders/single_file_utils.py

@@ -3925,3 +3925,16 @@ def create_z_image_controlnet_config(checkpoint, **kwargs):
         return v2_config
     else:
         raise ValueError("Unknown Z-Image Turbo ControlNet type.")
+
+
+def convert_z_image_controlnet_checkpoint_to_diffusers(checkpoint, **kwargs):
+    control_x_embedder_weight_shape = checkpoint["control_all_x_embedder.2-1.weight"].shape[1]
+    if control_x_embedder_weight_shape == 64:
+        return checkpoint
+    elif control_x_embedder_weight_shape == 132:
+        converted_state_dict = {
+            key: checkpoint.pop(key) for key in list(checkpoint.keys()) if not key.startswith("control_noise_refiner.")
+        }
+        return converted_state_dict
+    else:
+        raise ValueError("Unknown Z-Image Turbo ControlNet type.")

src/diffusers/models/controlnets/controlnet_z_image.py

@@ -432,21 +432,7 @@ def __init__(
 
         self.control_all_x_embedder = nn.ModuleDict(all_x_embedder)
         if self.add_control_noise_refiner:
-            self.control_noise_refiner = nn.ModuleList(
-                [
-                    ZImageControlTransformerBlock(
-                        1000 + layer_id,
-                        dim,
-                        n_heads,
-                        n_kv_heads,
-                        norm_eps,
-                        qk_norm,
-                        modulation=True,
-                        block_id=layer_id,
-                    )
-                    for layer_id in range(n_refiner_layers)
-                ]
-            )
+            self.control_noise_refiner = None
         else:
             self.control_noise_refiner = nn.ModuleList(
                 [