Add Z-Image ControlNet V2.0 support

VRAM usage is high.

- Auto-detect control_in_dim from adapter weights (16 for V1, 33 for V2.0)
- Auto-detect n_refiner_layers from state dict
- Add zero-padding for V2.0's additional channels
- Use accelerate.init_empty_weights() for efficient model creation
- Add ControlNet_Checkpoint_ZImage_Config to frontend schema

Author: Pfannkuchensack

invokeai/app/invocations/z_image_denoise.py

@@ -333,57 +333,74 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
             begin_step_percent = 0.0
             end_step_percent = 1.0
             cached_weights = None
+            control_in_dim = 16  # Default, will be set from adapter config if control is used
 
             if self.control is not None:
-                # Simplified approach: Create model and load weights directly
-                # 1. Get base transformer config (stays on CPU)
-                # 2. Create ZImageControlTransformer2DModel
-                # 3. Load base weights directly (strict=False)
-                # 4. Load adapter control weights on top
-                # 5. Move to GPU
-
-                # Get base transformer config (stays on CPU)
+                # Load base transformer config (NOT to GPU yet - just get the model reference)
                 base_transformer = cast(ZImageTransformer2DModel, transformer_info.model)
                 base_config = base_transformer.config
 
-                # Load control adapter (stays on CPU)
+                # Load control adapter
                 control_model_info = context.models.load(self.control.control_model)
                 control_adapter = control_model_info.model
                 assert isinstance(control_adapter, ZImageControlAdapter)
 
