Detect 2.0 vs 2.1 ZImageControlNetModel

Author: hlky

src/diffusers/loaders/single_file_utils.py

@@ -223,7 +223,8 @@
     "cosmos-2.0-v2w-14B": {"pretrained_model_name_or_path": "nvidia/Cosmos-Predict2-14B-Video2World"},
     "z-image-turbo": {"pretrained_model_name_or_path": "Tongyi-MAI/Z-Image-Turbo"},
     "z-image-turbo-controlnet": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union"},
-    "z-image-turbo-controlnet-2.x": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1"},
+    "z-image-turbo-controlnet-2.0": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0"},
+    "z-image-turbo-controlnet-2.1": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1"},
 }
 
 # Use to configure model sample size when original config is provided
@@ -784,7 +785,10 @@ def infer_diffusers_model_type(checkpoint):
             raise ValueError(f"Unexpected x_embedder shape: {x_embedder_shape} when loading Cosmos 2.0 model.")
 
     elif CHECKPOINT_KEY_NAMES["z-image-turbo-controlnet-2.x"] in checkpoint:
-        model_type = "z-image-turbo-controlnet-2.x"
+        if torch.all(checkpoint["control_noise_refiner.0.before_proj.weight"] == 0.0):
+            model_type = "z-image-turbo-controlnet-2.0"
+        else:
+            model_type = "z-image-turbo-controlnet-2.1"
 
     elif CHECKPOINT_KEY_NAMES["z-image-turbo-controlnet"] in checkpoint:
         model_type = "z-image-turbo-controlnet"

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py

@@ -58,14 +58,13 @@
         >>> #     torch_dtype=torch.bfloat16,
         >>> # )
 
-        >>> # 2.0 - `config` is required
+        >>> # 2.0
         >>> # controlnet = ZImageControlNetModel.from_single_file(
         >>> #     hf_hub_download(
         >>> #         "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0",
         >>> #         filename="Z-Image-Turbo-Fun-Controlnet-Union-2.0.safetensors",
         >>> #     ),
         >>> #     torch_dtype=torch.bfloat16,
-        >>> #     config="hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0",
         >>> # )
 
         >>> pipe = ZImageControlNetPipeline.from_pretrained(

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py

@@ -50,14 +50,13 @@
         ...     torch_dtype=torch.bfloat16,
         ... )
 
-        >>> # 2.0 - `config` is required
+        >>> # 2.0
         >>> # controlnet = ZImageControlNetModel.from_single_file(
         >>> #     hf_hub_download(
         >>> #         "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0",
         >>> #         filename="Z-Image-Turbo-Fun-Controlnet-Union-2.0.safetensors",
         >>> #     ),
         >>> #     torch_dtype=torch.bfloat16,
-        >>> #     config="hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0",
         >>> # )
 
         >>> pipe = ZImageControlNetInpaintPipeline.from_pretrained(