Z-Image-Turbo ControlNet

src/diffusers/__init__.py

@@ -279,6 +279,7 @@
             "WanTransformer3DModel",
             "WanVACETransformer3DModel",
             "ZImageTransformer2DModel",
+            "ZImageControlNetModel",
             "attention_backend",
         ]
     )
@@ -668,6 +669,7 @@
             "WuerstchenPriorPipeline",
             "ZImageImg2ImgPipeline",
             "ZImagePipeline",
+            "ZImageControlNetPipeline",
         ]
     )
 
@@ -1012,6 +1014,7 @@
             WanAnimateTransformer3DModel,
             WanTransformer3DModel,
             WanVACETransformer3DModel,
+            ZImageControlNetModel,
             ZImageTransformer2DModel,
             attention_backend,
         )
@@ -1369,6 +1372,7 @@
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,
             WuerstchenPriorPipeline,
+            ZImageControlNetPipeline,
             ZImageImg2ImgPipeline,
             ZImagePipeline,
         )

src/diffusers/loaders/single_file_model.py

@@ -53,6 +53,7 @@
     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,6 +173,10 @@
         "checkpoint_mapping_fn": convert_z_image_transformer_checkpoint_to_diffusers,
         "default_subfolder": "transformer",
     },
+    "ZImageControlNetModel": {
+        "checkpoint_mapping_fn": lambda x: x,
+        "config_create_fn": create_z_image_controlnet_config,
+    },
 }
 
 
@@ -369,6 +374,9 @@ 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"]
+            diffusers_model_config = config_create_fn()
         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,18 @@ 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():
+    return {
+        "all_f_patch_size": [1],
+        "all_patch_size": [2],
+        "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,
+    }

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