-                # Create ZImageControlTransformer2DModel
-                control_transformer = ZImageControlTransformer2DModel(
-                    all_patch_size=base_config.all_patch_size,
-                    all_f_patch_size=base_config.all_f_patch_size,
-                    in_channels=base_config.in_channels,
-                    dim=base_config.dim,
-                    n_layers=base_config.n_layers,
-                    n_refiner_layers=base_config.n_refiner_layers,
-                    n_heads=base_config.n_heads,
-                    n_kv_heads=base_config.n_kv_heads,
-                    norm_eps=base_config.norm_eps,
-                    qk_norm=base_config.qk_norm,
-                    cap_feat_dim=base_config.cap_feat_dim,
-                    rope_theta=base_config.rope_theta,
-                    t_scale=base_config.t_scale,
-                    axes_dims=base_config.axes_dims,
-                    axes_lens=base_config.axes_lens,
+                # Get control_in_dim from adapter config (16 for V1, 33 for V2.0)
+                adapter_config = control_adapter.config
+                control_in_dim = adapter_config.get("control_in_dim", 16)
+                num_control_blocks = adapter_config.get("num_control_blocks", 6)
+                n_refiner_layers = adapter_config.get("n_refiner_layers", 2)
+
+                # Calculate control_layers_places based on num_control_blocks
+                control_layers_places = [i * 2 for i in range(num_control_blocks)]
+
+                # Log control configuration for debugging
+                version = "V2.0" if control_in_dim > 16 else "V1"
+                context.util.signal_progress(
+                    f"Using Z-Image ControlNet {version}: control_in_dim={control_in_dim}, "
+                    f"num_blocks={num_control_blocks}, scale={self.control.control_context_scale}"
                 )
 
-                # Load base transformer weights directly (strict=False handles missing control keys)
+                # Create control transformer structure with empty weights
+                import accelerate
+
+                with accelerate.init_empty_weights():
+                    control_transformer = ZImageControlTransformer2DModel(
+                        control_layers_places=control_layers_places,
+                        control_in_dim=control_in_dim,
+                        all_patch_size=base_config.all_patch_size,
+                        all_f_patch_size=base_config.all_f_patch_size,
+                        in_channels=base_config.in_channels,
+                        dim=base_config.dim,
+                        n_layers=base_config.n_layers,
+                        n_refiner_layers=n_refiner_layers,
+                        n_heads=base_config.n_heads,
+                        n_kv_heads=base_config.n_kv_heads,
+                        norm_eps=base_config.norm_eps,
+                        qk_norm=base_config.qk_norm,
+                        cap_feat_dim=base_config.cap_feat_dim,
+                        rope_theta=base_config.rope_theta,
+                        t_scale=base_config.t_scale,
+                        axes_dims=base_config.axes_dims,
+                        axes_lens=base_config.axes_lens,
+                    )
+
+                # Load base weights with assign=True (assigns tensors directly, no copy of data)
                 control_transformer.load_state_dict(base_transformer.state_dict(), strict=False, assign=True)
 
-                # Load control adapter weights on top (only control-specific keys)
-                # Filter to only control_ prefixed keys to avoid overwriting x_pad_token
-                adapter_control_weights = {
-                    k: v for k, v in control_adapter.state_dict().items() if k.startswith("control_")
-                }
-                control_transformer.load_state_dict(adapter_control_weights, strict=False, assign=True)
+                # Load control adapter weights on top
+                adapter_weights = {k: v for k, v in control_adapter.state_dict().items() if k.startswith("control_")}
+                control_transformer.load_state_dict(adapter_weights, strict=False, assign=True)
 
-                # Move to device
+                # Move combined model to device
                 control_transformer = control_transformer.to(device=device, dtype=inference_dtype)
                 active_transformer = control_transformer
 
+                # Clean up to save memory # need to check
+                #del control_adapter
+                #if torch.cuda.is_available():
+                #    torch.cuda.empty_cache()
+
                 # Load and prepare control image - must be VAE-encoded!
                 if self.vae is None:
                     raise ValueError("VAE is required when using Z-Image Control. Connect a VAE to the 'vae' input.")
@@ -413,8 +430,23 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
 
                 # Add frame dimension: [B, C, H, W] -> [B, C, 1, H, W]
                 control_latents = control_latents.unsqueeze(2)
-                # Convert to list format expected by transformer
+
+                # Prepare control_context based on control_in_dim
+                # V1: 16 channels (just control latents)
+                # V2.0: 33 channels (control latents + zero padding)
+                # Following diffusers approach: simple zero-padding to match control_in_dim
+                b, c, f, h, w = control_latents.shape
+                if c < control_in_dim:
+                    # Pad with zeros to match control_in_dim (diffusers approach)
+                    padding_channels = control_in_dim - c
+                    zero_padding = torch.zeros(
+                        (b, padding_channels, f, h, w),
+                        device=device,
+                        dtype=inference_dtype,
+                    )
+                    control_latents = torch.cat([control_latents, zero_padding], dim=1)
                 control_context = list(control_latents.unbind(dim=0))
+
                 control_context_scale = self.control.control_context_scale
                 begin_step_percent = self.control.begin_step_percent
                 end_step_percent = self.control.end_step_percent
@@ -425,7 +457,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
 
             # Apply LoRA models to the active transformer.
             # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching.
-            # cached_weights is None when using control (since we create a new model), otherwise it's from model_on_device
+            # cached_weights is None when using control (since we create a new combined model),
+            # otherwise it comes from model_on_device() context.
             exit_stack.enter_context(
                 LayerPatcher.apply_smart_model_patches(
                     model=active_transformer,
@@ -457,15 +490,16 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
 
                 # Determine if control should be applied at this step
                 step_percent = step_idx / total_steps
+                use_control = self.control is not None
                 apply_control = (
-                    control_context is not None
+                    use_control
                     and step_percent >= begin_step_percent
                     and step_percent <= end_step_percent
                 )
 
                 # Transformer returns (List[torch.Tensor], dict) - we only need the tensor list
                 # If control is active, pass control_context to the control transformer
-                if apply_control:
+                if apply_control and control_context is not None:
                     model_output = active_transformer(
                         x=latent_model_input_list,
                         t=timestep,
@@ -486,7 +520,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
 
                 # Apply CFG if enabled
                 if do_classifier_free_guidance and neg_prompt_embeds is not None:
-                    if apply_control:
+                    if apply_control and control_context is not None:
                         model_output_uncond = active_transformer(
                             x=latent_model_input_list,
                             t=timestep,

invokeai/backend/model_manager/load/model_loaders/z_image.py

@@ -420,16 +420,46 @@ def _load_control_adapter(
                     control_block_indices.add(int(parts[1]))
         num_control_blocks = len(control_block_indices) if control_block_indices else 6
 
+        # Determine number of refiner layers from state dict
+        refiner_indices: set[int] = set()
+        for key in sd.keys():
+            if key.startswith("control_noise_refiner."):
+                parts = key.split(".")
+                if len(parts) > 1 and parts[1].isdigit():
+                    refiner_indices.add(int(parts[1]))
+        n_refiner_layers = len(refiner_indices) if refiner_indices else 2
+
+        # Determine control_in_dim from embedder weight shape
+        # control_in_dim = weight.shape[1] / (f_patch_size * patch_size * patch_size)
+        # For patch_size=2, f_patch_size=1: control_in_dim = weight.shape[1] / 4
+        control_in_dim = 16  # Default for V1
+        embedder_key = "control_all_x_embedder.2-1.weight"
+        if embedder_key in sd:
+            weight_shape = sd[embedder_key].shape
+            # weight_shape[1] = f_patch_size * patch_size * patch_size * control_in_dim
+            control_in_dim = weight_shape[1] // 4  # 4 = 1 * 2 * 2
+
+        # Log detected configuration for debugging
+        from invokeai.backend.util.logging import InvokeAILogger
+
+        logger = InvokeAILogger.get_logger(self.__class__.__name__)
+        version = "V2.0" if control_in_dim > 16 else "V1"
+        logger.info(
+            f"Z-Image ControlNet detected: {version} "
+            f"(control_in_dim={control_in_dim}, num_control_blocks={num_control_blocks}, "
+            f"n_refiner_layers={n_refiner_layers})"
+        )
+
         # Create an empty control adapter
         dim = 3840
         with accelerate.init_empty_weights():
             model = ZImageControlAdapter(
                 num_control_blocks=num_control_blocks,
-                control_in_dim=16,
+                control_in_dim=control_in_dim,
                 all_patch_size=(2,),
                 all_f_patch_size=(1,),
                 dim=dim,
-                n_refiner_layers=2,
+                n_refiner_layers=n_refiner_layers,
                 n_heads=30,
                 n_kv_heads=30,
                 norm_eps=1e-05,