Z-Image-Turbo ControlNet

src/diffusers/__init__.py

@@ -278,6 +278,7 @@
             "WanAnimateTransformer3DModel",
             "WanTransformer3DModel",
             "WanVACETransformer3DModel",
+            "ZImageControlNetModel",
             "ZImageTransformer2DModel",
             "attention_backend",
         ]
@@ -666,6 +667,8 @@
             "WuerstchenCombinedPipeline",
             "WuerstchenDecoderPipeline",
             "WuerstchenPriorPipeline",
+            "ZImageControlNetPipeline",
+            "ZImageControlNetInpaintPipeline",
             "ZImageImg2ImgPipeline",
             "ZImagePipeline",
         ]
@@ -1012,6 +1015,7 @@
             WanAnimateTransformer3DModel,
             WanTransformer3DModel,
             WanVACETransformer3DModel,
+            ZImageControlNetModel,
             ZImageTransformer2DModel,
             attention_backend,
         )
@@ -1369,6 +1373,8 @@
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,
             WuerstchenPriorPipeline,
+            ZImageControlNetInpaintPipeline,
+            ZImageControlNetPipeline,
             ZImageImg2ImgPipeline,
             ZImagePipeline,
         )

src/diffusers/loaders/single_file_model.py

@@ -49,10 +49,12 @@
     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,
     create_vae_diffusers_config_from_ldm,
+    create_z_image_controlnet_config,
     fetch_diffusers_config,
     fetch_original_config,
     load_single_file_checkpoint,
@@ -172,11 +174,19 @@
         "checkpoint_mapping_fn": convert_z_image_transformer_checkpoint_to_diffusers,
         "default_subfolder": "transformer",
     },
+    "ZImageControlNetModel": {
+        "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):
@@ -369,6 +379,10 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
             diffusers_model_config = config_mapping_fn(
                 original_config=original_config, checkpoint=checkpoint, **config_mapping_kwargs
             )
+        elif "config_create_fn" in mapping_functions:
+            config_create_fn = mapping_functions["config_create_fn"]
+            config_create_kwargs = _get_mapping_function_kwargs(config_create_fn, **kwargs)
+            diffusers_model_config = config_create_fn(checkpoint=checkpoint, **config_create_kwargs)
         else:
             if config is not None:
                 if isinstance(config, str):

src/diffusers/loaders/single_file_utils.py

@@ -121,6 +121,7 @@
     "instruct-pix2pix": "model.diffusion_model.input_blocks.0.0.weight",
     "lumina2": ["model.diffusion_model.cap_embedder.0.weight", "cap_embedder.0.weight"],
     "z-image-turbo": "cap_embedder.0.weight",
+    "z-image-turbo-controlnet": "control_all_x_embedder.2-1.weight",
     "sana": [
         "blocks.0.cross_attn.q_linear.weight",
         "blocks.0.cross_attn.q_linear.bias",
@@ -779,6 +780,9 @@ def infer_diffusers_model_type(checkpoint):
         else:
             raise ValueError(f"Unexpected x_embedder shape: {x_embedder_shape} when loading Cosmos 2.0 model.")
 
+    elif CHECKPOINT_KEY_NAMES["z-image-turbo-controlnet"] in checkpoint:
+        model_type = "z-image-turbo-controlnet"
+
     else:
         model_type = "v1"
 
@@ -3885,3 +3889,52 @@ def update_state_dict(state_dict: dict[str, object], old_key: str, new_key: str)
             handler_fn_inplace(key, converted_state_dict)
 
     return converted_state_dict
+
+
+def create_z_image_controlnet_config(checkpoint, **kwargs):
+    v1_config = {
+        "control_in_dim": 16,
+        "control_layers_places": [0, 5, 10, 15, 20, 25],
+        "dim": 3840,
+        "n_heads": 30,
+        "n_kv_heads": 30,
+        "n_refiner_layers": 2,
+        "norm_eps": 1e-05,
+        "qk_norm": True,
+        "all_f_patch_size": [1],
+        "all_patch_size": [2],
+    }
+    v2_config = {
+        "control_in_dim": 33,
+        "control_layers_places": [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28],
+        "control_refiner_layers_places": [0, 1],
+        "add_control_noise_refiner": True,
+        "dim": 3840,
+        "n_heads": 30,
+        "n_kv_heads": 30,
+        "n_refiner_layers": 2,
+        "norm_eps": 1e-05,
+        "qk_norm": True,
+        "all_f_patch_size": [1],
+        "all_patch_size": [2],
+    }
+    control_x_embedder_weight_shape = checkpoint["control_all_x_embedder.2-1.weight"].shape[1]
+    if control_x_embedder_weight_shape == 64:
+        return v1_config
+    elif control_x_embedder_weight_shape == 132:
+        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/__init__.py

@@ -66,6 +66,7 @@
     _import_structure["controlnets.controlnet_sparsectrl"] = ["SparseControlNetModel"]
     _import_structure["controlnets.controlnet_union"] = ["ControlNetUnionModel"]
     _import_structure["controlnets.controlnet_xs"] = ["ControlNetXSAdapter", "UNetControlNetXSModel"]
+    _import_structure["controlnets.controlnet_z_image"] = ["ZImageControlNetModel"]
     _import_structure["controlnets.multicontrolnet"] = ["MultiControlNetModel"]
     _import_structure["controlnets.multicontrolnet_union"] = ["MultiControlNetUnionModel"]
     _import_structure["embeddings"] = ["ImageProjection"]
@@ -180,6 +181,7 @@
             SD3MultiControlNetModel,
             SparseControlNetModel,
             UNetControlNetXSModel,
+            ZImageControlNetModel,
         )
         from .embeddings import ImageProjection
         from .modeling_utils import ModelMixin

src/diffusers/models/controlnets/__init__.py

@@ -19,6 +19,7 @@
     )
     from .controlnet_union import ControlNetUnionModel
     from .controlnet_xs import ControlNetXSAdapter, ControlNetXSOutput, UNetControlNetXSModel
+    from .controlnet_z_image import ZImageControlNetModel
     from .multicontrolnet import MultiControlNetModel
     from .multicontrolnet_union import MultiControlNetUnionModel