diff --git a/docs/src/content/docs/Users Guide/03.Models/Local Models/z-image.mdx b/docs/src/content/docs/Users Guide/03.Models/Local Models/z-image.mdx index b47d5d87d36..b0a9b638324 100644 --- a/docs/src/content/docs/Users Guide/03.Models/Local Models/z-image.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/Local Models/z-image.mdx @@ -110,18 +110,28 @@ Two Z-Image ControlNets are available as starters (both are part of the bundle): - **Z-Image ControlNet Tile** is useful for upscaling and adding detail. The **Upscale** tool can use it as an optional anchor; see [Upscaling](/users-guide/image-generation/upscale/). +Every Z-Image Turbo Fun-ControlNet alibaba-pai has published can be installed from its `.safetensors` file: +the first Union release, Union 2.0, the Union and Tile 2.1 builds (8-step, 2601 and 2602 included) and the 2.1 +Lite builds. Invoke tells them apart by their weights. Invoke does not inpaint through a Z-Image ControlNet: +a control layer takes a control image only. + On the Canvas, add a **Control Layer** while a Z-Image model is selected. The layer uses a Z-Image ControlNet with a default weight of `0.75`. The weight ranges from `0` to `2`, and `0.65` to `0.80` is the recommended range. **Only one** Z-Image control layer can be used per generation. +Earlier Invoke versions placed the hints of the Lite and first Union builds after the wrong layers, and left +out the refiner hints of 2.0 and 2.1 (2.0 also refined its control through the wrong blocks). Applied as they were trained, these ControlNets now follow the control image +more strongly at the same weight; lower the weight if a control image constrains the result too much. + In the workflow editor, the same thing is the **Z-Image ControlNet** node feeding the **Control** input of **Denoise - Z-Image**. That node also needs the **VAE** input connected, which it uses to encode the control image. ## Regional prompting -Z-Image supports **Regional Guidance** layers on the Canvas with **positive** prompts. Regional negative -prompts, auto-negative and regional reference images are not supported for Z-Image. +Z-Image supports **Regional Guidance** layers on the Canvas with **positive** prompts, together with a +control layer if you like. Regional negative prompts, auto-negative and regional reference images are not +supported for Z-Image. In the workflow editor, **Prompt - Z-Image** takes an optional mask, and **Denoise - Z-Image** accepts either one conditioning or a collection of them. diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index f736aabef3d..98e3f11516f 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -28,7 +28,7 @@ compared with PyTorch's default attention there: | --- | --- | --- | --- | | FLUX.1 dev (FP8) | 16.1 s → 14.0 s (−13 %) | 84.2 s → 54.8 s (−35 %) | Also serves ControlNets inside the denoising loop. | | FLUX.2 Klein 9B (FP8) | 15.6 s → 15.1 s (−3 %) | 27.1 s → 23.1 s (−15 %) | | -| Z-Image (FP8) | 10.4 s → 10.0 s (−5 %) | 33.7 s → 25.9 s (−23 %) | Z-Image Control masks its attention and keeps PyTorch's. | +| Z-Image (FP8) | 10.4 s → 10.0 s (−5 %) | 33.7 s → 25.9 s (−23 %) | Also serves ControlNets inside the denoising loop. | | Qwen-Image 2512 (GGUF Q4_K_M) | 78.8 s → 73.5 s (−7 %) | 208.9 s → 150.3 s (−28 %) | One block reaches about 17 % error per call; images were still judged equivalent. | | Krea-2 Turbo (GGUF Q4_K_M) | 11.5 s → 11.0 s (−4 %) | 42.7 s → 36.2 s (−15 %) | Off while `INVOKE_KREA2_SDPA_BACKEND` pins a PyTorch kernel. | | SDXL | 4.4 s → 4.2 s (−5 %) | 19.6 s → 16.1 s (−18 %) | Text cross-attention stays on PyTorch (too short to gain). Also serves ControlNets. | diff --git a/invokeai/app/invocations/vae/z_image_image_to_latents.py b/invokeai/app/invocations/vae/z_image_image_to_latents.py index 149116ae9a0..84d59f59c45 100644 --- a/invokeai/app/invocations/vae/z_image_image_to_latents.py +++ b/invokeai/app/invocations/vae/z_image_image_to_latents.py @@ -50,12 +50,14 @@ def vae_encode( image_tensor: torch.Tensor, tiled: bool = False, tile_size: int = 0, + deterministic: bool = False, ) -> torch.Tensor: - """Encode an image to a Z-Image latent. + """Encode an image to a Z-Image latent: a sample of the encoder's distribution, or its mode when + ``deterministic``. - `z_image_denoise` also encodes its control image here, at generation resolution. That is - usually small enough to encode whole, but an upscale's generation resolution is the - enlarged frame, so the control field carries its own `tiled` and `tile_size`. + `z_image_denoise` also encodes its control image here, deterministically as VideoX-Fun does, at + generation resolution. That is usually small enough to encode whole, but an upscale's generation + resolution is the enlarged frame, so the control field carries its own `tiled` and `tile_size`. """ if not isinstance(vae_info.model, AutoencoderKL): raise TypeError( @@ -87,7 +89,10 @@ def vae_encode( # other node that reaches this class. Tiling is a property of this one encode, so the # state is scoped and restored rather than switched off by hand and left behind. with torch.inference_mode(), scoped_vae_tiling(vae, effective_tile_size): - latents: torch.Tensor = vae.encode(image_tensor).latent_dist.sample().to(dtype=vae.dtype) + latent_dist = vae.encode(image_tensor).latent_dist + latents: torch.Tensor = (latent_dist.mode() if deterministic else latent_dist.sample()).to( + dtype=vae.dtype + ) # Apply scaling_factor and shift_factor from VAE config. # Z-Image uses: latents = (latents - shift_factor) * scaling_factor diff --git a/invokeai/app/invocations/z_image/z_image_denoise.py b/invokeai/app/invocations/z_image/z_image_denoise.py index a96d0962b40..d17adf69e27 100644 --- a/invokeai/app/invocations/z_image/z_image_denoise.py +++ b/invokeai/app/invocations/z_image/z_image_denoise.py @@ -3,9 +3,9 @@ from contextlib import ExitStack from typing import Callable, Iterator, Optional -import einops import torch import torchvision.transforms as tv_transforms +from diffusers.models.transformers.transformer_z_image import SEQ_MULTI_OF from diffusers.schedulers.scheduling_utils import SchedulerMixin from PIL import Image from torchvision.transforms.functional import resize as tv_resize @@ -28,6 +28,10 @@ from invokeai.app.invocations.z_image.z_image_control import ZImageControlField from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.schedulers import ZIMAGE_SCHEDULER_LABELS, ZIMAGE_SCHEDULER_MAP, ZIMAGE_SCHEDULER_NAME_VALUES +from invokeai.backend.fun_controlnet.conditioning import build_fun_control_context +from invokeai.backend.fun_controlnet.hooks import FunControlHooks, install_fun_control, should_apply +from invokeai.backend.fun_controlnet.z_image import ZImageFunControl +from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX @@ -35,7 +39,10 @@ from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes, requires_sidecar_patching from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.rectified_flow.schedule_shift import ShiftFit, time_shift -from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState +from invokeai.backend.stable_diffusion.diffusers_pipeline import ( + PipelineIntermediateState, + image_resized_to_grid_as_tensor, +) from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ZImageConditioningInfo from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice @@ -43,15 +50,9 @@ from invokeai.backend.util.sage_attention import sage_attention_scope from invokeai.backend.z_image.extensions.regional_prompting_extension import ZImageRegionalPromptingExtension from invokeai.backend.z_image.text_conditioning import ZImageTextConditioning -from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter -from invokeai.backend.z_image.z_image_controlnet_extension import ( - ZImageControlNetExtension, - z_image_forward_with_control, -) -from invokeai.backend.z_image.z_image_patchify_utils import SEQ_MULTI_OF from invokeai.backend.z_image.z_image_transformer_patch import patch_transformer_for_regional_prompting -# The transformer's width and head count, which its ControlNet adapters share: 30 heads of 128. +# The transformer's width and head count, which its ControlNets share: 30 heads of 128. Z_IMAGE_HIDDEN_SIZE = 3840 Z_IMAGE_ATTENTION_HEADS = 30 # Peak reserved activation bytes per attended token on a fused SDPA kernel; see `_estimate_working_memory`. @@ -292,12 +293,13 @@ def _estimate_working_memory( on ROCm under Windows, whose fused kernels are disabled, a 1024px generation spilled 7.7 GB into shared system memory while this node reserved only the default working memory. ``sdpa_score_matrix_bytes`` adds the term only where it is really built. An attention mask narrows which fused kernels a build can use, so - the call is told whether one is passed: regional prompting's additive bias, or the padding mask the - ControlNet forward always builds. The plain forward passes none at batch size 1. + the call is told whether one is passed: regional prompting's additive bias. The forward passes none at batch + size 1 otherwise, and a ControlNet's blocks take exactly the arguments the base blocks take. - A ControlNet adapter first runs its own blocks over a sequence of the same length and keeps one hint per - block, alive through the main layers. Measured over the plain forward at 1024px: 3.3, 5.9 and 18.9 - hint-sized tensors (sequence x width x element size) for 1, 6 and 15 blocks, so n + 4 are charged. + A ControlNet first runs its own blocks over a sequence of the same length and keeps one hint per block, + alive until the last layer it adds to. Measured on the hooked forward over the plain one at 1024px: 1.8 + hint-sized tensors (sequence x width x element size) for the 6-block 1.0 layout, 0.9 for the 3-block 2.1 + Lite and 11.6 for the 15-block 2.0 and 2.1, so the n + 4 charged bound every layout. """ GB = 1024**3 seq_len = _padded(image_seq_len) + _padded(text_seq_len) @@ -339,6 +341,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: inference_dtype = TorchDevice.choose_bfloat16_safe_dtype(device) transformer_info = context.models.load(self.transformer.transformer) + control_info = self._load_control(context) # Calculate image token grid dimensions patch_size = 2 # Z-Image uses patch_size=2 @@ -503,12 +506,15 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: ): raise ValueError(f"Unsupported Z-Image model format: {transformer_config.format}") - # Into RAM ahead of the transformer's VRAM lock, so the estimate below can size the adapter's hints. - control_model_info = context.models.load(self.control.control_model) if self.control is not None else None + # Encoded ahead of the transformer's lock, so the VAE does not have to fit beside it. + control_context: torch.Tensor | None = None num_control_blocks = 0 - if control_model_info is not None: - assert isinstance(control_model_info.model, ZImageControlAdapter) - num_control_blocks = len(control_model_info.model.control_layers) + if control_info is not None: + assert isinstance(control_info.model, ZImageFunControl) + control_context = self._build_control_context(context, control_info.model).to( + device=device, dtype=inference_dtype + ) + num_control_blocks = len(control_info.model.control_layers) # Ask the model cache to keep the forward's working memory free; it offloads as much of the # transformer as that takes. Without it the cache reserves only `device_working_mem_gb`, which a @@ -524,7 +530,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: regional_attention_bias_bytes=self._regional_attention_bias_bytes( regional_attn_mask, img_seq_len, inference_dtype ), - has_attention_mask=regional_attn_mask is not None or control_model_info is not None, + has_attention_mask=regional_attn_mask is not None, device=device, dtype=inference_dtype, ) @@ -533,7 +539,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # every other build. working_mem_bytes += peak_dequant_transient_bytes(transformer_info.model, inference_dtype) - # Load transformer - always use base transformer, control is handled via extension (cached_weights, transformer) = exit_stack.enter_context( transformer_info.model_on_device(working_mem_bytes=working_mem_bytes) ) @@ -545,100 +550,19 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # patching on a module that has no patchable weights. model_is_quantized = requires_sidecar_patching(transformer, transformer_config.format) - # Prepare control extension if control is provided - control_extension: ZImageControlNetExtension | None = None - - if self.control is not None: - assert control_model_info is not None - # Locked with the same working memory: the transformer can no longer be offloaded, so an - # adapter placed with only the default reservation would fill the room kept for the forward. - (_, control_adapter) = exit_stack.enter_context( - control_model_info.model_on_device(working_mem_bytes=working_mem_bytes) + control_hooks: FunControlHooks | None = None + if control_info is not None: + assert self.control is not None and control_context is not None + # Locked with the same working memory: the transformer can no longer be offloaded, so a control + # placed with only the default reservation would fill the room kept for the forward. The loader cast + # it to the device's bfloat16-safe dtype, `inference_dtype` here; it casts what crosses into its chain + # and the hints back (`ZImageFunControl._block_runner` / `add_hint`), so nothing cache-managed is cast. + (_, control) = exit_stack.enter_context( + control_info.model_on_device(working_mem_bytes=working_mem_bytes) ) - assert isinstance(control_adapter, ZImageControlAdapter) - - # Get control_in_dim from adapter config (16 for V1, 33 for V2.0) - control_in_dim = control_adapter.config.get("control_in_dim", 16) - - # 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} (Extension): control_in_dim={control_in_dim}, " - f"num_blocks={num_control_blocks}, scale={self.control.control_context_scale}" - ) - - # 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.") - - control_image = context.images.get_pil(self.control.image_name) - - # Resize control image to match output dimensions - control_image = control_image.convert("RGB") - control_image = control_image.resize((self.width, self.height), Image.Resampling.LANCZOS) - - # Convert to tensor format for VAE encoding - from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor - - control_image_tensor = image_resized_to_grid_as_tensor(control_image) - if control_image_tensor.dim() == 3: - control_image_tensor = einops.rearrange(control_image_tensor, "c h w -> 1 c h w") - - # Encode control image through VAE to get latents - vae_info = context.models.load(self.vae.vae) - control_latents = ZImageImageToLatentsInvocation.vae_encode( - vae_info=vae_info, - image_tensor=control_image_tensor, - tiled=self.control.tiled or context.config.get().force_tiled_decode, - tile_size=self.control.tile_size, - ) - - # Move to inference device/dtype - control_latents = control_latents.to(device=device, dtype=inference_dtype) - - # Add frame dimension: [B, C, H, W] -> [C, 1, H, W] (single image) - control_latents = control_latents.squeeze(0).unsqueeze(1) - - # Prepare control_cond based on control_in_dim - # V1: 16 channels (just control latents) - # V2.0: 33 channels = 16 control + 16 reference + 1 mask - # - Channels 0-15: control image latents (from VAE encoding) - # - Channels 16-31: reference/inpaint image latents (zeros for pure control) - # - Channel 32: inpaint mask (1.0 = don't inpaint, 0.0 = inpaint region) - # For pure control (no inpainting), we set mask=1 to tell model "use control, don't inpaint" - c, f, h, w = control_latents.shape - if c < control_in_dim: - padding_channels = control_in_dim - c - if padding_channels == 17: - # V2.0: 16 reference channels (zeros) + 1 mask channel (ones) - ref_padding = torch.zeros( - (16, f, h, w), - device=device, - dtype=inference_dtype, - ) - # Mask channel = 1.0 means "don't inpaint this region, use control signal" - mask_channel = torch.ones( - (1, f, h, w), - device=device, - dtype=inference_dtype, - ) - control_latents = torch.cat([control_latents, ref_padding, mask_channel], dim=0) - else: - # Generic padding with zeros for other cases - zero_padding = torch.zeros( - (padding_channels, f, h, w), - device=device, - dtype=inference_dtype, - ) - control_latents = torch.cat([control_latents, zero_padding], dim=0) - - # Create control extension (adapter is already on device from model_on_device) - control_extension = ZImageControlNetExtension( - control_adapter=control_adapter, - control_cond=control_latents, - weight=self.control.control_context_scale, - begin_step_percent=self.control.begin_step_percent, - end_step_percent=self.control.end_step_percent, + assert isinstance(control, ZImageFunControl) + control_hooks = exit_stack.enter_context( + install_fun_control(transformer, control, control_context, scale=self.control.control_context_scale) ) # Apply LoRA models to the transformer. @@ -700,27 +624,14 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: latent_model_input = latent_model_input.unsqueeze(2) # Add frame dimension latent_model_input_list = list(latent_model_input.unbind(dim=0)) - # Determine if control should be applied at this step - apply_control = control_extension is not None and control_extension.should_apply( - user_step, total_steps - ) + # Heun's two forwards of a step share its user step, and so its control gate. + self._gate_control(control_hooks, total_steps, user_step) - # Run forward pass - if apply_control: - model_out_list, _ = z_image_forward_with_control( - transformer=transformer, - x=latent_model_input_list, - t=timestep, - cap_feats=[pos_prompt_embeds], - control_extension=control_extension, - ) - else: - model_output = transformer( - x=latent_model_input_list, - t=timestep, - cap_feats=[pos_prompt_embeds], - ) - model_out_list = model_output[0] + model_out_list = transformer( + x=latent_model_input_list, + t=timestep, + cap_feats=[pos_prompt_embeds], + )[0] noise_pred_cond = torch.stack([t.float() for t in model_out_list], dim=0) noise_pred_cond = noise_pred_cond.squeeze(2) @@ -728,21 +639,11 @@ 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: - model_out_list_uncond, _ = z_image_forward_with_control( - transformer=transformer, - x=latent_model_input_list, - t=timestep, - cap_feats=[neg_prompt_embeds], - control_extension=control_extension, - ) - else: - model_output_uncond = transformer( - x=latent_model_input_list, - t=timestep, - cap_feats=[neg_prompt_embeds], - ) - model_out_list_uncond = model_output_uncond[0] + model_out_list_uncond = transformer( + x=latent_model_input_list, + t=timestep, + cap_feats=[neg_prompt_embeds], + )[0] noise_pred_uncond = torch.stack([t.float() for t in model_out_list_uncond], dim=0) noise_pred_uncond = noise_pred_uncond.squeeze(2) @@ -817,27 +718,13 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: latent_model_input = latent_model_input.unsqueeze(2) # Add frame dimension latent_model_input_list = list(latent_model_input.unbind(dim=0)) - # Determine if control should be applied at this step - apply_control = control_extension is not None and control_extension.should_apply( - step_idx, total_steps - ) + self._gate_control(control_hooks, total_steps, step_idx) - # Run forward pass - use custom forward with control if extension is active - if apply_control: - model_out_list, _ = z_image_forward_with_control( - transformer=transformer, - x=latent_model_input_list, - t=timestep, - cap_feats=[pos_prompt_embeds], - control_extension=control_extension, - ) - else: - model_output = transformer( - x=latent_model_input_list, - t=timestep, - cap_feats=[pos_prompt_embeds], - ) - model_out_list = model_output[0] # Extract list of tensors from tuple + model_out_list = transformer( + x=latent_model_input_list, + t=timestep, + cap_feats=[pos_prompt_embeds], + )[0] # The list of output tensors noise_pred_cond = torch.stack([t.float() for t in model_out_list], dim=0) noise_pred_cond = noise_pred_cond.squeeze(2) # Remove frame dimension @@ -845,21 +732,11 @@ 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: - model_out_list_uncond, _ = z_image_forward_with_control( - transformer=transformer, - x=latent_model_input_list, - t=timestep, - cap_feats=[neg_prompt_embeds], - control_extension=control_extension, - ) - else: - model_output_uncond = transformer( - x=latent_model_input_list, - t=timestep, - cap_feats=[neg_prompt_embeds], - ) - model_out_list_uncond = model_output_uncond[0] # Extract list of tensors from tuple + model_out_list_uncond = transformer( + x=latent_model_input_list, + t=timestep, + cap_feats=[neg_prompt_embeds], + )[0] noise_pred_uncond = torch.stack([t.float() for t in model_out_list_uncond], dim=0) noise_pred_uncond = noise_pred_uncond.squeeze(2) @@ -889,6 +766,53 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: return latents + def _load_control(self, context: InvocationContext) -> LoadedModel | None: + """The control model in RAM, ahead of the transformer's lock so the reservation can size its hints.""" + if self.control is None: + return None + if self.vae is None: + raise ValueError( + "Z-Image ControlNet needs the VAE that encodes its control image: connect one to the 'vae' input." + ) + control_info = context.models.load(self.control.control_model) + if not isinstance(control_info.model, ZImageFunControl): + raise TypeError( + f"Expected a Z-Image ControlNet for '{self.control.control_model.name}', " + f"got {type(control_info.model).__name__}." + ) + return control_info + + def _build_control_context(self, context: InvocationContext, control: ZImageFunControl) -> torch.Tensor: + """The control image at the generation size as the tokens ``control`` embeds. + + VideoX-Fun's context for control alone: the control image's latents, a zero keep mask and zero masked-image + latents, packed 2x2 as the transformer packs its latents. Encoded to the mode of the VAE's distribution as + VideoX-Fun encodes it, not to a sample of it. + """ + assert self.control is not None and self.vae is not None + field = self.control + vae_info = context.models.load(self.vae.vae) + tiled = field.tiled or context.config.get().force_tiled_decode + + def encode(pixels: torch.Tensor) -> torch.Tensor: + return ZImageImageToLatentsInvocation.vae_encode( + vae_info=vae_info, image_tensor=pixels, tiled=tiled, tile_size=field.tile_size, deterministic=True + ) + + image = context.images.get_pil(field.image_name).convert("RGB") + image = image.resize((self.width, self.height), Image.Resampling.LANCZOS) + pixels = image_resized_to_grid_as_tensor(image).unsqueeze(0) + return control.pack(build_fun_control_context(encode=encode, control_image=pixels)) + + def _gate_control(self, hooks: FunControlHooks | None, total_steps: int, step_index: int) -> None: + # One gate per user step: both guidance passes run with the control or neither does. + if hooks is None: + return + assert self.control is not None + hooks.active = should_apply( + step_index, total_steps, self.control.begin_step_percent, self.control.end_step_percent + ) + def _prepare_noise_tensor( self, context: InvocationContext, inference_dtype: torch.dtype, device: torch.device ) -> torch.Tensor: diff --git a/invokeai/backend/fun_controlnet/__init__.py b/invokeai/backend/fun_controlnet/__init__.py index 239cbac6f91..138b5b81868 100644 --- a/invokeai/backend/fun_controlnet/__init__.py +++ b/invokeai/backend/fun_controlnet/__init__.py @@ -2,8 +2,8 @@ Every Fun control model follows one recipe: the control context is projected to the block width, runs through a short chain of blocks that are the base transformer's block plus two projections, and each block's ``after_proj`` output is -added to the base image stream after a listed base layer. ``branch.py`` holds the recipe and the seams a family fills -in, ``hooks.py`` attaches a branch to a transformer without touching its forward, and a module per family -(``qwen_image.py``, ``flux2.py``) supplies the block class, the layer map and the calling convention. -``conditioning.py`` builds the control context all of them were trained on. +added to the base image stream after a listed base layer. ``branch.py`` holds the recipe, the stages a control runs in +and the seams a single-chain family fills in; ``hooks.py`` attaches a control to a transformer without touching its +forward; a module per family (``qwen_image.py``, ``flux2.py``, ``z_image.py``) supplies the block class, the layer map +and the calling convention. ``conditioning.py`` builds the control context all of them were trained on. """ diff --git a/invokeai/backend/fun_controlnet/branch.py b/invokeai/backend/fun_controlnet/branch.py index 8717ad6d3c4..50cec75110d 100644 --- a/invokeai/backend/fun_controlnet/branch.py +++ b/invokeai/backend/fun_controlnet/branch.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from typing import Any import torch @@ -9,17 +10,102 @@ """Runs one control block on the control image stream and returns the stream it produced. Whatever else the block takes and updates (its own text stream, modulation, RoPE) the family closes over.""" +StageRunner = Callable[[Mapping[str, Any], list[torch.Tensor]], tuple[list[torch.Tensor], torch.Tensor | None]] +"""Runs a stage's chain from its entry block's bound forward arguments and ``carry``, which holds the control stream +the previous stage handed on (empty for the first stage). The stage pops the stream it continues from, so nothing else +keeps it alive while the chain runs. Returns one hint per hint block and the control stream for the next stage.""" -class FunControlBranch(nn.Module, ABC): - """``control_img_in`` and ``control_blocks``, with the base layers their hints land on. + +@dataclass(frozen=True) +class FunControlStage: + """One place in the base forward where a control chain runs. + + ``entry`` is the base block whose pre-hook runs ``run`` on that block's bound arguments; the stage's hint ``k`` is + added to the output of ``hint_blocks[k]``. A family with one chain has one stage; a family whose chain is split + across the base's own stages (a refiner ahead of the main blocks) has several, which one forward reaches in order, + each handing its control stream to the next. + """ + + entry: nn.Module + hint_blocks: tuple[nn.Module, ...] + run: StageRunner + + +def run_vace_chain( + make_input: Callable[[], torch.Tensor], + x: torch.Tensor, + blocks: Sequence[nn.Module], + run_block: BlockRunner, + *, + num_hints: int, +) -> tuple[list[torch.Tensor], torch.Tensor]: + """The VACE recipe every Fun control chain follows: block 0 adds ``before_proj(c)`` to the base stream ``x``, every + block runs on the control stream ``c``, and the first ``num_hints`` blocks' ``after_proj`` outputs are the hints. + + The chain builds its input ``c`` with ``make_input`` and holds the only reference to it, so each stream is freed as + soon as the next one replaces it. An input the caller had built and kept would stay alive, an image-length hidden + state, through every block. Returns the hints and the control stream after the last block. The chain runs in + ``c``'s dtype; ``x`` is cast to it. + """ + c = make_input() + hints: list[torch.Tensor] = [] + for index, block in enumerate(blocks): + if index == 0: + c = block.before_proj(c) + x.to(c.dtype) + c = run_block(block, c) + if index < num_hints: + hints.append(block.after_proj(c)) + return hints, c + + +def hint_blocks(base_blocks: Sequence[nn.Module], layers: Sequence[int]) -> tuple[nn.Module, ...]: + """The base blocks at ``layers``, refused when the transformer is too shallow for them.""" + layers = tuple(layers) + if len(base_blocks) <= layers[-1]: + raise ValueError( + f"The control branch adds hints after base layers {layers}; this transformer has {len(base_blocks)} blocks." + ) + return tuple(base_blocks[layer] for layer in layers) + + +class FunControl(nn.Module, ABC): + """A Fun control model as the hooks see it: its width and context features, the stages it runs inside a base + transformer's forward, and how a hint goes onto a base block's output.""" + + @property + @abstractmethod + def hidden_size(self) -> int: + """The width of the hints.""" + + @property + @abstractmethod + def control_in_dim(self) -> int: + """The features per token of the control context.""" + + @abstractmethod + def base_hidden_size(self, transformer: nn.Module) -> int: + """The width of the base image stream the hints are added to.""" + + @abstractmethod + def stages(self, transformer: nn.Module, control_context: torch.Tensor) -> Sequence[FunControlStage]: + """The stages to hook into ``transformer``, in the order one forward reaches their entry blocks. Refuses a + transformer whose blocks the hints cannot land on.""" + + @abstractmethod + def add_hint(self, output: Any, hint: torch.Tensor, scale: float) -> Any: + """A base block's ``output`` with ``hint * scale`` added to its image stream and nothing else touched.""" + + +class FunControlBranch(FunControl): + """``control_img_in`` and ``control_blocks``, with the base layers their hints land on: one chain, entered at base + block 0. ``control_layers[k]`` is the index of the base block after whose output ``control_blocks[k]``'s hint is added. Block 0 carries ``before_proj`` and every block ``after_proj``, both ``Linear(hidden, hidden)``; the family's block class adds them to the base block it subclasses, and keeps that block's ``forward``. - The abstract methods are the family seams the hooks go through: which base blocks to hook, how to read the base - image stream from block 0's arguments, how to call a control block with them, and where in a block's output the hint - goes. + The abstract methods are the family seams: which base blocks to hook, how to read the base image stream from block + 0's arguments, how to call a control block with them, and where in a block's output the hint goes. """ def __init__( @@ -53,23 +139,28 @@ def compute_hints( ) -> list[torch.Tensor]: """One hint per control block, from the control context and the base image stream ``x`` as it enters base block 0. The chain runs in the branch's dtype; ``add_hint`` casts back.""" - c = self.control_img_in(control_context.to(self.dtype)) - hints: list[torch.Tensor] = [] - for index, block in enumerate(self.control_blocks): - if index == 0: - c = block.before_proj(c) + x.to(c.dtype) - c = run_block(block, c) - hints.append(block.after_proj(c)) + hints, _ = run_vace_chain( + lambda: self.control_img_in(control_context.to(self.dtype)), + x, + self.control_blocks, + run_block, + num_hints=len(self.control_blocks), + ) return hints + def stages(self, transformer: nn.Module, control_context: torch.Tensor) -> Sequence[FunControlStage]: + blocks = self.base_blocks(transformer) + targets = hint_blocks(blocks, self.control_layers) + + def run(inputs: Mapping[str, Any], _carry: list[torch.Tensor]) -> tuple[list[torch.Tensor], None]: + return self.compute_hints(control_context, self.image_stream(inputs), self.block_runner(inputs)), None + + return (FunControlStage(entry=blocks[0], hint_blocks=targets, run=run),) + @abstractmethod def base_blocks(self, transformer: nn.Module) -> Sequence[nn.Module]: """The base transformer's blocks, indexed as ``control_layers`` are.""" - @abstractmethod - def base_hidden_size(self, transformer: nn.Module) -> int: - """The width of the base image stream the hints are added to.""" - @abstractmethod def image_stream(self, block_inputs: Mapping[str, Any]) -> torch.Tensor: """The base image stream among base block 0's bound forward arguments.""" @@ -78,7 +169,3 @@ def image_stream(self, block_inputs: Mapping[str, Any]) -> torch.Tensor: def block_runner(self, block_inputs: Mapping[str, Any]) -> BlockRunner: """A runner that calls a control block the way the base called its block 0, threading the branch's own copy of whatever else the block updates from call to call.""" - - @abstractmethod - def add_hint(self, output: Any, hint: torch.Tensor, scale: float) -> Any: - """A base block's ``output`` with ``hint * scale`` added to its image stream and nothing else touched.""" diff --git a/invokeai/backend/fun_controlnet/hooks.py b/invokeai/backend/fun_controlnet/hooks.py index 0f18378f414..0a7e7c5311c 100644 --- a/invokeai/backend/fun_controlnet/hooks.py +++ b/invokeai/backend/fun_controlnet/hooks.py @@ -7,66 +7,86 @@ from torch import nn from torch.utils.hooks import RemovableHandle -from invokeai.backend.fun_controlnet.branch import FunControlBranch +from invokeai.backend.fun_controlnet.branch import FunControl, FunControlStage + +HintAdder = Callable[[Any, torch.Tensor, float], Any] +"""``FunControl.add_hint``: a base block's output with a scaled hint added to its image stream.""" def should_apply(step_index: int, total_steps: int, begin_step_percent: float, end_step_percent: float) -> bool: - """Whether the control applies at a step: ``begin <= step_index / total_steps <= end``, both ends included, as the - Z-Image control extension counts.""" + """Whether the control applies at a step: ``begin <= step_index / total_steps <= end``, both ends included.""" if total_steps <= 0: return True return begin_step_percent <= step_index / total_steps <= end_step_percent class FunControlHooks: - """The hooks a branch is installed with. ``active`` gates the forwards that follow without re-installing. + """The hooks a control model is installed with. ``active`` gates the forwards that follow without re-installing. - Base block 0's pre-hook computes the hints from that forward's real arguments; the post-hook of each listed base - block adds its hint and the last one drops them, so no hint of one forward reaches the next and nothing stays alive - through the layers that come after the last listed one. + Each stage's entry block runs that stage's chain in a pre-hook, on the block's real arguments and the control stream + the stage before it carried; the post-hook of each of the stage's hint blocks adds its hint, and the last one drops + them. The first stage's entry opens every forward and drops whatever an earlier one left behind, so no hint or + control stream of one forward reaches the next, and none stays alive past the block that consumes it. """ - def __init__( - self, branch: FunControlBranch, blocks: Sequence[nn.Module], control_context: torch.Tensor, scale: float - ) -> None: + def __init__(self, stages: Sequence[FunControlStage], add_hint: HintAdder, scale: float) -> None: self.active = True - self._branch = branch - self._control_context = control_context + self._stages = tuple(stages) + self._add_hint = add_hint self._scale = scale - self._hints: list[torch.Tensor] | None = None - # Bound once: the base calls its blocks by keyword, a checkpointing caller by position, and the family reads - # the arguments by name either way. - self._signature = inspect.signature(blocks[0].forward) - last_layer = branch.control_layers[-1] - self._handles: list[RemovableHandle] = [ - blocks[0].register_forward_pre_hook(self._compute_hints, with_kwargs=True) - ] - for position, layer in enumerate(branch.control_layers): + self._hints: list[list[torch.Tensor] | None] = [None] * len(self._stages) + # What the stage that ran last hands to the one after it, which pops it: held here, and only here, until then. + self._carry: list[torch.Tensor] = [] + self._handles: list[RemovableHandle] = [] + for index, stage in enumerate(self._stages): + # Bound once: the base calls its blocks by keyword or by position, and the family reads the arguments by + # name either way. + signature = inspect.signature(stage.entry.forward) self._handles.append( - blocks[layer].register_forward_hook(self._hint_adder(position, clear=layer == last_layer)) + stage.entry.register_forward_pre_hook(self._stage_runner(index, signature), with_kwargs=True) ) + last = len(stage.hint_blocks) - 1 + for position, block in enumerate(stage.hint_blocks): + self._handles.append( + block.register_forward_hook(self._hint_adder(index, position, clear=position == last)) + ) - def _compute_hints(self, _module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None: - # Block 0 runs first in every forward: whatever a previous forward left behind is stale here. - self._hints = None - if not self.active: + def _stage_runner( + self, index: int, signature: inspect.Signature + ) -> Callable[[nn.Module, tuple[Any, ...], dict[str, Any]], None]: + def run_stage(_module: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None: + if index == 0: + # The first stage's entry runs first in every forward: whatever a previous forward left is stale here. + self._hints = [None] * len(self._stages) + self._carry.clear() + if not self.active: + return None + if index > 0 and not self._carry: + raise RuntimeError( + f"Control stage {index} was entered without the control stream of stage {index - 1}: the forward " + "did not run the block that starts the control." + ) + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + hints, handed_on = self._stages[index].run(bound.arguments, self._carry) + self._hints[index] = hints + self._carry.clear() + # The last stage has no one to hand its stream to; holding it would keep it alive into the next forward. + if handed_on is not None and index + 1 < len(self._stages): + self._carry.append(handed_on) return None - bound = self._signature.bind(*args, **kwargs) - bound.apply_defaults() - inputs = bound.arguments - self._hints = self._branch.compute_hints( - self._control_context, self._branch.image_stream(inputs), self._branch.block_runner(inputs) - ) - return None - def _hint_adder(self, position: int, clear: bool) -> Callable[[nn.Module, tuple[Any, ...], Any], Any]: + return run_stage + + def _hint_adder(self, stage: int, position: int, clear: bool) -> Callable[[nn.Module, tuple[Any, ...], Any], Any]: def add_hint(_module: nn.Module, _args: tuple[Any, ...], output: Any) -> Any: - if self._hints is None: + hints = self._hints[stage] + if hints is None: return None - hint = self._hints[position] + hint = hints[position] if clear: - self._hints = None - return self._branch.add_hint(output, hint, self._scale) + self._hints[stage] = None + return self._add_hint(output, hint, self._scale) return add_hint @@ -74,36 +94,31 @@ def remove(self) -> None: for handle in self._handles: handle.remove() self._handles.clear() - self._hints = None + self._hints = [None] * len(self._stages) + self._carry.clear() @contextmanager def install_fun_control( - transformer: nn.Module, branch: FunControlBranch, control_context: torch.Tensor, *, scale: float + transformer: nn.Module, control: FunControl, control_context: torch.Tensor, *, scale: float ) -> Iterator[FunControlHooks]: - """Run ``transformer`` with ``branch``'s hints for the duration of the context. + """Run ``transformer`` with ``control``'s hints for the duration of the context. The hints are recomputed on every forward (each denoising step, each guidance branch) from the control context, which is the same for both guidance branches. The hooks come off on exit however the body ends: the blocks belong to the model cache and outlive this call. """ - blocks = branch.base_blocks(transformer) - base_hidden_size = branch.base_hidden_size(transformer) - if base_hidden_size != branch.hidden_size: + base_hidden_size = control.base_hidden_size(transformer) + if base_hidden_size != control.hidden_size: raise ValueError( - f"The control branch is {branch.hidden_size} wide and the transformer {base_hidden_size}: " + f"The control branch is {control.hidden_size} wide and the transformer {base_hidden_size}: " "this is not a control model for this transformer." ) - if len(blocks) <= branch.control_layers[-1]: - raise ValueError( - f"The control branch adds hints after base layers {branch.control_layers}; " - f"this transformer has {len(blocks)} blocks." - ) - if control_context.dim() != 3 or control_context.shape[-1] != branch.control_in_dim: + if control_context.dim() != 3 or control_context.shape[-1] != control.control_in_dim: raise ValueError( - f"The control context must be (batch, tokens, {branch.control_in_dim}), got {tuple(control_context.shape)}." + f"The control context must be (batch, tokens, {control.control_in_dim}), got {tuple(control_context.shape)}." ) - hooks = FunControlHooks(branch, blocks, control_context, scale) + hooks = FunControlHooks(control.stages(transformer, control_context), control.add_hint, scale) try: yield hooks finally: diff --git a/invokeai/backend/fun_controlnet/z_image.py b/invokeai/backend/fun_controlnet/z_image.py new file mode 100644 index 00000000000..8ec56b65a38 --- /dev/null +++ b/invokeai/backend/fun_controlnet/z_image.py @@ -0,0 +1,322 @@ +"""Z-Image Fun-ControlNet: the control alibaba-pai trained for the Z-Image transformer, in the four layouts it published. + +Reference: VideoX-Fun ``videox_fun/models/z_image_transformer2d_control.py``, ``pipeline_z_image_control.py`` and +``config/z_image/z_image_control{,_2.0,_2.1,_2.1_lite}.yaml``. + +The chain is split across the base's own two stages, so it runs as two hooked stages: + +- **Refiner stage**, entered at ``noise_refiner[0]`` on the image stream alone. The control context is embedded, padded + the way the image stream is, and run through the refiner chain. From 2.0 on that chain joins the image stream like a + VACE chain and adds its two hints after both base refiner blocks; 1.0 runs plain refiner blocks and adds nothing here. +- **Main stage**, entered at ``layers[0]`` on the unified ``[image | caption]`` stream. The refined control stream is + joined with the caption part of that stream and runs through ``control_layers``, one hint per block after the listed + base layers. + +The checkpoints carry no config, and 2.0 and 2.1 have the same keys and shapes; the variant is read off the embedder +width, the block count and whether the refiner projections were trained (2.0's are all zero). +""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from typing import Any, Literal + +import accelerate +import torch +import torch.nn.functional as F +from diffusers.models.transformers.transformer_z_image import SEQ_MULTI_OF, ZImageTransformerBlock +from einops import rearrange +from torch import nn + +from invokeai.backend.fun_controlnet.branch import BlockRunner, FunControl, FunControlStage, hint_blocks, run_vace_chain + +# Z-Image attention heads are 128 wide. The checkpoint does not say, so the head count follows from the width. +Z_IMAGE_ATTENTION_HEAD_DIM = 128 +# VideoX-Fun `control_refiner_layers_places`: the refiner chain's hints land after both base refiner blocks. +Z_IMAGE_FUN_CONTROL_REFINER_LAYERS: tuple[int, ...] = (0, 1) +_NUM_REFINER_BLOCKS = 2 +# The transformer's only patch size, 2x2 spatial by 1 frame, names the embedder. +_EMBEDDER_KEY = "2-1" +_EMBEDDER_WEIGHT = f"control_all_x_embedder.{_EMBEDDER_KEY}.weight" +# Every Z-Image block takes its stream as `x`; the other arguments are handed to the control blocks as they are. +_STREAM = "x" +_REFINER_PROJECTIONS = tuple( + f"control_noise_refiner.{block}.{projection}.{tensor}" + for block, projection in ((0, "before_proj"), (0, "after_proj"), (1, "after_proj")) + for tensor in ("weight", "bias") +) + +RefinerChain = Literal["plain", "refiner", "main"] + + +@dataclass(frozen=True) +class ZImageFunControlVariant: + """One published layout. + + ``refiner_chain`` is what runs in the refiner stage: ``plain`` refiner blocks with no projections, which neither + join the image stream nor add hints (1.0); the ``refiner`` blocks with their projections (2.1, Lite); or the + ``main`` ``control_layers`` (2.0, whose config ran the refiner chain over them; its own refiner blocks were never + trained and are not built). + """ + + name: str + control_in_channels: int + main_layers: tuple[int, ...] + refiner_chain: RefinerChain + + +Z_IMAGE_FUN_CONTROL_1_0 = ZImageFunControlVariant("1.0", 16, (0, 5, 10, 15, 20, 25), "plain") +Z_IMAGE_FUN_CONTROL_2_0 = ZImageFunControlVariant("2.0", 33, tuple(range(0, 30, 2)), "main") +Z_IMAGE_FUN_CONTROL_2_1 = ZImageFunControlVariant("2.1", 33, tuple(range(0, 30, 2)), "refiner") +Z_IMAGE_FUN_CONTROL_2_1_LITE = ZImageFunControlVariant("2.1 Lite", 33, (0, 10, 20), "refiner") + +# The refiner projections each refiner chain leaves in the file: 1.0 has none; 2.0 carries them all zero, as initialised, +# because its config never ran them; 2.1 and Lite trained them. +_REFINER_PROJECTION_STATE: dict[RefinerChain, str] = {"plain": "absent", "main": "untrained", "refiner": "trained"} +# What tells the layouts apart: embedder input features (2x2 patches), control blocks, and the refiner projections. +_LAYOUTS: dict[tuple[int, int, str], ZImageFunControlVariant] = { + ( + 4 * variant.control_in_channels, + len(variant.main_layers), + _REFINER_PROJECTION_STATE[variant.refiner_chain], + ): variant + for variant in ( + Z_IMAGE_FUN_CONTROL_1_0, + Z_IMAGE_FUN_CONTROL_2_0, + Z_IMAGE_FUN_CONTROL_2_1, + Z_IMAGE_FUN_CONTROL_2_1_LITE, + ) +} + + +class ZImageFunControlBlock(ZImageTransformerBlock): + """A Z-Image block plus the projections of the control recipe it carries; ``forward`` is the base block's.""" + + def __init__( + self, *, layer_id: int, dim: int, num_attention_heads: int, before_proj: bool, after_proj: bool + ) -> None: + # Z-Image's own block settings: RMSNorm eps 1e-5, query/key norm, timestep modulation. + super().__init__( + layer_id, + dim, + n_heads=num_attention_heads, + n_kv_heads=num_attention_heads, + norm_eps=1e-5, + qk_norm=True, + modulation=True, + ) + if before_proj: + self.before_proj = nn.Linear(dim, dim) + if after_proj: + self.after_proj = nn.Linear(dim, dim) + + +class ZImageFunControl(FunControl): + """The control for ``ZImageTransformer2DModel``, under the checkpoint's own names: ``control_all_x_embedder``, + ``control_noise_refiner`` and ``control_layers``.""" + + def __init__(self, variant: ZImageFunControlVariant, *, hidden_size: int, num_attention_heads: int) -> None: + super().__init__() + self.variant = variant + self.control_all_x_embedder = nn.ModuleDict( + {_EMBEDDER_KEY: nn.Linear(4 * variant.control_in_channels, hidden_size)} + ) + self.control_layers = nn.ModuleList( + ZImageFunControlBlock( + layer_id=index, + dim=hidden_size, + num_attention_heads=num_attention_heads, + before_proj=index == 0, + after_proj=True, + ) + for index in range(len(variant.main_layers)) + ) + joined = variant.refiner_chain == "refiner" + self.control_noise_refiner = nn.ModuleList( + ZImageFunControlBlock( + layer_id=1000 + index, + dim=hidden_size, + num_attention_heads=num_attention_heads, + before_proj=joined and index == 0, + after_proj=joined, + ) + for index in range(_NUM_REFINER_BLOCKS if variant.refiner_chain != "main" else 0) + ) + + @classmethod + def from_state_dict( + cls, sd: Mapping[str, torch.Tensor], *, dtype: torch.dtype, attention_head_dim: int = Z_IMAGE_ATTENTION_HEAD_DIM + ) -> "ZImageFunControl": + """Build the control the checkpoint describes and load it, cast to ``dtype``. + + The width comes off the embedder, the variant off the embedder's input features, the control block count and + the refiner projections (2.0's are all zero). Anything that is not one of the published layouts is refused by + name rather than loaded into a control that would run and produce noise. + """ + weight = sd.get(_EMBEDDER_WEIGHT) + if weight is None or weight.dim() != 2: + raise ValueError(f"Not a Z-Image Fun-ControlNet checkpoint: `{_EMBEDDER_WEIGHT}` is missing.") + hidden_size, in_features = weight.shape + if hidden_size % attention_head_dim: + raise ValueError( + f"A {hidden_size}-wide control cannot be split into {attention_head_dim}-wide attention heads." + ) + variant = _variant_of(sd, in_features) + if variant.refiner_chain == "main": + sd = {key: tensor for key, tensor in sd.items() if not key.startswith("control_noise_refiner.")} + + with accelerate.init_empty_weights(): + control = cls(variant, hidden_size=hidden_size, num_attention_heads=hidden_size // attention_head_dim) + missing, unexpected = control.load_state_dict(dict(sd), strict=False, assign=True) + if missing or unexpected: + raise ValueError( + f"This checkpoint does not match the Z-Image Fun-ControlNet {variant.name} layout: " + f"missing {sorted(missing)[:5]}, unexpected {sorted(unexpected)[:5]}." + ) + return control.to(dtype=dtype).eval() + + @property + def hidden_size(self) -> int: + return self._embedder.out_features + + @property + def control_in_dim(self) -> int: + return self._embedder.in_features + + @property + def dtype(self) -> torch.dtype: + return self._embedder.weight.dtype + + @property + def _embedder(self) -> nn.Linear: + return self.control_all_x_embedder[_EMBEDDER_KEY] + + def pack(self, context: torch.Tensor) -> torch.Tensor: + """The control context ``(B, 33, h, w)`` as the tokens the embedder takes, ``(B, h/2 * w/2, 4 * C)``. + + Patched 2x2 in the order the transformer patches its latents (``_patchify_image``: row-major tokens, features + in ``(pH, pW, C)`` order). The 1.0 layout takes the control latents alone, so only their 16 channels are kept. + """ + return rearrange( + context[:, : self.variant.control_in_channels], "b c (h ph) (w pw) -> b (h w) (ph pw c)", ph=2, pw=2 + ) + + def base_hidden_size(self, transformer: nn.Module) -> int: + return transformer.dim + + def stages(self, transformer: nn.Module, control_context: torch.Tensor) -> Sequence[FunControlStage]: + refiner_targets = ( + hint_blocks(transformer.noise_refiner, Z_IMAGE_FUN_CONTROL_REFINER_LAYERS) + if self.variant.refiner_chain != "plain" + else () + ) + return ( + FunControlStage( + entry=transformer.noise_refiner[0], + hint_blocks=refiner_targets, + run=partial(self._run_refiner_stage, control_context), + ), + FunControlStage( + entry=transformer.layers[0], + hint_blocks=hint_blocks(transformer.layers, self.variant.main_layers), + run=self._run_main_stage, + ), + ) + + def add_hint(self, output: torch.Tensor, hint: torch.Tensor, scale: float) -> torch.Tensor: + return output + (hint * scale).to(output.dtype) + + def _run_refiner_stage( + self, control_context: torch.Tensor, inputs: Mapping[str, Any], _carry: list[torch.Tensor] + ) -> tuple[list[torch.Tensor], torch.Tensor]: + x = inputs[_STREAM] + num_tokens, padded_tokens = control_context.shape[1], x.shape[1] + if padded_tokens != num_tokens + (-num_tokens) % SEQ_MULTI_OF: + raise ValueError( + f"The control context covers {num_tokens} image tokens and the transformer's image stream is " + f"{padded_tokens} long: the control image must be encoded at the size of the latents being denoised." + ) + run_block = self._block_runner(inputs) + if self.variant.refiner_chain == "plain": + c = self._embed(control_context, x) + for block in self.control_noise_refiner: + c = run_block(block, c) + return [], c + chain = self.control_layers if self.variant.refiner_chain == "main" else self.control_noise_refiner + return run_vace_chain( + partial(self._embed, control_context, x), + x, + chain, + run_block, + num_hints=len(Z_IMAGE_FUN_CONTROL_REFINER_LAYERS), + ) + + def _run_main_stage(self, inputs: Mapping[str, Any], carry: list[torch.Tensor]) -> tuple[list[torch.Tensor], None]: + unified = inputs[_STREAM] + hints, _ = run_vace_chain( + partial(_beside_the_caption, carry, unified), + unified, + self.control_layers, + self._block_runner(inputs), + num_hints=len(self.control_layers), + ) + return hints, None + + def _embed(self, control_context: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """The control stream the refiner stage starts from: the context embedded, padded as the image stream ``x``.""" + num_tokens, padded_tokens = control_context.shape[1], x.shape[1] + # Embedded at the padded length, as VideoX-Fun embeds its padded patches: what a GEMM computes for one row + # depends on how many rows it computes at once, so this keeps the rows the reference computes. + tokens = F.pad(control_context.to(self.dtype), (0, 0, 0, padded_tokens - num_tokens)) + c = self._embedder(tokens).expand(x.shape[0], -1, -1) + # The image stream's padding rows are the base's `x_pad_token` here, which is what VideoX-Fun writes into the + # control stream's padding rows. Taken from the stream, they are exactly what the base computed, however its + # weights are stored. + return torch.cat([c[:, :num_tokens], x[:, num_tokens:].to(c.dtype)], dim=1) + + def _block_runner(self, inputs: Mapping[str, Any]) -> BlockRunner: + # Mask, RoPE, modulation input: the control blocks take exactly what the base block took, regional prompting's + # attention bias included. The chain runs in the control's dtype, so the floating inputs are cast on the way in + # and the hint on the way out (`add_hint`). RoPE is complex and passes untouched. + passthrough = { + name: value.to(self.dtype) if isinstance(value, torch.Tensor) and value.is_floating_point() else value + for name, value in inputs.items() + if name != _STREAM + } + + def run(block: nn.Module, c: torch.Tensor) -> torch.Tensor: + return block(c, **passthrough) + + return run + + +def _beside_the_caption(carry: list[torch.Tensor], unified: torch.Tensor) -> torch.Tensor: + """The control stream the main stage starts from: the refiner stage's, popped from ``carry``, beside the refined + caption. The unified stream is ``[image | caption]``, its image part as long as the control stream.""" + carried = carry.pop() + return torch.cat([carried, unified[:, carried.shape[1] :].to(carried.dtype)], dim=1) + + +def _variant_of(sd: Mapping[str, torch.Tensor], in_features: int) -> ZImageFunControlVariant: + num_blocks = 1 + max( + (int(parts[1]) for key in sd if (parts := key.split("."))[0] == "control_layers" and parts[1].isdigit()), + default=-1, + ) + projections = [sd[key] for key in _REFINER_PROJECTIONS if key in sd] + if not projections: + refiner = "absent" + elif len(projections) < len(_REFINER_PROJECTIONS): + refiner = "incomplete" + else: + refiner = "trained" if any(bool(tensor.any()) for tensor in projections) else "untrained" + variant = _LAYOUTS.get((in_features, num_blocks, refiner)) + if variant is None: + known = "; ".join( + f"{layout.name}: {features} features, {blocks} blocks, {projection}" + for (features, blocks, projection), layout in _LAYOUTS.items() + ) + raise ValueError( + f"Not a Z-Image Fun-ControlNet layout Invoke knows: `control_all_x_embedder` takes {in_features} " + f"features, with {num_blocks} control blocks and {refiner} refiner projections ({known})." + ) + return variant diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py index e1624e08c41..eadd672b38d 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/z_image.py @@ -7,9 +7,9 @@ import torch from transformers import AutoTokenizer, Qwen3ForCausalLM +from invokeai.backend.fun_controlnet.z_image import ZImageFunControl from invokeai.backend.model_manager.checkpoint_prefix import CheckpointPrefix from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Diffusers_Config_Base -from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_ZImage_Config from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.configs.main import ( Main_Checkpoint_ZImage_Config, @@ -83,7 +83,7 @@ from invokeai.backend.quantization.sdnq.loaders import raise_on_incomplete_sdnq_load, sdnq_sd_loader from invokeai.backend.qwen3.qwen3_tokenizer import load_bundled_qwen3_tokenizer from invokeai.backend.util.devices import TorchDevice -from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras, log_unexpected_keys +from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras def _remap_z_image_layer_paths(layer_names: Any) -> dict[str, list[str]]: @@ -1016,115 +1016,30 @@ def _load_model( type=ModelType.ControlNet, format=ModelFormat.Checkpoint, fp8_storage=Unimplemented( - "the adapter is loaded straight from the file and never cast at all; the Tile build is ~6.7 GB, " - "so this one is worth closing" + "the control is cast to the compute dtype at load and runs as its own module tree, so the layerwise cast " + "would fit; the Union and Tile 2.1 builds are 6.7 GB in bf16, so this one is worth closing" ), ) -class ZImageControlCheckpointModel(ModelLoader): - """Class to load Z-Image Control adapter models from safetensors checkpoint. +class ZImageFunControlNetLoader(ModelLoader): + """Loads a Z-Image Fun-ControlNet (Union 1.0, 2.0, 2.1, 2.1 Lite, Tile 2.1) from its safetensors file. - Z-Image Control models are standalone adapters containing control layers - (control_layers, control_all_x_embedder, control_noise_refiner) that can be - combined with a base ZImageTransformer2DModel at runtime for spatial conditioning - (Canny, HED, Depth, Pose, MLSD). + The file carries no config: `ZImageFunControl.from_state_dict` reads the width and the variant off the tensors. + Cast to the device's bfloat16-safe dtype, the policy the transformer loaders follow. """ - def _load_model( - self, - config: AnyModelConfig, - submodel_type: Optional[SubModelType] = None, - ) -> AnyModel: - if not isinstance(config, Checkpoint_Config_Base): - raise ValueError("Only CheckpointConfigBase models are supported here.") - - # ControlNet type models don't use submodel_type - load the adapter directly - return self._load_control_adapter(config) - - def _load_control_adapter( - self, - config: AnyModelConfig, - ) -> AnyModel: + def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> AnyModel: from safetensors.torch import load_file - from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter - - assert isinstance(config, ControlNet_Checkpoint_ZImage_Config) - model_path = Path(config.path) - - # Load the safetensors state dict - sd = load_file(model_path) - # Before the geometry probe, because the shapes it reads are meaningless on a packed weight. - reject_quantized_side_channel(sd, f"Z-Image ControlNet checkpoint {model_path.name}") - - # Determine number of control blocks from state dict - # Control blocks are named control_layers.0, control_layers.1, etc. - control_block_indices = set() - for key in sd.keys(): - if key.startswith("control_layers."): - parts = key.split(".") - if len(parts) > 1 and parts[1].isdigit(): - 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=control_in_dim, - all_patch_size=(2,), - all_f_patch_size=(1,), - dim=dim, - n_refiner_layers=n_refiner_layers, - n_heads=30, - n_kv_heads=30, - norm_eps=1e-05, - qk_norm=True, - ) - - # Load state dict with strict=False to handle missing keys like x_pad_token - # Some control adapters may not include x_pad_token in their checkpoint - missing_keys, unexpected_keys = model.load_state_dict(sd, assign=True, strict=False) - log_unexpected_keys("Z-Image ControlNet checkpoint", unexpected_keys) - - # Initialize x_pad_token if it was missing from the checkpoint - if "x_pad_token" in missing_keys: - import torch.nn as nn - - model.x_pad_token = nn.Parameter(torch.empty(dim)) - nn.init.normal_(model.x_pad_token, std=0.02) - - return model + if not isinstance(config, Checkpoint_Config_Base): + raise ValueError("Only checkpoint configs are supported here.") + path = Path(config.path) + sd = load_file(path) + # Before the layout is read: the shapes mean nothing on a packed weight, and the tensors are assigned as they + # are, so no scale would be folded in. alibaba-pai publishes dense bf16 files only. + reject_quantized_side_channel(sd, f"Z-Image ControlNet checkpoint {path.name}") + control = ZImageFunControl.from_state_dict(sd, dtype=TorchDevice.choose_bfloat16_safe_dtype(self._torch_device)) + self._logger.info(f"Z-Image ControlNet {path.name}: the {control.variant.name} layout") + return control def _fold_comfy_scaled_weights(sd: dict[str, Any], dtype: torch.dtype) -> int: diff --git a/invokeai/backend/model_manager/starter_models/z_image.py b/invokeai/backend/model_manager/starter_models/z_image.py index 7f7ae9a42e3..3d7d12f9b92 100644 --- a/invokeai/backend/model_manager/starter_models/z_image.py +++ b/invokeai/backend/model_manager/starter_models/z_image.py @@ -63,7 +63,7 @@ name="Z-Image ControlNet Union", base=BaseModelType.ZImage, source="https://huggingface.co/alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.1/resolve/main/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.safetensors", - description="Unified ControlNet for Z-Image Turbo supporting Canny, HED, Depth, Pose, MLSD, and Inpainting modes.", + description="Unified ControlNet for Z-Image Turbo supporting Canny, HED, Depth, Pose and MLSD control images.", type=ModelType.ControlNet, ) diff --git a/invokeai/backend/z_image/__init__.py b/invokeai/backend/z_image/__init__.py index 7fe48dd2cb2..556841e7471 100644 --- a/invokeai/backend/z_image/__init__.py +++ b/invokeai/backend/z_image/__init__.py @@ -1,16 +1 @@ # Z-Image backend utilities -from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter -from invokeai.backend.z_image.z_image_control_transformer import ZImageControlTransformer2DModel -from invokeai.backend.z_image.z_image_controlnet_extension import ( - ZImageControlNetExtension, - z_image_forward_with_control, -) -from invokeai.backend.z_image.z_image_patchify_utils import patchify_control_context - -__all__ = [ - "ZImageControlAdapter", - "ZImageControlTransformer2DModel", - "ZImageControlNetExtension", - "z_image_forward_with_control", - "patchify_control_context", -] diff --git a/invokeai/backend/z_image/z_image_control_adapter.py b/invokeai/backend/z_image/z_image_control_adapter.py deleted file mode 100644 index b8af273232d..00000000000 --- a/invokeai/backend/z_image/z_image_control_adapter.py +++ /dev/null @@ -1,236 +0,0 @@ -# Adapted from https://github.com/aigc-apps/VideoX-Fun/blob/main/videox_fun/models/z_image_transformer2d_control.py -# Copyright (c) Alibaba, Inc. and its affiliates. -# Apache License 2.0 - -""" -Z-Image Control Adapter for InvokeAI. - -This module provides a standalone control adapter that can be combined with -a base ZImageTransformer2DModel at runtime. The adapter contains only the -control-specific layers (control_layers, control_all_x_embedder, control_noise_refiner). -""" - -from typing import List, Optional - -import torch -import torch.nn as nn -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.modeling_utils import ModelMixin -from diffusers.models.transformers.transformer_z_image import ( - SEQ_MULTI_OF, - ZImageTransformerBlock, -) -from torch.nn.utils.rnn import pad_sequence - - -class ZImageControlTransformerBlock(ZImageTransformerBlock): - """Control-specific transformer block with skip connections for hint generation.""" - - def __init__( - self, - layer_id: int, - dim: int, - n_heads: int, - n_kv_heads: int, - norm_eps: float, - qk_norm: bool, - modulation: bool = True, - block_id: int = 0, - ): - super().__init__(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation) - self.block_id = block_id - if block_id == 0: - self.before_proj = nn.Linear(dim, dim) - nn.init.zeros_(self.before_proj.weight) - nn.init.zeros_(self.before_proj.bias) - self.after_proj = nn.Linear(dim, dim) - nn.init.zeros_(self.after_proj.weight) - nn.init.zeros_(self.after_proj.bias) - - def forward( - self, - c: torch.Tensor, - x: torch.Tensor, - attn_mask: torch.Tensor, - freqs_cis: torch.Tensor, - adaln_input: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Advance the control state by one block; returns ``(hint, control_state)``. - - Upstream stacks every hint onto the running state and copies that stack in each block. Every copy is - larger than any block the caching allocator has freed, so reserved memory grew by (n+1)(n+2)/2 - hint-sized tensors -- measured 4.5 GiB for a 15-block adapter at 1024px, 0.6 GiB with the hints kept apart. - """ - if self.block_id == 0: - c = self.before_proj(c) + x - c = super().forward(c, attn_mask=attn_mask, freqs_cis=freqs_cis, adaln_input=adaln_input) - return self.after_proj(c), c - - -class ZImageControlAdapter(ModelMixin, ConfigMixin): - """Standalone Z-Image Control Adapter. - - This adapter contains only the control-specific layers and can be combined - with a base ZImageTransformer2DModel at runtime. It computes control hints - that are added to the transformer's hidden states. - - The adapter supports 5 control modes: Canny, HED, Depth, Pose, MLSD. - Recommended control_context_scale: 0.65-0.80. - """ - - @register_to_config - def __init__( - self, - num_control_blocks: int = 6, # Number of control layer blocks - control_in_dim: int = 16, - all_patch_size: tuple[int, ...] = (2,), - all_f_patch_size: tuple[int, ...] = (1,), - dim: int = 3840, - n_refiner_layers: int = 2, - n_heads: int = 30, - n_kv_heads: int = 30, - norm_eps: float = 1e-5, - qk_norm: bool = True, - ): - super().__init__() - - self.dim = dim - self.control_in_dim = control_in_dim - self.all_patch_size = all_patch_size - self.all_f_patch_size = all_f_patch_size - - # Control patch embeddings - all_x_embedder = {} - for patch_size, f_patch_size in zip(all_patch_size, all_f_patch_size, strict=True): - x_embedder = nn.Linear( - f_patch_size * patch_size * patch_size * control_in_dim, - dim, - bias=True, - ) - all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder - - self.control_all_x_embedder = nn.ModuleDict(all_x_embedder) - - # Control noise refiner - self.control_noise_refiner = nn.ModuleList( - [ - ZImageTransformerBlock( - 1000 + layer_id, - dim, - n_heads, - n_kv_heads, - norm_eps, - qk_norm, - modulation=True, - ) - for layer_id in range(n_refiner_layers) - ] - ) - - # Control transformer blocks - self.control_layers = nn.ModuleList( - [ - ZImageControlTransformerBlock( - i, - dim, - n_heads, - n_kv_heads, - norm_eps, - qk_norm, - block_id=i, - ) - for i in range(num_control_blocks) - ] - ) - - # Padding token for control context - self.x_pad_token = nn.Parameter(torch.empty(dim)) - nn.init.normal_(self.x_pad_token, std=0.02) - - def forward( - self, - control_context: List[torch.Tensor], - unified_hidden_states: torch.Tensor, - cap_feats: torch.Tensor, - timestep_emb: torch.Tensor, - attn_mask: torch.Tensor, - freqs_cis: torch.Tensor, - rope_embedder, - patchify_fn, - patch_size: int = 2, - f_patch_size: int = 1, - ) -> tuple[torch.Tensor, ...]: - """Compute control hints from control context. - - Args: - control_context: List of control image latents [C, 1, H, W] - unified_hidden_states: Combined image+caption embeddings from main path - cap_feats: Caption feature embeddings - timestep_emb: Timestep embeddings - attn_mask: Attention mask - freqs_cis: RoPE frequencies - rope_embedder: RoPE embedder from base model - patchify_fn: Patchify function from base model - patch_size: Spatial patch size - f_patch_size: Frame patch size - - Returns: - Tuple of hint tensors to be added at each control layer position - """ - bsz = len(control_context) - device = control_context[0].device - - # Patchify control context using base model's patchify - ( - control_context_patches, - x_size, - x_pos_ids, - x_inner_pad_mask, - ) = patchify_fn(control_context, patch_size, f_patch_size, cap_feats.size(1)) - - # Embed control context - x_item_seqlens = [len(_) for _ in control_context_patches] - assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) - x_max_item_seqlen = max(x_item_seqlens) - - control_context_cat = torch.cat(control_context_patches, dim=0) - control_context_cat = self.control_all_x_embedder[f"{patch_size}-{f_patch_size}"](control_context_cat) - - # Match timestep dtype - adaln_input = timestep_emb.type_as(control_context_cat) - control_context_cat[torch.cat(x_inner_pad_mask)] = self.x_pad_token - control_context_list = list(control_context_cat.split(x_item_seqlens, dim=0)) - x_freqs_cis = list(rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) - - control_context_padded = pad_sequence(control_context_list, batch_first=True, padding_value=0.0) - x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) - x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(x_item_seqlens): - x_attn_mask[i, :seq_len] = 1 - - # Refine control context - for layer in self.control_noise_refiner: - control_context_padded = layer(control_context_padded, x_attn_mask, x_freqs_cis, adaln_input) - - # Unify with caption features - cap_item_seqlens = [cap_feats.size(1)] * bsz - control_context_unified = [] - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - control_context_unified.append(torch.cat([control_context_padded[i][:x_len], cap_feats[i][:cap_len]])) - control_context_unified = pad_sequence(control_context_unified, batch_first=True, padding_value=0.0) - c = control_context_unified - - # Process through control layers - hints = [] - for layer in self.control_layers: - hint, c = layer( - c, - x=unified_hidden_states, - attn_mask=attn_mask, - freqs_cis=freqs_cis, - adaln_input=adaln_input, - ) - hints.append(hint) - return tuple(hints) diff --git a/invokeai/backend/z_image/z_image_control_transformer.py b/invokeai/backend/z_image/z_image_control_transformer.py deleted file mode 100644 index ab64c64582f..00000000000 --- a/invokeai/backend/z_image/z_image_control_transformer.py +++ /dev/null @@ -1,643 +0,0 @@ -# Adapted from https://github.com/aigc-apps/VideoX-Fun/blob/main/videox_fun/models/z_image_transformer2d_control.py -# Copyright (c) Alibaba, Inc. and its affiliates. -# Apache License 2.0 - -""" -Z-Image Control Transformer for InvokeAI. - -This module provides the ZImageControlTransformer2DModel which extends the base -ZImageTransformer2DModel with control conditioning capabilities (Canny, HED, Depth, Pose, MLSD). -""" - -from typing import Any, Dict, List, Optional - -import torch -import torch.nn as nn -from diffusers.configuration_utils import register_to_config -from diffusers.models.transformers.transformer_z_image import ( - SEQ_MULTI_OF, - ZImageTransformer2DModel, - ZImageTransformerBlock, -) -from diffusers.utils import is_torch_version -from torch.nn.utils.rnn import pad_sequence - - -class ZImageControlTransformerBlock(ZImageTransformerBlock): - """Control-specific transformer block with skip connections for hint generation. - - This block extends ZImageTransformerBlock with before_proj and after_proj layers - that create skip connections for the control signal. The hints are accumulated - across blocks and used to condition the main transformer. - """ - - def __init__( - self, - layer_id: int, - dim: int, - n_heads: int, - n_kv_heads: int, - norm_eps: float, - qk_norm: bool, - modulation: bool = True, - block_id: int = 0, - ): - super().__init__(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation) - self.block_id = block_id - if block_id == 0: - self.before_proj = nn.Linear(dim, dim) - nn.init.zeros_(self.before_proj.weight) - nn.init.zeros_(self.before_proj.bias) - self.after_proj = nn.Linear(dim, dim) - nn.init.zeros_(self.after_proj.weight) - nn.init.zeros_(self.after_proj.bias) - - def forward( - self, - c: torch.Tensor, - x: torch.Tensor, - attn_mask: torch.Tensor, - freqs_cis: torch.Tensor, - adaln_input: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if self.block_id == 0: - c = self.before_proj(c) + x - all_c: list[torch.Tensor] = [] - else: - all_c = list(torch.unbind(c)) - c = all_c.pop(-1) - - c = super().forward(c, attn_mask=attn_mask, freqs_cis=freqs_cis, adaln_input=adaln_input) - c_skip = self.after_proj(c) - all_c += [c_skip, c] - c = torch.stack(all_c) - return c - - -class BaseZImageTransformerBlock(ZImageTransformerBlock): - """Modified transformer block that accepts control hints. - - This block extends ZImageTransformerBlock to add control hints to the - hidden states at specific positions in the network. - """ - - def __init__( - self, - layer_id: int, - dim: int, - n_heads: int, - n_kv_heads: int, - norm_eps: float, - qk_norm: bool, - modulation: bool = True, - block_id: Optional[int] = 0, - ): - super().__init__(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation) - self.block_id = block_id - - def forward( - self, - hidden_states: torch.Tensor, - attn_mask: torch.Tensor, - freqs_cis: torch.Tensor, - adaln_input: Optional[torch.Tensor] = None, - hints: Optional[tuple[torch.Tensor, ...]] = None, - context_scale: float = 1.0, - ) -> torch.Tensor: - hidden_states = super().forward( - hidden_states, - attn_mask=attn_mask, - freqs_cis=freqs_cis, - adaln_input=adaln_input, - ) - if self.block_id is not None and hints is not None: - hidden_states = hidden_states + hints[self.block_id] * context_scale - return hidden_states - - -class ZImageControlTransformer2DModel(ZImageTransformer2DModel): - """Z-Image Control Transformer for spatial conditioning. - - This model extends ZImageTransformer2DModel with control layers that process - a control image (e.g., Canny edges, depth map) and inject control signals - into the main transformer at every other layer. - - The control model supports 5 modes: Canny, HED, Depth, Pose, MLSD. - Recommended control_context_scale: 0.65-0.80. - - Args: - control_layers_places: List of layer indices where control is applied. - Defaults to every other layer [0, 2, 4, ...]. - control_in_dim: Input dimension for control context. Defaults to in_channels. - All other args are passed to ZImageTransformer2DModel. - """ - - @register_to_config - def __init__( - self, - control_layers_places: Optional[List[int]] = None, - control_in_dim: Optional[int] = None, - all_patch_size: tuple[int, ...] = (2,), - all_f_patch_size: tuple[int, ...] = (1,), - in_channels: int = 16, - dim: int = 3840, - n_layers: int = 30, - n_refiner_layers: int = 2, - n_heads: int = 30, - n_kv_heads: int = 30, - norm_eps: float = 1e-5, - qk_norm: bool = True, - cap_feat_dim: int = 2560, - rope_theta: float = 256.0, - t_scale: float = 1000.0, - axes_dims: tuple[int, ...] = (32, 48, 48), - axes_lens: tuple[int, ...] = (1024, 512, 512), - ): - super().__init__( - all_patch_size=all_patch_size, - all_f_patch_size=all_f_patch_size, - in_channels=in_channels, - dim=dim, - n_layers=n_layers, - n_refiner_layers=n_refiner_layers, - n_heads=n_heads, - n_kv_heads=n_kv_heads, - norm_eps=norm_eps, - qk_norm=qk_norm, - cap_feat_dim=cap_feat_dim, - rope_theta=rope_theta, - t_scale=t_scale, - axes_dims=axes_dims, - axes_lens=axes_lens, - ) - - # Control layer configuration - self.control_layers_places = ( - list(range(0, n_layers, 2)) if control_layers_places is None else control_layers_places - ) - self.control_in_dim = in_channels if control_in_dim is None else control_in_dim - - assert 0 in self.control_layers_places - self.control_layers_mapping = {i: n for n, i in enumerate(self.control_layers_places)} - - # Replace standard layers with control-aware layers - del self.layers - self.layers = nn.ModuleList( - [ - BaseZImageTransformerBlock( - i, - dim, - n_heads, - n_kv_heads, - norm_eps, - qk_norm, - block_id=self.control_layers_mapping[i] if i in self.control_layers_places else None, - ) - for i in range(n_layers) - ] - ) - - # Control transformer blocks - self.control_layers = nn.ModuleList( - [ - ZImageControlTransformerBlock( - i, - dim, - n_heads, - n_kv_heads, - norm_eps, - qk_norm, - block_id=i, - ) - for i in range(len(self.control_layers_places)) - ] - ) - - # Control patch embeddings - all_x_embedder = {} - for patch_size, f_patch_size in zip(all_patch_size, all_f_patch_size, strict=True): - x_embedder = nn.Linear( - f_patch_size * patch_size * patch_size * self.control_in_dim, - dim, - bias=True, - ) - all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder - - self.control_all_x_embedder = nn.ModuleDict(all_x_embedder) - - # Control noise refiner - self.control_noise_refiner = nn.ModuleList( - [ - ZImageTransformerBlock( - 1000 + layer_id, - dim, - n_heads, - n_kv_heads, - norm_eps, - qk_norm, - modulation=True, - ) - for layer_id in range(n_refiner_layers) - ] - ) - - def patchify( - self, - all_image: List[torch.Tensor], - patch_size: int, - f_patch_size: int, - cap_seq_len: int, - ) -> tuple[List[torch.Tensor], List[tuple], List[torch.Tensor], List[torch.Tensor]]: - """Patchify images without embedding. - - This method extracts patches from images for control context processing. - Unlike patchify_and_embed, this only processes images without caption features. - - Args: - all_image: List of image tensors [C, F, H, W] - patch_size: Spatial patch size (height and width) - f_patch_size: Frame patch size - cap_seq_len: Caption sequence length (for position ID offset) - - Returns: - Tuple of: - - all_image_out: List of patchified image tensors - - all_image_size: List of (F, H, W) tuples - - all_image_pos_ids: List of position ID tensors - - all_image_pad_mask: List of padding mask tensors - """ - pH = pW = patch_size - pF = f_patch_size - device = all_image[0].device - - all_image_out = [] - all_image_size = [] - all_image_pos_ids = [] - all_image_pad_mask = [] - - # Calculate padded caption length for position offset - cap_padding_len = (-cap_seq_len) % SEQ_MULTI_OF - cap_padded_len = cap_seq_len + cap_padding_len - - for image in all_image: - C, F, H, W = image.size() - all_image_size.append((F, H, W)) - F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW - - # Patchify: [C, F, H, W] -> [(F*H*W)/(patch), patch_elements * C] - image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) - image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C) - - image_ori_len = len(image) - image_padding_len = (-image_ori_len) % SEQ_MULTI_OF - - # Create position IDs - image_ori_pos_ids = self.create_coordinate_grid( - size=(F_tokens, H_tokens, W_tokens), - start=(cap_padded_len + 1, 0, 0), - device=device, - ).flatten(0, 2) - image_padding_pos_ids = ( - self.create_coordinate_grid( - size=(1, 1, 1), - start=(0, 0, 0), - device=device, - ) - .flatten(0, 2) - .repeat(image_padding_len, 1) - ) - image_padded_pos_ids = torch.cat([image_ori_pos_ids, image_padding_pos_ids], dim=0) - all_image_pos_ids.append(image_padded_pos_ids) - - # Padding mask - all_image_pad_mask.append( - torch.cat( - [ - torch.zeros((image_ori_len,), dtype=torch.bool, device=device), - torch.ones((image_padding_len,), dtype=torch.bool, device=device), - ], - dim=0, - ) - ) - - # Padded feature - image_padded_feat = torch.cat([image, image[-1:].repeat(image_padding_len, 1)], dim=0) - all_image_out.append(image_padded_feat) - - return all_image_out, all_image_size, all_image_pos_ids, all_image_pad_mask - - def forward_control( - self, - x: torch.Tensor, - cap_feats: torch.Tensor, - control_context: List[torch.Tensor], - kwargs: Dict[str, Any], - t: torch.Tensor, - patch_size: int = 2, - f_patch_size: int = 1, - ) -> tuple[torch.Tensor, ...]: - """Process control context and generate hints for the main transformer. - - Args: - x: Unified image+caption embeddings from main path - cap_feats: Caption feature embeddings - control_context: List of control images (VAE-encoded latents) - kwargs: Additional kwargs including attn_mask, freqs_cis - t: Timestep embeddings - patch_size: Spatial patch size - f_patch_size: Frame patch size - - Returns: - Tuple of hint tensors to be added at each control layer position - """ - bsz = len(control_context) - device = control_context[0].device - - # Patchify control context - ( - control_context_patches, - x_size, - x_pos_ids, - x_inner_pad_mask, - ) = self.patchify(control_context, patch_size, f_patch_size, cap_feats.size(1)) - - # Embed control context - x_item_seqlens = [len(_) for _ in control_context_patches] - assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) - x_max_item_seqlen = max(x_item_seqlens) - - control_context_cat = torch.cat(control_context_patches, dim=0) - control_context_cat = self.control_all_x_embedder[f"{patch_size}-{f_patch_size}"](control_context_cat) - - # Match t_embedder output dtype - adaln_input = t.type_as(control_context_cat) - control_context_cat[torch.cat(x_inner_pad_mask)] = self.x_pad_token - control_context_list = list(control_context_cat.split(x_item_seqlens, dim=0)) - x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) - - control_context_padded = pad_sequence(control_context_list, batch_first=True, padding_value=0.0) - x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) - x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(x_item_seqlens): - x_attn_mask[i, :seq_len] = 1 - - # Refine control context - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer in self.control_noise_refiner: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - control_context_padded = torch.utils.checkpoint.checkpoint( - create_custom_forward(layer), - control_context_padded, - x_attn_mask, - x_freqs_cis, - adaln_input, - **ckpt_kwargs, - ) - else: - for layer in self.control_noise_refiner: - control_context_padded = layer(control_context_padded, x_attn_mask, x_freqs_cis, adaln_input) - - # Unify with caption features - cap_item_seqlens = [cap_feats.size(1)] * bsz # Assume same length for batch - control_context_unified = [] - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - control_context_unified.append(torch.cat([control_context_padded[i][:x_len], cap_feats[i][:cap_len]])) - control_context_unified = pad_sequence(control_context_unified, batch_first=True, padding_value=0.0) - c = control_context_unified - - # Process through control layers - for layer in self.control_layers: - if torch.is_grad_enabled() and self.gradient_checkpointing: - - def create_custom_forward(module, **static_kwargs): - def custom_forward(*inputs): - return module(*inputs, **static_kwargs) - - return custom_forward - - ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - c = torch.utils.checkpoint.checkpoint( - create_custom_forward( - layer, - x=x, - attn_mask=kwargs["attn_mask"], - freqs_cis=kwargs["freqs_cis"], - adaln_input=kwargs["adaln_input"], - ), - c, - **ckpt_kwargs, - ) - else: - c = layer( - c, - x=x, - attn_mask=kwargs["attn_mask"], - freqs_cis=kwargs["freqs_cis"], - adaln_input=kwargs["adaln_input"], - ) - - hints = torch.unbind(c)[:-1] - return hints - - def forward( - self, - x: List[torch.Tensor], - t: torch.Tensor, - cap_feats: List[torch.Tensor], - patch_size: int = 2, - f_patch_size: int = 1, - control_context: Optional[List[torch.Tensor]] = None, - control_context_scale: float = 1.0, - ) -> tuple[List[torch.Tensor], dict]: - """Forward pass with control conditioning. - - Args: - x: List of image tensors [B, C, 1, H, W] - t: Timestep tensor - cap_feats: List of caption feature tensors - patch_size: Spatial patch size (default 2) - f_patch_size: Frame patch size (default 1) - control_context: List of control image latents (VAE-encoded) - control_context_scale: Strength of control signal (0.65-0.80 recommended) - - Returns: - Tuple of (output tensors, empty dict) - """ - assert patch_size in self.all_patch_size - assert f_patch_size in self.all_f_patch_size - - if control_context is None: - # Fall back to base model behavior without control - return super().forward(x, t, cap_feats, patch_size, f_patch_size) - - bsz = len(x) - device = x[0].device - t = t * self.t_scale - t = self.t_embedder(t) - - ( - x, - cap_feats, - x_size, - x_pos_ids, - cap_pos_ids, - x_inner_pad_mask, - cap_inner_pad_mask, - ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) - - # Image embedding and refinement - x_item_seqlens = [len(_) for _ in x] - assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) - x_max_item_seqlen = max(x_item_seqlens) - - x = torch.cat(x, dim=0) - x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) - - adaln_input = t.type_as(x) - x[torch.cat(x_inner_pad_mask)] = self.x_pad_token - x = list(x.split(x_item_seqlens, dim=0)) - x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) - - x = pad_sequence(x, batch_first=True, padding_value=0.0) - x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) - x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(x_item_seqlens): - x_attn_mask[i, :seq_len] = 1 - - # Noise refiner - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer in self.noise_refiner: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - x = torch.utils.checkpoint.checkpoint( - create_custom_forward(layer), - x, - x_attn_mask, - x_freqs_cis, - adaln_input, - **ckpt_kwargs, - ) - else: - for layer in self.noise_refiner: - x = layer(x, x_attn_mask, x_freqs_cis, adaln_input) - - # Caption embedding and refinement - cap_item_seqlens = [len(_) for _ in cap_feats] - assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) - cap_max_item_seqlen = max(cap_item_seqlens) - - cap_feats = torch.cat(cap_feats, dim=0) - cap_feats = self.cap_embedder(cap_feats) - cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token - cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) - cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0)) - - cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0) - cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) - cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(cap_item_seqlens): - cap_attn_mask[i, :seq_len] = 1 - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer in self.context_refiner: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - cap_feats = torch.utils.checkpoint.checkpoint( - create_custom_forward(layer), - cap_feats, - cap_attn_mask, - cap_freqs_cis, - **ckpt_kwargs, - ) - else: - for layer in self.context_refiner: - cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis) - - # Unified processing - unified = [] - unified_freqs_cis = [] - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) - unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]])) - unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens, strict=True)] - unified_max_item_seqlen = max(unified_item_seqlens) - - unified = pad_sequence(unified, batch_first=True, padding_value=0.0) - unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) - unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - unified_attn_mask[i, :seq_len] = 1 - - # Generate control hints - kwargs = { - "attn_mask": unified_attn_mask, - "freqs_cis": unified_freqs_cis, - "adaln_input": adaln_input, - } - hints = self.forward_control( - unified, - cap_feats, - control_context, - kwargs, - t=t, - patch_size=patch_size, - f_patch_size=f_patch_size, - ) - - # Main transformer with control hints - for layer in self.layers: - layer_kwargs = { - "attn_mask": unified_attn_mask, - "freqs_cis": unified_freqs_cis, - "adaln_input": adaln_input, - "hints": hints, - "context_scale": control_context_scale, - } - if torch.is_grad_enabled() and self.gradient_checkpointing: - - def create_custom_forward(module, **static_kwargs): - def custom_forward(*inputs): - return module(*inputs, **static_kwargs) - - return custom_forward - - ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - - unified = torch.utils.checkpoint.checkpoint( - create_custom_forward(layer, **layer_kwargs), - unified, - **ckpt_kwargs, - ) - else: - unified = layer(unified, **layer_kwargs) - - # Final layer and unpatchify - unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input) - unified = list(unified.unbind(dim=0)) - x = self.unpatchify(unified, x_size, patch_size, f_patch_size) - - x = torch.stack(x) - return x, {} diff --git a/invokeai/backend/z_image/z_image_controlnet_extension.py b/invokeai/backend/z_image/z_image_controlnet_extension.py deleted file mode 100644 index aa037afcf24..00000000000 --- a/invokeai/backend/z_image/z_image_controlnet_extension.py +++ /dev/null @@ -1,407 +0,0 @@ -"""Z-Image ControlNet Extension for spatial conditioning. - -This module provides an extension-based approach to Z-Image ControlNet, -similar to how FLUX ControlNet works. Instead of duplicating the entire -transformer, we compute control hints separately and inject them into -the base transformer's forward pass. -""" - -import logging -from typing import List, Optional, Tuple - -import torch -from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel -from torch.nn.utils.rnn import pad_sequence - -from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter -from invokeai.backend.z_image.z_image_patchify_utils import SEQ_MULTI_OF, patchify_control_context - -logger = logging.getLogger(__name__) - - -class ZImageControlNetExtension: - """Extension for Z-Image ControlNet - computes control hints without duplicating the transformer. - - This class follows the same pattern as FLUX ControlNet extensions: - - The control adapter is loaded separately - - Control hints are computed per step - - Hints are injected into the transformer's layer outputs - - Attributes: - control_adapter: The Z-Image control adapter model - control_cond: VAE-encoded control image latents - weight: Control strength (recommended: 0.65-0.80) - begin_step_percent: When to start applying control (0.0 = start) - end_step_percent: When to stop applying control (1.0 = end) - """ - - def __init__( - self, - control_adapter: ZImageControlAdapter, - control_cond: torch.Tensor, - weight: float = 0.75, - begin_step_percent: float = 0.0, - end_step_percent: float = 1.0, - skip_layers: int = 0, # Skip first N control injection layers - ): - self._adapter = control_adapter - self._control_cond = control_cond - self._weight = weight - self._begin_step_percent = begin_step_percent - self._end_step_percent = end_step_percent - self._skip_layers = skip_layers - - # Get actual number of control blocks from loaded model (not config!) - # The safetensors may have more blocks than the config suggests - self._num_control_blocks = len(control_adapter.control_layers) - - # Control layers are applied at every other layer (0, 2, 4, ...) - # This matches the default configuration in the original implementation - self._control_places = [i * 2 for i in range(self._num_control_blocks)] - - logger.debug("Actual num_control_blocks: %s", self._num_control_blocks) - logger.debug("control_places: %s", self._control_places) - - first_layer = control_adapter.control_layers[0] - if hasattr(first_layer, "after_proj"): - after_proj_norm = first_layer.after_proj.weight.norm().item() - logger.debug("First control layer after_proj weight norm: %s", after_proj_norm) - if after_proj_norm < 1e-6: - logger.warning("after_proj weights are near-zero! Weights may not be loaded correctly.") - - @property - def weight(self) -> float: - return self._weight - - @property - def control_places(self) -> List[int]: - return self._control_places - - def should_apply(self, step_index: int, total_steps: int) -> bool: - """Check if control should be applied at this step.""" - if total_steps == 0: - return True - step_percent = step_index / total_steps - return self._begin_step_percent <= step_percent <= self._end_step_percent - - def compute_hints( - self, - base_transformer: ZImageTransformer2DModel, - unified_hidden_states: torch.Tensor, - cap_feats: torch.Tensor, - timestep_emb: torch.Tensor, - attn_mask: torch.Tensor, - freqs_cis: torch.Tensor, - x_item_seqlens: List[int], - cap_item_seqlens: List[int], - x_freqs_cis: torch.Tensor, - patch_size: int = 2, - f_patch_size: int = 1, - ) -> Tuple[torch.Tensor, ...]: - """Compute control hints using the adapter. - - This method processes the control condition through the adapter's - control_noise_refiner and control_layers to produce hints that - will be added to the transformer's hidden states. - - Args: - base_transformer: The base Z-Image transformer (for rope_embedder) - unified_hidden_states: Combined image+caption hidden states - cap_feats: Caption feature embeddings (padded) - timestep_emb: Timestep embeddings (adaln_input) - attn_mask: Unified attention mask - freqs_cis: RoPE frequencies - x_item_seqlens: Image sequence lengths per batch item - cap_item_seqlens: Caption sequence lengths per batch item - patch_size: Spatial patch size - f_patch_size: Frame patch size - - Returns: - Tuple of hint tensors to add at each control layer position - """ - # control_cond is always [C, F, H, W] format (single control image) - # where C = control_in_dim (16 for V1, 33 for V2.0), F = 1 frame - bsz = 1 - device = self._control_cond.device - - # Wrap control_cond in a list for patchify_control_context - # Expected input: List of [C, F, H, W] tensors - control_context = [self._control_cond] - - # Patchify control context - # Note: We don't use control_pos_ids anymore - we use x_freqs_cis from main path instead - ( - control_patches, - _, - _control_pos_ids, # Not used - we use main path's position encoding - control_pad_mask, - ) = patchify_control_context( - control_context, - patch_size, - f_patch_size, - cap_feats.size(1), - ) - - # Embed control context - ctrl_item_seqlens = [len(p) for p in control_patches] - assert all(s % SEQ_MULTI_OF == 0 for s in ctrl_item_seqlens) - ctrl_max_seqlen = max(ctrl_item_seqlens) - - control_cat = torch.cat(control_patches, dim=0) - embedder_key = f"{patch_size}-{f_patch_size}" - control_cat = self._adapter.control_all_x_embedder[embedder_key](control_cat) - - # Apply padding token (ensure dtype matches) - adaln_input = timestep_emb.type_as(control_cat) - x_pad_token = self._adapter.x_pad_token.to(dtype=control_cat.dtype) - control_cat[torch.cat(control_pad_mask)] = x_pad_token - - control_list = list(control_cat.split(ctrl_item_seqlens, dim=0)) - - control_padded = pad_sequence(control_list, batch_first=True, padding_value=0.0) - - # Use x_freqs_cis from main path for control patches (same spatial structure) - # This ensures control and image have aligned position encodings - ctrl_freqs_cis_for_refiner = x_freqs_cis[:, : control_padded.shape[1]] - - ctrl_attn_mask = torch.zeros((bsz, ctrl_max_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(ctrl_item_seqlens): - ctrl_attn_mask[i, :seq_len] = 1 - - # Refine control context through control_noise_refiner - # Using x_freqs_cis to match main path's position encoding - for layer in self._adapter.control_noise_refiner: - control_padded = layer(control_padded, ctrl_attn_mask, ctrl_freqs_cis_for_refiner, adaln_input) - - # Unify control with caption features - control_unified = [] - for i in range(bsz): - ctrl_len = ctrl_item_seqlens[i] - cap_len = cap_item_seqlens[i] - control_unified.append(torch.cat([control_padded[i][:ctrl_len], cap_feats[i][:cap_len]])) - - control_unified = pad_sequence(control_unified, batch_first=True, padding_value=0.0) - c = control_unified - - if not hasattr(self, "_debug_printed"): - self._debug_printed = True - logger.debug("control_unified shape: %s", control_unified.shape) - logger.debug("unified_hidden_states shape: %s", unified_hidden_states.shape) - logger.debug("ctrl_item_seqlens: %s, x_item_seqlens: %s", ctrl_item_seqlens, x_item_seqlens) - - layer0 = self._adapter.control_layers[0] - if hasattr(layer0, "before_proj"): - logger.debug("before_proj weight norm: %.6f", layer0.before_proj.weight.norm().item()) - if hasattr(layer0, "after_proj"): - logger.debug("after_proj weight norm: %.6f", layer0.after_proj.weight.norm().item()) - - if len(self._adapter.control_noise_refiner) > 0: - refiner0 = self._adapter.control_noise_refiner[0] - if hasattr(refiner0, "attn"): - logger.debug("noise_refiner[0] attn.wq norm: %.6f", refiner0.attn.wq.weight.norm().item()) - - hint_list = [] - for layer in self._adapter.control_layers: - hint, c = layer( - c, - x=unified_hidden_states, - attn_mask=attn_mask, - freqs_cis=freqs_cis, - adaln_input=adaln_input, - ) - hint_list.append(hint) - hints = tuple(hint_list) - - if not hasattr(self, "_hints_printed"): - self._hints_printed = True - logger.debug("Number of hints: %s", len(hints)) - if hints: - logger.debug("First hint shape: %s", hints[0].shape) - for i, h in enumerate(hints[:3]): # First 3 hints - logger.debug( - "Hint[%s] mean: %.6f, std: %.6f, min: %.6f, max: %.6f", - i, - h.mean().item(), - h.std().item(), - h.min().item(), - h.max().item(), - ) - - return hints - - -def z_image_forward_with_control( - transformer: ZImageTransformer2DModel, - x: List[torch.Tensor], - t: torch.Tensor, - cap_feats: List[torch.Tensor], - control_extension: Optional[ZImageControlNetExtension] = None, - patch_size: int = 2, - f_patch_size: int = 1, -) -> Tuple[List[torch.Tensor], dict]: - """Forward pass through Z-Image transformer with optional control injection. - - This function replicates the base transformer's forward pass but allows - injecting control hints at specific layer positions. It uses the base - transformer's weights directly without duplicating them. - - Args: - transformer: The base Z-Image transformer model - x: List of image tensors [C, F, H, W] - t: Timestep tensor - cap_feats: List of caption feature tensors - control_extension: Optional control extension for hint injection - patch_size: Spatial patch size (default: 2) - f_patch_size: Frame patch size (default: 1) - - Returns: - Tuple of (output tensors list, empty dict for compatibility) - """ - assert patch_size in transformer.all_patch_size - assert f_patch_size in transformer.all_f_patch_size - - bsz = len(x) - device = x[0].device - t_scaled = t * transformer.t_scale - t_emb = transformer.t_embedder(t_scaled) - - # Patchify and embed using base transformer's method - ( - x_patches, - cap_feats_patches, - x_size, - x_pos_ids, - cap_pos_ids, - x_inner_pad_mask, - cap_inner_pad_mask, - ) = transformer.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) - - # === X embed & refine === - x_item_seqlens = [len(p) for p in x_patches] - assert all(s % SEQ_MULTI_OF == 0 for s in x_item_seqlens) - x_max_item_seqlen = max(x_item_seqlens) - - embedder_key = f"{patch_size}-{f_patch_size}" - x_cat = torch.cat(x_patches, dim=0) - x_cat = transformer.all_x_embedder[embedder_key](x_cat) - - adaln_input = t_emb.type_as(x_cat) - x_cat[torch.cat(x_inner_pad_mask)] = transformer.x_pad_token - - x_list = list(x_cat.split(x_item_seqlens, dim=0)) - x_freqs_cis = list(transformer.rope_embedder(torch.cat(x_pos_ids, dim=0)).split([len(p) for p in x_pos_ids], dim=0)) - - x_padded = pad_sequence(x_list, batch_first=True, padding_value=0.0) - x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) - x_freqs_cis = x_freqs_cis[:, : x_padded.shape[1]] - - x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(x_item_seqlens): - x_attn_mask[i, :seq_len] = 1 - - # Noise refiner - for layer in transformer.noise_refiner: - x_padded = layer(x_padded, x_attn_mask, x_freqs_cis, adaln_input) - - # === Cap embed & refine === - cap_item_seqlens = [len(p) for p in cap_feats_patches] - cap_max_item_seqlen = max(cap_item_seqlens) - - cap_cat = torch.cat(cap_feats_patches, dim=0) - cap_cat = transformer.cap_embedder(cap_cat) - cap_cat[torch.cat(cap_inner_pad_mask)] = transformer.cap_pad_token - - cap_list = list(cap_cat.split(cap_item_seqlens, dim=0)) - cap_freqs_cis = list( - transformer.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split([len(p) for p in cap_pos_ids], dim=0) - ) - - cap_padded = pad_sequence(cap_list, batch_first=True, padding_value=0.0) - cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) - cap_freqs_cis = cap_freqs_cis[:, : cap_padded.shape[1]] - - cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(cap_item_seqlens): - cap_attn_mask[i, :seq_len] = 1 - - # Context refiner - for layer in transformer.context_refiner: - cap_padded = layer(cap_padded, cap_attn_mask, cap_freqs_cis) - - # === Unified === - unified = [] - unified_freqs_cis = [] - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - unified.append(torch.cat([x_padded[i][:x_len], cap_padded[i][:cap_len]])) - unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]])) - - unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens, strict=False)] - unified_max_item_seqlen = max(unified_item_seqlens) - - unified = pad_sequence(unified, batch_first=True, padding_value=0.0) - unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) - - unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - unified_attn_mask[i, :seq_len] = 1 - - # === Compute control hints if extension provided === - # IMPORTANT: Hints are computed ONCE using the INITIAL unified state (before main layers) - # This matches the original VideoX-Fun architecture - control_places: List[int] = [] - control_weight: float = 1.0 - hints: Optional[Tuple[torch.Tensor, ...]] = None - - if not hasattr(z_image_forward_with_control, "_layers_printed"): - z_image_forward_with_control._layers_printed = True - logger.debug("Base transformer has %s layers", len(transformer.layers)) - - if control_extension is not None: - # Compute ALL hints at once using the INITIAL unified state (before main layers run) - hints = control_extension.compute_hints( - base_transformer=transformer, - unified_hidden_states=unified, # INITIAL unified state! - cap_feats=cap_padded, - timestep_emb=adaln_input, - attn_mask=unified_attn_mask, - freqs_cis=unified_freqs_cis, - x_item_seqlens=x_item_seqlens, - cap_item_seqlens=cap_item_seqlens, - x_freqs_cis=x_freqs_cis, - patch_size=patch_size, - f_patch_size=f_patch_size, - ) - control_places = control_extension.control_places - control_weight = control_extension.weight - - # === Main transformer layers with pre-computed hint injection === - skip_layers = control_extension._skip_layers if control_extension is not None else 0 - control_layer_idx = 0 - for layer_idx, layer in enumerate(transformer.layers): - unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input) - - # Inject pre-computed control hint at designated positions - if hints is not None and layer_idx in control_places and control_layer_idx < len(hints): - # Skip first N hints if configured - if control_layer_idx >= skip_layers: - hint = hints[control_layer_idx] - - if not hasattr(z_image_forward_with_control, "_injection_printed"): - z_image_forward_with_control._injection_printed = True - logger.debug("Injection at layer %s (control_layer %s)", layer_idx, control_layer_idx) - logger.debug("Hint mean: %.6f, std: %.6f", hint.mean().item(), hint.std().item()) - logger.debug("Unified mean: %.6f, std: %.6f", unified.mean().item(), unified.std().item()) - logger.debug("control_weight: %s, skip_layers: %s", control_weight, skip_layers) - - unified = unified + hint * control_weight - - control_layer_idx += 1 - - # === Final layer and unpatchify === - unified = transformer.all_final_layer[embedder_key](unified, adaln_input) - unified = list(unified.unbind(dim=0)) - output = transformer.unpatchify(unified, x_size, patch_size, f_patch_size) - - return output, {} diff --git a/invokeai/backend/z_image/z_image_patchify_utils.py b/invokeai/backend/z_image/z_image_patchify_utils.py deleted file mode 100644 index ddbe99a3a82..00000000000 --- a/invokeai/backend/z_image/z_image_patchify_utils.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Utility functions for Z-Image patchify operations.""" - -from typing import List, Tuple - -import torch - -# Sequence must be multiple of this value (from diffusers transformer_z_image) -SEQ_MULTI_OF = 32 - - -def create_coordinate_grid( - size: Tuple[int, ...], - start: Tuple[int, ...] | None = None, - device: torch.device | None = None, -) -> torch.Tensor: - """Create a coordinate grid for position embeddings. - - Args: - size: Size of the grid (e.g., (F, H, W)) - start: Starting coordinates (default: all zeros) - device: Target device - - Returns: - Coordinate grid tensor of shape (*size, len(size)) - """ - if start is None: - start = tuple(0 for _ in size) - - axes = [ - torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size, strict=False) - ] - grids = torch.meshgrid(axes, indexing="ij") - return torch.stack(grids, dim=-1) - - -def patchify_control_context( - all_image: List[torch.Tensor], - patch_size: int, - f_patch_size: int, - cap_seq_len: int, -) -> Tuple[List[torch.Tensor], List[Tuple[int, int, int]], List[torch.Tensor], List[torch.Tensor]]: - """Patchify control images without embedding. - - This function extracts patches from control images for control context processing. - It handles padding and position ID creation for the control signal. - - Args: - all_image: List of control image tensors [C, F, H, W] - patch_size: Spatial patch size (height and width) - f_patch_size: Frame patch size - cap_seq_len: Caption sequence length (for position ID offset) - - Returns: - Tuple of: - - all_image_out: List of patchified image tensors - - all_image_size: List of (F, H, W) tuples - - all_image_pos_ids: List of position ID tensors - - all_image_pad_mask: List of padding mask tensors - """ - pH = pW = patch_size - pF = f_patch_size - device = all_image[0].device - - all_image_out: List[torch.Tensor] = [] - all_image_size: List[Tuple[int, int, int]] = [] - all_image_pos_ids: List[torch.Tensor] = [] - all_image_pad_mask: List[torch.Tensor] = [] - - # Calculate padded caption length for position offset - cap_padding_len = (-cap_seq_len) % SEQ_MULTI_OF - cap_padded_len = cap_seq_len + cap_padding_len - - for image in all_image: - C, F, H, W = image.size() - all_image_size.append((F, H, W)) - F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW - - # Patchify: [C, F, H, W] -> [(F_tokens*H_tokens*W_tokens), (pF*pH*pW*C)] - # Step 1: Rearrange to put spatial dims together for proper patching - # [C, F, H, W] -> [F, H, W, C] - image = image.permute(1, 2, 3, 0).contiguous() - - # Step 2: Split H and W into tokens and patch sizes - # [F, H, W, C] -> [F, H_tokens, pH, W_tokens, pW, C] - image = image.view(F, H_tokens, pH, W_tokens, pW, C) - - # Step 3: Rearrange to group patches and features - # [F, H_tokens, pH, W_tokens, pW, C] -> [F, H_tokens, W_tokens, pH, pW, C] - image = image.permute(0, 1, 3, 2, 4, 5).contiguous() - - # Step 4: For F > 1, we'd need to handle F similarly, but for F=1 this is simpler - # Final reshape: [F*H_tokens*W_tokens, pH*pW*C] - num_patches = F_tokens * H_tokens * W_tokens - patch_features = pF * pH * pW * C - image = image.reshape(num_patches, patch_features) - - image_ori_len = len(image) - image_padding_len = (-image_ori_len) % SEQ_MULTI_OF - - # Create position IDs - image_ori_pos_ids = create_coordinate_grid( - size=(F_tokens, H_tokens, W_tokens), - start=(cap_padded_len + 1, 0, 0), - device=device, - ).flatten(0, 2) - - image_padding_pos_ids = ( - create_coordinate_grid( - size=(1, 1, 1), - start=(0, 0, 0), - device=device, - ) - .flatten(0, 2) - .repeat(image_padding_len, 1) - ) - image_padded_pos_ids = torch.cat([image_ori_pos_ids, image_padding_pos_ids], dim=0) - all_image_pos_ids.append(image_padded_pos_ids) - - # Padding mask - all_image_pad_mask.append( - torch.cat( - [ - torch.zeros((image_ori_len,), dtype=torch.bool, device=device), - torch.ones((image_padding_len,), dtype=torch.bool, device=device), - ], - dim=0, - ) - ) - - # Padded feature - image_padded_feat = torch.cat([image, image[-1:].repeat(image_padding_len, 1)], dim=0) - all_image_out.append(image_padded_feat) - - return all_image_out, all_image_size, all_image_pos_ids, all_image_pad_mask diff --git a/scripts/calibrate_z_image_working_memory.py b/scripts/calibrate_z_image_working_memory.py index 7ed5c90fb97..a2e067ee43d 100644 --- a/scripts/calibrate_z_image_working_memory.py +++ b/scripts/calibrate_z_image_working_memory.py @@ -10,14 +10,17 @@ 2. the materialized score matrix, only where SDPA has no fused kernel for these shapes -- ROCm/Windows without the experimental AOTriton kernels, MPS -- which `sdpa_score_matrix_bytes` charges at the shipped `SDPA_MATH_BYTES_PER_SCORE_ELEMENT`; -3. with a ControlNet attached, the adapter's hints: one per control block, live next to the transformer's own - activations. +3. with a ControlNet attached, its hints: one per control block, live next to the transformer's own activations + until the last layer the control adds to. This measures all three on the real geometry (dim 3840, 30 heads of 128, caption features 2560 wide) with randomly initialized models and a reduced transformer block count: memory depends on shapes, not on weight values, and a -no-grad forward frees each block's intermediates. `--extra-blocks` re-checks that on the build at hand. Every point -runs in a fresh subprocess, so the caching allocator's history cannot contaminate the peak-reserved reading, and a -point that runs out of memory is recorded as such rather than aborting the run. +no-grad forward frees each block's intermediates. `--extra-blocks` re-checks that on the build at hand. The ControlNet +point runs the control through the denoise's hooks with its hints on consecutive main blocks: they are all computed as +the first main block starts and dropped after the last one they land on, so where in between they land does not move +the peak, and packing them keeps the transformer at one block per hint. Every point runs in a fresh subprocess, so the +caching allocator's history cannot contaminate the peak-reserved reading, and a point that runs out of memory is +recorded as such rather than aborting the run. Peak *reserved*, not allocated: the estimate is compared against free memory, so it has to bound what the allocator actually takes off the card. @@ -31,11 +34,12 @@ Usage ----- python scripts/calibrate_z_image_working_memory.py - python scripts/calibrate_z_image_working_memory.py --extra-blocks 4 --control-blocks 6 - python scripts/calibrate_z_image_working_memory.py --max-px 1024 --control-blocks 0 + python scripts/calibrate_z_image_working_memory.py --extra-blocks 4 --control-variant 2.1-lite + python scripts/calibrate_z_image_working_memory.py --max-px 1024 --control-variant none """ import argparse +import dataclasses import json import subprocess import sys @@ -51,6 +55,13 @@ ZImageDenoiseInvocation, _padded, ) +from invokeai.backend.fun_controlnet.z_image import ( + Z_IMAGE_FUN_CONTROL_1_0, + Z_IMAGE_FUN_CONTROL_2_0, + Z_IMAGE_FUN_CONTROL_2_1, + Z_IMAGE_FUN_CONTROL_2_1_LITE, + ZImageFunControlVariant, +) from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT GIB = 1024**3 @@ -58,7 +69,7 @@ CAP_FEAT_DIM = 2560 PATCH_SIZE = 2 -CONTROL_IN_DIM = 33 # V2.0 adapters; the channel count only sizes the patch embedder +REFINER_BLOCKS = 2 # One transformer block at full width: attention 4 dim^2, SwiGLU 8 dim^2, adaLN and norms, rounded up. BLOCK_PARAMS = 13 * Z_IMAGE_HIDDEN_SIZE**2 @@ -66,6 +77,13 @@ TEXT_TOKENS = 512 DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32} +CONTROL_VARIANTS: dict[str, ZImageFunControlVariant | None] = { + "none": None, + "1.0": Z_IMAGE_FUN_CONTROL_1_0, + "2.0": Z_IMAGE_FUN_CONTROL_2_0, + "2.1": Z_IMAGE_FUN_CONTROL_2_1, + "2.1-lite": Z_IMAGE_FUN_CONTROL_2_1_LITE, +} def image_tokens(px: int) -> int: @@ -77,22 +95,39 @@ def unified_tokens(px: int, text_tokens: int) -> int: return _padded(image_tokens(px)) + _padded(text_tokens) -def predicted_bytes(px: int, blocks: int, dtype: torch.dtype, force_math: bool, control_blocks: int) -> int: +def control_blocks(variant: ZImageFunControlVariant | None) -> int: + return len(variant.main_layers) if variant is not None else 0 + + +def base_blocks(blocks: int, variant: ZImageFunControlVariant | None) -> int: + """The main blocks a point builds: one per hint at least, which land on consecutive blocks.""" + return max(blocks, control_blocks(variant)) + + +def packed(variant: ZImageFunControlVariant) -> ZImageFunControlVariant: + """``variant`` with its hints on main blocks 0, 1, 2, ...: the same peak on the fewest blocks.""" + return dataclasses.replace(variant, main_layers=tuple(range(control_blocks(variant)))) + + +def predicted_bytes( + px: int, blocks: int, dtype: torch.dtype, force_math: bool, variant: ZImageFunControlVariant | None +) -> int: """What a point is expected to take: working memory as the node estimates it, plus the weights.""" working = ZImageDenoiseInvocation._estimate_working_memory( image_seq_len=image_tokens(px), text_seq_len=TEXT_TOKENS, num_loras=0, - num_control_blocks=control_blocks, - has_attention_mask=control_blocks > 0, + num_control_blocks=control_blocks(variant), device=torch.device("cuda"), dtype=dtype, ) if force_math: tokens = unified_tokens(px, TEXT_TOKENS) working += Z_IMAGE_ATTENTION_HEADS * tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - # Main blocks plus refiners and embedders; adapter blocks plus its refiners. - modules = blocks + 4 + (control_blocks + 2 if control_blocks else 0) + # Main blocks plus refiners and embedders; the control's blocks plus its refiners. + modules = base_blocks(blocks, variant) + 2 * REFINER_BLOCKS + if variant is not None: + modules += control_blocks(variant) + (REFINER_BLOCKS if variant.refiner_chain != "main" else 0) return working + modules * BLOCK_PARAMS * dtype.itemsize @@ -112,35 +147,29 @@ def _peak_reserved(fn) -> int | None: return torch.cuda.max_memory_reserved() - baseline -def _control_adapter(blocks: int, device: torch.device, dtype: torch.dtype) -> torch.nn.Module: - """A randomly initialized adapter, allocated directly in ``dtype`` on the device: 15 blocks are ~3B parameters, +def _control(variant: ZImageFunControlVariant, device: torch.device, dtype: torch.dtype) -> torch.nn.Module: + """A randomly initialized control, allocated directly in ``dtype`` on the device: 17 blocks are ~3B parameters, which would otherwise pass through the host in float32 first.""" - from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter + from invokeai.backend.fun_controlnet.z_image import ZImageFunControl with torch.device("meta"): - adapter = ZImageControlAdapter( - num_control_blocks=blocks, - control_in_dim=CONTROL_IN_DIM, - dim=Z_IMAGE_HIDDEN_SIZE, - n_heads=Z_IMAGE_ATTENTION_HEADS, - n_kv_heads=Z_IMAGE_ATTENTION_HEADS, + control = ZImageFunControl( + variant, hidden_size=Z_IMAGE_HIDDEN_SIZE, num_attention_heads=Z_IMAGE_ATTENTION_HEADS ) - adapter = adapter.to(dtype).to_empty(device=device) - for param in adapter.parameters(): + control = control.to(dtype).to_empty(device=device) + for param in control.parameters(): torch.nn.init.normal_(param, std=0.02) - return adapter.eval() + return control.eval() @torch.inference_mode() -def measure(px: int, text_tokens: int, blocks: int, dtype: torch.dtype, force_math: bool, control_blocks: int) -> dict: +def measure(px: int, text_tokens: int, blocks: int, dtype: torch.dtype, force_math: bool, variant_name: str) -> dict: """Peak reserved memory for one transformer forward at one resolution.""" from diffusers import ZImageTransformer2DModel - from invokeai.backend.z_image.z_image_controlnet_extension import ( - ZImageControlNetExtension, - z_image_forward_with_control, - ) + from invokeai.backend.fun_controlnet.hooks import install_fun_control + variant = CONTROL_VARIANTS[variant_name] device = torch.device("cuda") model = ( ZImageTransformer2DModel( @@ -148,8 +177,8 @@ def measure(px: int, text_tokens: int, blocks: int, dtype: torch.dtype, force_ma all_f_patch_size=(1,), in_channels=16, dim=Z_IMAGE_HIDDEN_SIZE, - n_layers=blocks, - n_refiner_layers=1, + n_layers=base_blocks(blocks, variant), + n_refiner_layers=REFINER_BLOCKS, n_heads=Z_IMAGE_ATTENTION_HEADS, n_kv_heads=Z_IMAGE_ATTENTION_HEADS, norm_eps=1e-05, @@ -168,28 +197,24 @@ def measure(px: int, text_tokens: int, blocks: int, dtype: torch.dtype, force_ma t = torch.full((1,), 0.5, device=device, dtype=dtype) cap_feats = [torch.randn(text_tokens, CAP_FEAT_DIM, device=device, dtype=dtype)] - extension = None - if control_blocks: - extension = ZImageControlNetExtension( - control_adapter=_control_adapter(control_blocks, device, dtype), - control_cond=torch.randn(CONTROL_IN_DIM, 1, latent, latent, device=device, dtype=dtype), - ) + control = None + if variant is not None: + control = _control(packed(variant), device, dtype) + context = torch.randn(1, image_tokens(px), control.control_in_dim, device=device, dtype=dtype) def forward() -> None: - with sdpa_kernel(SDPBackend.MATH) if force_math else nullcontext(): - if extension is None: - model(x=x, t=t, cap_feats=cap_feats) - else: - z_image_forward_with_control( - transformer=model, x=x, t=t, cap_feats=cap_feats, control_extension=extension - ) + with ( + sdpa_kernel(SDPBackend.MATH) if force_math else nullcontext(), + install_fun_control(model, control, context, scale=0.75) if control is not None else nullcontext(), + ): + model(x=x, t=t, cap_feats=cap_feats) row = { "px": px, "tokens": unified_tokens(px, text_tokens), "blocks": blocks, "force_math": force_math, - "control_blocks": control_blocks, + "control_variant": variant_name, } # Warm up first: the cold call pays one-off workspace costs the per-token slope must not absorb. if _peak_reserved(forward) is None: @@ -198,8 +223,8 @@ def forward() -> None: return row | ({"oom": True} if peak is None else {"oom": False, "reserved_delta": peak}) -def _run_point(px: int, blocks: int, dtype_name: str, force_math: bool, control_blocks: int) -> dict | None: - args = [str(px), str(TEXT_TOKENS), str(blocks), dtype_name, "1" if force_math else "0", str(control_blocks)] +def _run_point(px: int, blocks: int, dtype_name: str, force_math: bool, variant_name: str) -> dict | None: + args = [str(px), str(TEXT_TOKENS), str(blocks), dtype_name, "1" if force_math else "0", variant_name] proc = subprocess.run([sys.executable, __file__, "--single", *args], capture_output=True, text=True) line = proc.stdout.strip().splitlines()[-1] if proc.stdout.strip() else "" try: @@ -210,17 +235,15 @@ def _run_point(px: int, blocks: int, dtype_name: str, force_math: bool, control_ return None -def report( - dtype_name: str, blocks: int, extra_blocks: int | None, control_blocks: int, max_px: int, force: bool -) -> None: +def report(dtype_name: str, blocks: int, extra_blocks: int | None, variant_name: str, max_px: int, force: bool) -> None: dtype = DTYPES[dtype_name] sizes = [px for px in SIZES if px <= max_px] free, _ = torch.cuda.mem_get_info() print(f"{'kernel':>7} {'blocks':>6} {'control':>7} {'px':>5} {'tokens':>7} {'reserved(GiB)':>14}") print("-" * 52) - def run(kernel: str, count: int, px: int, force_math: bool, control: int) -> int | None: - need = predicted_bytes(px, count, dtype, force_math, control) + def run(kernel: str, count: int, px: int, force_math: bool, control: str) -> int | None: + need = predicted_bytes(px, count, dtype, force_math, CONTROL_VARIANTS[control]) if need > free and not force: print( f"{kernel:>7} {count:>6} {control:>7} {px:>5} {'':>7} {'skipped':>14}" @@ -238,10 +261,11 @@ def run(kernel: str, count: int, px: int, force_math: bool, control: int) -> int fused: dict[int, dict[int, int]] = {} for count in counts: for px in sizes: - if (peak := run("auto", count, px, False, 0)) is not None: + if (peak := run("auto", count, px, False, "none")) is not None: fused.setdefault(count, {})[px] = peak - math_peak = run("math", blocks, SIZES[0], True, 0) - control_peak = run("auto", blocks, SIZES[0], False, control_blocks) if control_blocks else None + math_peak = run("math", blocks, SIZES[0], True, "none") + variant = CONTROL_VARIANTS[variant_name] + control_peak = run("auto", blocks, SIZES[0], False, variant_name) if variant is not None else None print() for count, peaks in fused.items(): @@ -267,9 +291,10 @@ def run(kernel: str, count: int, px: int, force_math: bool, control: int) -> int if control_peak is not None and reference is not None: extra = control_peak - reference hint_bytes = tokens * Z_IMAGE_HIDDEN_SIZE * dtype.itemsize + count = control_blocks(variant) print( - f"control, {control_blocks} blocks, {SIZES[0]}px: {extra / GIB:.3f} GiB over no control, " - f"{extra / hint_bytes:.2f} hint-sized tensors; the estimate charges n + 4 = {control_blocks + 4}" + f"control {variant_name}, {count} blocks, {SIZES[0]}px: {extra / GIB:.3f} GiB over no control, " + f"{extra / hint_bytes:.2f} hint-sized tensors; the estimate charges n + 4 = {count + 4}" ) @@ -278,7 +303,9 @@ def main() -> None: parser.add_argument("--dtype", choices=list(DTYPES), default="bfloat16") parser.add_argument("--blocks", type=int, default=2, help="Main transformer blocks. Default 2.") parser.add_argument("--extra-blocks", type=int, default=None, help="Measure a second block count as well.") - parser.add_argument("--control-blocks", type=int, default=15, help="ControlNet blocks; 0 skips. Default 15.") + parser.add_argument( + "--control-variant", choices=list(CONTROL_VARIANTS), default="2.1", help="ControlNet layout. Default 2.1." + ) parser.add_argument("--max-px", type=int, default=SIZES[-1], help="Largest fused-kernel size to measure.") parser.add_argument("--force", action="store_true", help="Run points predicted not to fit free device memory.") parser.add_argument("--single", nargs="*", default=None, help=argparse.SUPPRESS) @@ -287,16 +314,14 @@ def main() -> None: if not torch.cuda.is_available(): raise SystemExit("No CUDA/HIP device available.") if args.single: - px, text_tokens, blocks, dtype_name, force_math, control_blocks = args.single - row = measure( - int(px), int(text_tokens), int(blocks), DTYPES[dtype_name], force_math == "1", int(control_blocks) - ) + px, text_tokens, blocks, dtype_name, force_math, variant_name = args.single + row = measure(int(px), int(text_tokens), int(blocks), DTYPES[dtype_name], force_math == "1", variant_name) print(json.dumps(row)) return print( f"torch {torch.__version__} | device {torch.cuda.get_device_name(0)} | hip={torch.version.hip} | {args.dtype}" ) - report(args.dtype, args.blocks, args.extra_blocks, args.control_blocks, args.max_px, args.force) + report(args.dtype, args.blocks, args.extra_blocks, args.control_variant, args.max_px, args.force) if __name__ == "__main__": diff --git a/tests/app/invocations/test_z_image_control.py b/tests/app/invocations/test_z_image_control.py new file mode 100644 index 00000000000..f2ac713c990 --- /dev/null +++ b/tests/app/invocations/test_z_image_control.py @@ -0,0 +1,267 @@ +"""What the Z-Image denoise does with its control field: refuses a run without the VAE or with another kind of model, +locks the control with the transformer's reservation and hooks it inside the exit stack, gates it once per step for +both guidance passes and both of Heun's forwards, and builds the packed control context VideoX-Fun builds for control +alone. + +The denoise runs its transformer in bfloat16, which the CPU test fleet cannot matmul, so every run here stops at the +LoRA patcher or runs the loop with the transformer's forward stood in. The forward with the hooks is covered in +`tests/backend/fun_controlnet/test_z_image_fun_control.py`; the reservation and the encode's tiling in +`test_z_image_denoise_working_memory.py`. +""" + +from contextlib import contextmanager, nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel +from PIL import Image + +from invokeai.app.invocations.fields import ZImageConditioningField +from invokeai.app.invocations.model import ModelIdentifierField, VAEField +from invokeai.app.invocations.z_image import z_image_denoise +from invokeai.app.invocations.z_image.z_image_control import ZImageControlField +from invokeai.app.invocations.z_image.z_image_denoise import ZImageDenoiseInvocation +from invokeai.backend.fun_controlnet.hooks import FunControlHooks +from invokeai.backend.fun_controlnet.z_image import Z_IMAGE_FUN_CONTROL_2_1_LITE, ZImageFunControl +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType +from invokeai.backend.patches.layer_patcher import LayerPatcher +from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor +from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( + ConditioningFieldData, + ZImageConditioningInfo, +) +from invokeai.backend.util.devices import TorchDevice + +SIDE = 64 # an 8x8 latent: 16 image tokens +LATENT_SIDE = SIDE // 8 +CAPTION_TOKENS = 4 +CAPTION_DIM = 32 +HIDDEN = 64 +HEADS = 2 +# The Lite layout's hints land after base layers 0, 10 and 20, which a 21-block transformer is just deep enough for. +LAYERS = 21 + + +class _Stop(Exception): + """Raised where the node has made the decision a test looks at.""" + + +@pytest.fixture(autouse=True) +def cpu(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(TorchDevice, "choose_bfloat16_safe_dtype", lambda _device=None: torch.bfloat16) + + +def _model(type_: ModelType) -> ModelIdentifierField: + return ModelIdentifierField(key=type_.value, hash="hash", name=type_.value, base=BaseModelType.ZImage, type=type_) + + +_TRANSFORMER_CONFIG = { + "all_patch_size": (2,), + "all_f_patch_size": (1,), + "in_channels": 16, + "dim": HIDDEN, + "n_layers": LAYERS, + "n_refiner_layers": 2, + "n_heads": HEADS, + "n_kv_heads": HEADS, + "cap_feat_dim": CAPTION_DIM, + "axes_dims": [8, 12, 12], + "axes_lens": [128, 32, 32], +} + + +def _tiny_transformer() -> ZImageTransformer2DModel: + return ZImageTransformer2DModel(**_TRANSFORMER_CONFIG).eval() + + +def _tiny_control() -> ZImageFunControl: + return ZImageFunControl(Z_IMAGE_FUN_CONTROL_2_1_LITE, hidden_size=HIDDEN, num_attention_heads=HEADS).eval() + + +def _loaded(model, lock) -> MagicMock: + """A `LoadedModel` stand-in: `model` unlocked, `model_on_device` as `lock` says (a context or a raise).""" + info = MagicMock() + info.model = model + info.compute_device = torch.device("cpu") + info.model_on_device = MagicMock(side_effect=lock) + return info + + +def _yielding(model): + # The transformer and the control are locked with `working_mem_bytes`, the VAE with `working_memory`. + return lambda **_kwargs: nullcontext((None, model)) + + +def _control_field(**overrides) -> ZImageControlField: + return ZImageControlField(image_name="canny", control_model=_model(ModelType.ControlNet), **overrides) + + +def _denoise(**overrides) -> ZImageDenoiseInvocation: + fields = { + "latents": None, + "noise": None, + "denoise_mask": None, + "denoising_start": 0.0, + "denoising_end": 1.0, + "add_noise": True, + "transformer": SimpleNamespace(transformer=_model(ModelType.Main), loras=[]), + "positive_conditioning": ZImageConditioningField(conditioning_name="pos"), + "negative_conditioning": ZImageConditioningField(conditioning_name="neg"), + "guidance_scale": 3.0, + "width": SIDE, + "height": SIDE, + "steps": 2, + "seed": 0, + "shift": None, + "scheduler": "euler", + "control": _control_field(), + "vae": VAEField(vae=_model(ModelType.VAE)), + } + fields.update(overrides) + return ZImageDenoiseInvocation.model_construct(**fields) + + +def _control_image() -> Image.Image: + return Image.linear_gradient("L").resize((SIDE // 2, SIDE // 2)).convert("RGB") + + +def _context(transformer_info=None, control_info=None, vae_info=None) -> MagicMock: + infos = {ModelType.Main: transformer_info, ModelType.ControlNet: control_info, ModelType.VAE: vae_info} + context = MagicMock() + context.models.load.side_effect = lambda identifier: infos[identifier.type] + context.models.get_config.return_value = SimpleNamespace( + base=BaseModelType.ZImage, type=ModelType.Main, format=ModelFormat.Checkpoint + ) + context.conditioning.load.side_effect = lambda name: ConditioningFieldData( + conditionings=[ZImageConditioningInfo(prompt_embeds=torch.randn(CAPTION_TOKENS, CAPTION_DIM))] + ) + context.images.get_pil.side_effect = lambda name: {"canny": _control_image()}[name] + context.config.get.return_value.force_tiled_decode = False + return context + + +def _hook_count(transformer) -> int: + blocks = [*transformer.noise_refiner, *transformer.layers] + return sum(len(block._forward_hooks) + len(block._forward_pre_hooks) for block in blocks) + + +def test_control_without_its_vae_is_refused_before_anything_else_is_loaded() -> None: + context = _context(transformer_info=_loaded(_tiny_transformer(), _Stop)) + + with pytest.raises(ValueError, match="needs the VAE"): + _denoise(vae=None)._run_diffusion(context) + + assert context.models.load.call_count == 1 + + +def test_another_kind_of_control_model_is_refused() -> None: + context = _context( + transformer_info=_loaded(_tiny_transformer(), _Stop), control_info=_loaded(torch.nn.Linear(2, 2), _Stop) + ) + + with pytest.raises(TypeError, match="Expected a Z-Image ControlNet"): + _denoise()._run_diffusion(context) + + +def test_the_control_is_locked_with_the_same_reservation_and_hooked_inside_the_exit_stack( + monkeypatch: pytest.MonkeyPatch, flux_shaped_vae +) -> None: + transformer, control, vae = _tiny_transformer(), _tiny_control(), flux_shaped_vae() + transformer_info = _loaded(transformer, _yielding(transformer)) + control_info = _loaded(control, _yielding(control)) + context = _context(transformer_info, control_info, _loaded(vae, _yielding(vae))) + hooks_at_patching: list[int] = [] + + def stop(**_kwargs): + hooks_at_patching.append(_hook_count(transformer)) + raise _Stop + + monkeypatch.setattr(LayerPatcher, "apply_smart_model_patches", stop) + + with pytest.raises(_Stop): + _denoise()._run_diffusion(context) + + reserved = transformer_info.model_on_device.call_args.kwargs["working_mem_bytes"] + assert control_info.model_on_device.call_args.kwargs["working_mem_bytes"] == reserved + # Two entry pre-hooks, the refiner pair and the three Lite layers. + assert hooks_at_patching == [2 + 2 + 3] + assert _hook_count(transformer) == 0 + + +class _GateRecordingTransformer(ZImageTransformer2DModel): + """The real class, so the real hooks land on its blocks, with the forward replaced: no bfloat16 matmul on the CPU, + and a record of the control gate at every call.""" + + def __init__(self) -> None: + super().__init__(**_TRANSFORMER_CONFIG) + self.hooks: FunControlHooks | None = None + self.gates: list[bool] = [] + + def forward(self, x, t, cap_feats, *_args, **_kwargs): + assert self.hooks is not None + self.gates.append(self.hooks.active) + return ([torch.zeros_like(item) for item in x],) + + +@pytest.mark.parametrize( + ("scheduler", "begin", "end", "expected_gates"), + [ + # Four Euler steps with CFG are eight forwards, two per step. + ("euler", 0.5, 1.0, [False] * 4 + [True] * 4), + ("euler", 0.0, 0.5, [True] * 6 + [False] * 2), + # Heun runs two forwards per step but the last; with CFG, four calls per step: user steps 0 and 1 off, 2 and + # 3 on. + ("heun", 0.5, 1.0, [False] * 8 + [True] * 6), + # User steps 0 and 1 on: both forwards of step 1 run with the control, though the second is the fourth call. + ("heun", 0.0, 0.25, [True] * 8 + [False] * 6), + ], + ids=["euler-second-half", "euler-first-half", "heun-second-half", "heun-first-quarter"], +) +def test_the_gate_is_set_once_per_step_for_both_guidance_passes( + monkeypatch: pytest.MonkeyPatch, flux_shaped_vae, scheduler, begin, end, expected_gates +) -> None: + """The gate follows `user step / total steps` with both ends included and is re-evaluated every step; the positive + and the negative pass of a step, and Heun's two forwards of a step, see the same one.""" + transformer, control, vae = _GateRecordingTransformer(), _tiny_control(), flux_shaped_vae() + install = z_image_denoise.install_fun_control + + @contextmanager + def capturing_install(*args, **kwargs): + with install(*args, **kwargs) as hooks: + transformer.hooks = hooks + yield hooks + + monkeypatch.setattr(z_image_denoise, "install_fun_control", capturing_install) + context = _context( + _loaded(transformer, _yielding(transformer)), + _loaded(control, _yielding(control)), + _loaded(vae, _yielding(vae)), + ) + + _denoise( + steps=4, scheduler=scheduler, control=_control_field(begin_step_percent=begin, end_step_percent=end) + )._run_diffusion(context) + + assert transformer.gates == expected_gates + + +@torch.no_grad() +def test_the_control_context_is_the_mode_encoded_control_image_packed_like_the_latents(flux_shaped_vae) -> None: + """VideoX-Fun's context for control alone, `[control | keep mask | masked image]` with the last two zero, from the + mode of the VAE's distribution, normalized as the generation latents are, and patched 2x2 the way the transformer + patches its latents.""" + vae, control, transformer = flux_shaped_vae(), _tiny_control(), _tiny_transformer() + context = _context(vae_info=_loaded(vae, _yielding(vae))) + + packed = _denoise()._build_control_context(context, control) + + pixels = image_resized_to_grid_as_tensor(_control_image().resize((SIDE, SIDE), Image.Resampling.LANCZOS)) + mode = vae.encode(pixels.unsqueeze(0)).latent_dist.mode() + latents = (mode - vae.config.shift_factor) * vae.config.scaling_factor + context_4d = torch.cat([latents, torch.zeros(1, 17, LATENT_SIDE, LATENT_SIDE)], dim=1) + expected, _, _ = transformer._patchify_image(context_4d[0].unsqueeze(1), 2, 1) + assert packed.shape == (1, (LATENT_SIDE // 2) ** 2, 132) + assert torch.equal(packed[0], expected) diff --git a/tests/app/invocations/test_z_image_denoise_working_memory.py b/tests/app/invocations/test_z_image_denoise_working_memory.py index 0fe5698d0bb..3f0e426f66a 100644 --- a/tests/app/invocations/test_z_image_denoise_working_memory.py +++ b/tests/app/invocations/test_z_image_denoise_working_memory.py @@ -26,13 +26,13 @@ from invokeai.app.invocations.vae.z_image_image_to_latents import ZImageImageToLatentsInvocation from invokeai.app.invocations.z_image.z_image_control import ZImageControlField, ZImageControlInvocation from invokeai.app.invocations.z_image.z_image_denoise import ZImageDenoiseInvocation +from invokeai.backend.fun_controlnet.z_image import ZImageFunControl from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType from invokeai.backend.quantization.int8_convrot import CONVROT_GROUP_SIZE, Int8ConvrotLinear from invokeai.backend.util import attention from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT from invokeai.backend.z_image.extensions.regional_prompting_extension import ZImageRegionalPromptingExtension from invokeai.backend.z_image.text_conditioning import ZImageRegionalTextConditioning -from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter MB = 1024**2 GB = 1024**3 @@ -77,9 +77,11 @@ class TestEstimate: (1024, 0, 1012), (1440, 0, 2004), (2048, 0, 3942), - (1024, 1, 1122), - (1024, 6, 1212), - (1024, 15, 1650), + # The ControlNet layouts on the hooked forward, rounded up from the GiB the script prints. + (1024, 6, 1073), # 1.0 + (1024, 3, 1043), # 2.1 Lite + (1024, 15, 1405), # 2.0 + (1024, 15, 1403), # 2.1 ] @pytest.mark.parametrize("px, control_blocks, measured_mb", MEASURED_FUSED) @@ -169,22 +171,27 @@ def _run( control = None control_info = MagicMock() + order: list[str] = [] if control_blocks is None: transformer_info.model_on_device = MagicMock(side_effect=_StopBeforeLoad) else: - # Past the transformer's lock, up to the adapter's. - transformer_info.model_on_device = MagicMock(return_value=nullcontext((None, transformer))) + + def lock_transformer(**_kwargs): + order.append("transformer lock") + return nullcontext((None, transformer)) + + # Past the transformer's lock, up to the control's. + transformer_info.model_on_device = MagicMock(side_effect=lock_transformer) control = ZImageControlField.model_construct( control_model=MagicMock(), image_name="control", **(control_tiling or {}) ) - control_info.model = MagicMock(spec=ZImageControlAdapter) + control_info.model = MagicMock(spec=ZImageFunControl) control_info.model.control_layers = [torch.nn.Identity()] * control_blocks - if control_tiling is None: - control_info.model_on_device = MagicMock(side_effect=_StopBeforeLoad) - else: - # Past the adapter's lock too, up to the control image's encode. - control_info.model.config = {"control_in_dim": 33} - control_info.model_on_device = MagicMock(return_value=nullcontext((None, control_info.model))) + control_info.model_on_device = MagicMock(side_effect=_StopBeforeLoad) + + def encode(**_kwargs): + order.append("encode") + return torch.zeros(1, 16, 32, 32) caption_tokens = {"pos": POSITIVE_TOKENS, "neg": negative_tokens} context = MagicMock() @@ -240,7 +247,7 @@ def _run( patch.object(ZImageDenoiseInvocation, "_estimate_working_memory", return_value=ACTIVATION_ESTIMATE) ) control_encode = stack.enter_context( - patch.object(ZImageImageToLatentsInvocation, "vae_encode", side_effect=_StopBeforeLoad) + patch.object(ZImageImageToLatentsInvocation, "vae_encode", side_effect=encode) ) if regional_mask is not None: regional = ZImageRegionalPromptingExtension( @@ -258,10 +265,11 @@ def _run( return SimpleNamespace( reserved=transformer_info.model_on_device.call_args.kwargs["working_mem_bytes"], estimate=estimate.call_args.kwargs, - adapter_reserved=( + control_reserved=( control_info.model_on_device.call_args.kwargs.get("working_mem_bytes") if control is not None else None ), control_encode=control_encode.call_args.kwargs if control_encode.called else None, + order=order, ) def test_an_unquantized_transformer_reserves_the_estimate_for_its_sequence(self): @@ -302,14 +310,22 @@ def test_a_regional_mask_reaches_the_estimate(self): True, ) - def test_a_control_adapter_reaches_the_estimate_and_keeps_the_reservation(self): - """The adapter's block count sizes the hints and its forward always passes a padding mask. It is locked after - the transformer, which can then no longer be offloaded -- so it must ask for the same working memory, or its - weights fill the room kept for the forward.""" + def test_a_control_reaches_the_estimate_and_keeps_the_reservation(self): + """The control's block count sizes the hints; its blocks take the base blocks' arguments, so it adds no + attention mask. It is locked after the transformer, which can then no longer be offloaded -- so it must ask for + the same working memory, or its weights fill the room kept for the forward.""" + run = self._run(control_blocks=15) + + assert (run.estimate["num_control_blocks"], run.estimate["has_attention_mask"]) == (15, False) + assert run.control_reserved == run.reserved == ACTIVATION_ESTIMATE + + def test_the_control_image_is_encoded_before_the_transformer_is_locked(self): + """So the VAE does not have to fit beside a transformer the cache can no longer offload. The encode takes the + mode of the VAE's distribution, as VideoX-Fun encodes control images: a sample would add noise of its own.""" run = self._run(control_blocks=15) - assert (run.estimate["num_control_blocks"], run.estimate["has_attention_mask"]) == (15, True) - assert run.adapter_reserved == run.reserved == ACTIVATION_ESTIMATE + assert run.order == ["encode", "transformer lock"] + assert run.control_encode["deterministic"] is True def test_the_control_fields_tiling_reaches_the_control_encode(self): """The control image is encoded at generation size. For an upscale that is the enlarged frame, whose untiled diff --git a/tests/backend/fun_controlnet/test_qwen_image_fun_control.py b/tests/backend/fun_controlnet/test_qwen_image_fun_control.py index 78de91abeb4..94b754665d9 100644 --- a/tests/backend/fun_controlnet/test_qwen_image_fun_control.py +++ b/tests/backend/fun_controlnet/test_qwen_image_fun_control.py @@ -5,6 +5,8 @@ the hooks beyond the modules' own forwards. Everything runs in float32 on the CPU. """ +import weakref + import accelerate import pytest import torch @@ -146,6 +148,30 @@ def test_scale_zero_and_an_inactive_handle_leave_the_forward_untouched( assert torch.equal(_forward(model, inputs), plain) +@torch.no_grad() +def test_the_projected_context_is_released_before_the_last_control_block( + tiny_qwen_image_transformer, tiny_branch, inputs, control_context +) -> None: + """The chain owns its input: the projected context is freed once block 0 has joined it with the image stream, not + held through every block -- an image-length hidden state more than the hint reservation allows for.""" + projected: list[weakref.ref] = [] + alive_at_last_block: list[bool] = [] + taps = [ + tiny_branch.control_img_in.register_forward_hook(lambda _m, _a, out: projected.append(weakref.ref(out))), + tiny_branch.control_blocks[-1].register_forward_pre_hook( + lambda *_args: alive_at_last_block.append(projected[-1]() is not None) + ), + ] + try: + with install_fun_control(tiny_qwen_image_transformer, tiny_branch, control_context, scale=SCALE): + _forward(tiny_qwen_image_transformer, inputs) + finally: + for tap in taps: + tap.remove() + + assert alive_at_last_block == [False] + + @torch.no_grad() def test_the_hooks_come_off_on_exit_and_when_the_body_raises( tiny_qwen_image_transformer, tiny_branch, inputs, control_context diff --git a/tests/backend/fun_controlnet/test_z_image_fun_control.py b/tests/backend/fun_controlnet/test_z_image_fun_control.py new file mode 100644 index 00000000000..edb67188344 --- /dev/null +++ b/tests/backend/fun_controlnet/test_z_image_fun_control.py @@ -0,0 +1,697 @@ +"""The Z-Image Fun control, hooked into the real transformer, computes what VideoX-Fun's control transformer does. + +`_videox_forward` restates VideoX-Fun's `ZImageControlTransformer2DModel.forward` at batch size 1 against diffusers' +`ZImageTransformer2DModel`: the base preamble, the control context patched, embedded and padded with the base's +`x_pad_token` (`patchify`), the refiner chain with its hints after the base refiner blocks (`forward_control_2_0_refiner`) +or the plain control refiner (`forward_control_1_0`), the main chain over `[control | caption]` +(`forward_control_2_0_layers`), and an explicit loop over the base blocks adding each hint after its listed layer. The +control blocks run in VideoX-Fun's own stacked formulation. Its masks are all-True at batch size 1; the restatement +passes `None`, as diffusers does. Where the hints land, which chain refines the control and how many context channels +it reads come from `VIDEOX_CONFIGS`, VideoX-Fun's four config files restated, never from the implementation's variants; +the restatement shares only the modules' own forwards with the hooks. Everything runs in float32 on the CPU, every +parameter drawn at random, the zero-initialised projections and the pad tokens included. +""" + +import weakref + +import accelerate +import pytest +import torch +from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel, ZImageTransformerBlock + +from invokeai.backend.fun_controlnet.hooks import install_fun_control +from invokeai.backend.fun_controlnet.z_image import ( + Z_IMAGE_FUN_CONTROL_1_0, + Z_IMAGE_FUN_CONTROL_2_0, + Z_IMAGE_FUN_CONTROL_2_1, + Z_IMAGE_FUN_CONTROL_2_1_LITE, + ZImageFunControl, + ZImageFunControlVariant, +) +from invokeai.backend.z_image.z_image_transformer_patch import patch_transformer_for_regional_prompting +from tests.backend.model_manager.load.state_dicts.z_image_fun_controlnet_1_0_keys import ( + state_dict_keys as keys_1_0, +) +from tests.backend.model_manager.load.state_dicts.z_image_fun_controlnet_2_1_keys import ( + state_dict_keys as keys_2_1, +) +from tests.backend.model_manager.load.state_dicts.z_image_fun_controlnet_2_1_lite_keys import ( + state_dict_keys as keys_2_1_lite, +) + +HIDDEN = 64 # two heads of 32 +HEADS = 2 +HEAD_DIM = HIDDEN // HEADS +CAPTION_DIM = 32 +LATENT_SIDE = 10 # 5x5 = 25 image tokens: the transformer pads them to 32 with seven `x_pad_token` rows +IMAGE_TOKENS = (LATENT_SIDE // 2) ** 2 +CAPTION_TOKENS = 20 # padded to 32 with `cap_pad_token` +SCALE = 0.75 + +VARIANTS = [Z_IMAGE_FUN_CONTROL_1_0, Z_IMAGE_FUN_CONTROL_2_0, Z_IMAGE_FUN_CONTROL_2_1, Z_IMAGE_FUN_CONTROL_2_1_LITE] + +# VideoX-Fun `config/z_image/z_image_control.yaml`, `_2.0.yaml`, `_2.1.yaml` and `_2.1_lite.yaml`, as written there. +# Keys a file leaves out take `ZImageControlTransformer2DModel`'s defaults: no control refiner, and refiner hints after +# base refiner blocks `range(n_refiner_layers)`. +VIDEOX_CONFIGS: dict[str, dict] = { + "1.0": {"control_layers_places": [0, 5, 10, 15, 20, 25], "control_in_dim": 16}, + "2.0": { + "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, + "add_control_noise_refiner_correctly": False, + "control_in_dim": 33, + }, + "2.1": { + "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, + "add_control_noise_refiner_correctly": True, + "control_in_dim": 33, + }, + "2.1 Lite": { + "control_layers_places": [0, 10, 20], + "control_refiner_layers_places": [0, 1], + "add_control_noise_refiner": True, + "add_control_noise_refiner_correctly": True, + "control_in_dim": 33, + }, +} +REFINER_BLOCKS = 2 + + +def _config(variant: ZImageFunControlVariant) -> dict: + config = { + "control_refiner_layers_places": list(range(REFINER_BLOCKS)), + "add_control_noise_refiner": False, + "add_control_noise_refiner_correctly": False, + } + return config | VIDEOX_CONFIGS[variant.name] + + +def _randomized(module: torch.nn.Module, seed: int) -> torch.nn.Module: + """Every parameter from N(0, 0.1). diffusers zero-initialises the pad tokens and VideoX-Fun the projections, which + would leave a pad row or a hint that is zero and a parity test that cannot fail.""" + torch.manual_seed(seed) + for parameter in module.parameters(): + torch.nn.init.normal_(parameter, std=0.1) + return module.eval() + + +@pytest.fixture(scope="module") +def transformer() -> ZImageTransformer2DModel: + """The real Z-Image transformer at toy width: 30 main blocks and two refiner blocks, so every variant's real layer + map fits.""" + model = ZImageTransformer2DModel( + all_patch_size=(2,), + all_f_patch_size=(1,), + in_channels=16, + dim=HIDDEN, + n_layers=30, + n_refiner_layers=2, + n_heads=HEADS, + n_kv_heads=HEADS, + norm_eps=1e-05, + qk_norm=True, + cap_feat_dim=CAPTION_DIM, + rope_theta=256.0, + t_scale=1000.0, + axes_dims=[8, 12, 12], + axes_lens=[128, 32, 32], + ) + return _randomized(model, seed=0) + + +def _control(variant: ZImageFunControlVariant, seed: int = 1) -> ZImageFunControl: + return _randomized(ZImageFunControl(variant, hidden_size=HIDDEN, num_attention_heads=HEADS), seed) + + +@pytest.fixture +def inputs(transformer) -> dict: + torch.manual_seed(2) + return { + "x": [torch.randn(16, 1, LATENT_SIDE, LATENT_SIDE)], + "t": torch.tensor([0.4]), + "cap_feats": [torch.randn(CAPTION_TOKENS, CAPTION_DIM)], + } + + +@pytest.fixture +def context_4d() -> torch.Tensor: + """A control context `[control | keep mask | masked image]`, every channel non-zero so a dropped one shows.""" + torch.manual_seed(3) + return torch.randn(1, 33, LATENT_SIDE, LATENT_SIDE) + + +def _forward(model, inputs) -> torch.Tensor: + # Indexed rather than `return_dict=False`, which regional prompting's patched forward does not take. + return model(**inputs)[0][0] + + +def _hook_count(model) -> int: + blocks = [*model.noise_refiner, *model.layers] + return sum(len(block._forward_hooks) + len(block._forward_pre_hooks) for block in blocks) + + +def _videox_control_blocks(blocks, c, x, **kwargs) -> tuple[list[torch.Tensor], torch.Tensor]: + """`ZImageControlTransformerBlock.forward` as VideoX-Fun writes it: every block stacks its hint and its stream + onto what it was given, and the next one unbinds them again.""" + for index, block in enumerate(blocks): + if index == 0: + c = block.before_proj(c) + x + all_c = [] + else: + all_c = list(torch.unbind(c)) + c = all_c.pop(-1) + c = ZImageTransformerBlock.forward(block, c, **kwargs) + all_c += [block.after_proj(c), c] + c = torch.stack(all_c) + *hints, c = torch.unbind(c) + return hints, c + + +def _videox_patchify(context: torch.Tensor, pad_to: int) -> tuple[torch.Tensor, torch.Tensor]: + """VideoX-Fun's `patchify` of a `(C, H, W)` context: 2x2 patches in `(pH, pW, C)` order, padded with copies of the + last token, and the mask of those padding rows.""" + channels, height, width = context.shape + tokens = context.view(channels, height // 2, 2, width // 2, 2).permute(1, 3, 2, 4, 0).reshape(-1, 4 * channels) + padding = pad_to - len(tokens) + padded = torch.cat([tokens, tokens[-1:].repeat(padding, 1)]) + return padded, torch.cat([torch.zeros(len(tokens), dtype=torch.bool), torch.ones(padding, dtype=torch.bool)]) + + +def _videox_forward(model, control, config, context_4d, scale, *, x, t, cap_feats, drop_refiner_hints=False): + places, refiner_places = config["control_layers_places"], config["control_refiner_layers_places"] + adaln = model.t_embedder(t * model.t_scale) + (xs, caps, sizes, x_pos, cap_pos, x_pad, cap_pad) = model.patchify_and_embed(x, cap_feats, 2, 1) + + hidden = model.all_x_embedder["2-1"](xs[0]) + hidden[x_pad[0]] = model.x_pad_token + hidden = hidden[None] + x_freqs = model.rope_embedder(x_pos[0])[None] + + # `patchify` + embed: the control context's padding rows are the base's `x_pad_token`. + tokens, pad_mask = _videox_patchify(context_4d[0, : config["control_in_dim"]], hidden.shape[1]) + c = control.control_all_x_embedder["2-1"](tokens) + c[pad_mask] = model.x_pad_token + c = c[None] + + refiner_hints = None + if config["add_control_noise_refiner"]: + # `forward_control_2_0_refiner`: `add_control_noise_refiner_correctly` picks the refiner blocks (2.1) or the + # main control blocks (2.0) for the chain over the image stream. + correctly = config["add_control_noise_refiner_correctly"] + chain = control.control_noise_refiner if correctly else control.control_layers + refiner_hints, c = _videox_control_blocks( + chain, c, hidden, attn_mask=None, freqs_cis=x_freqs, adaln_input=adaln + ) + for index, block in enumerate(model.noise_refiner): + hidden = block(hidden, None, x_freqs, adaln) + if refiner_hints is not None and index in refiner_places and not drop_refiner_hints: + hidden = hidden + refiner_hints[refiner_places.index(index)] * scale + + caption = model.cap_embedder(caps[0]) + caption[cap_pad[0]] = model.cap_pad_token + caption = caption[None] + # diffusers numbers the caption's padding rows on from its tokens and appends position-0 ids it then cuts off. + cap_freqs = model.rope_embedder(cap_pos[0])[None, : caption.shape[1]] + for block in model.context_refiner: + caption = block(caption, None, cap_freqs) + + if not config["add_control_noise_refiner"]: + # `forward_control_1_0`: plain refiner blocks over the control context alone. + for block in control.control_noise_refiner: + c = block(c, None, x_freqs, adaln) + + unified = torch.cat([hidden, caption], dim=1) + unified_freqs = torch.cat([x_freqs, cap_freqs], dim=1) + hints, _ = _videox_control_blocks( + control.control_layers, + torch.cat([c, caption], dim=1), + unified, + attn_mask=None, + freqs_cis=unified_freqs, + adaln_input=adaln, + ) + for index, block in enumerate(model.layers): + unified = block(unified, None, unified_freqs, adaln) + if index in places: + unified = unified + hints[places.index(index)] * scale + unified = model.all_final_layer["2-1"](unified, adaln) + return model.unpatchify(list(unified.unbind(dim=0)), sizes, 2, 1)[0] + + +@torch.no_grad() +@pytest.mark.parametrize("variant", VARIANTS, ids=lambda v: v.name) +def test_the_hooked_forward_computes_the_videox_formulation(transformer, inputs, context_4d, variant) -> None: + control = _control(variant) + expected = _videox_forward(transformer, control, _config(variant), context_4d, SCALE, **inputs) + plain = _forward(transformer, inputs) + + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE): + hooked = _forward(transformer, inputs) + + assert torch.equal(hooked, expected) + assert not torch.allclose(hooked, plain) + + +@torch.no_grad() +def test_both_guidance_passes_take_the_same_context(transformer, inputs, context_4d) -> None: + """The positive and the negative pass of a step run one after the other with the hooks installed once; each must + match its own restatement, with nothing of the first carried into the second.""" + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + # A longer caption: the unified stream, and with it the control chain's, is longer on this pass. + negative = {**inputs, "cap_feats": [torch.randn(CAPTION_TOKENS + 20, CAPTION_DIM)]} + + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE): + positive_out = _forward(transformer, inputs) + negative_out = _forward(transformer, negative) + + config = _config(Z_IMAGE_FUN_CONTROL_2_1) + assert torch.equal(positive_out, _videox_forward(transformer, control, config, context_4d, SCALE, **inputs)) + assert torch.equal(negative_out, _videox_forward(transformer, control, config, context_4d, SCALE, **negative)) + + +@torch.no_grad() +def test_the_parity_can_tell_where_the_hints_land(transformer, inputs, context_4d) -> None: + """The comparison above is not vacuous: the Lite weights placed after base layers 0, 2 and 4 -- the map Invoke used + for every build before -- compute something else, and so does a forward without the refiner hints.""" + control = _control(Z_IMAGE_FUN_CONTROL_2_1_LITE) + misplaced = ZImageFunControl( + ZImageFunControlVariant("misplaced", 33, (0, 2, 4), "refiner"), hidden_size=HIDDEN, num_attention_heads=HEADS + ).eval() + misplaced.load_state_dict(control.state_dict()) + config = _config(Z_IMAGE_FUN_CONTROL_2_1_LITE) + expected = _videox_forward(transformer, control, config, context_4d, SCALE, **inputs) + + with install_fun_control(transformer, misplaced, misplaced.pack(context_4d), scale=SCALE): + assert not torch.allclose(_forward(transformer, inputs), expected) + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE): + hooked = _forward(transformer, inputs) + assert not torch.allclose( + hooked, _videox_forward(transformer, control, config, context_4d, SCALE, drop_refiner_hints=True, **inputs) + ) + + +@torch.no_grad() +@pytest.mark.parametrize("variant", [Z_IMAGE_FUN_CONTROL_1_0, Z_IMAGE_FUN_CONTROL_2_1], ids=lambda v: v.name) +def test_hints_land_on_the_listed_blocks_only(transformer, inputs, context_4d, variant) -> None: + """Each base block's output is recorded before the hint hook sees it and after: it moves at the listed refiner and + main blocks and nowhere else, and 1.0 leaves the refiner untouched.""" + control = _control(variant) + blocks = {f"refiner.{i}": b for i, b in enumerate(transformer.noise_refiner)} + blocks.update({f"layer.{i}": b for i, b in enumerate(transformer.layers)}) + before: dict[str, torch.Tensor] = {} + after: dict[str, torch.Tensor] = {} + handles = [ + block.register_forward_hook(lambda _m, _a, out, name=name: before.__setitem__(name, out)) + for name, block in blocks.items() + ] + try: + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE): + handles += [ + block.register_forward_hook(lambda _m, _a, out, name=name: after.__setitem__(name, out)) + for name, block in blocks.items() + ] + _forward(transformer, inputs) + finally: + for handle in handles: + handle.remove() + + config = _config(variant) + refiner_places = config["control_refiner_layers_places"] if config["add_control_noise_refiner"] else [] + hinted = {f"refiner.{i}" for i in refiner_places} | {f"layer.{i}" for i in config["control_layers_places"]} + moved = {name for name in blocks if after[name] is not before[name]} + assert moved == hinted + assert all(not torch.equal(after[name], before[name]) for name in hinted) + + +@torch.no_grad() +def test_scale_zero_and_an_inactive_handle_leave_the_forward_untouched(transformer, inputs, context_4d) -> None: + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + context = control.pack(context_4d) + plain = _forward(transformer, inputs) + + with install_fun_control(transformer, control, context, scale=0.0): + assert torch.equal(_forward(transformer, inputs), plain) + + with install_fun_control(transformer, control, context, scale=SCALE) as hooks: + hooks.active = False + assert torch.equal(_forward(transformer, inputs), plain) + hooks.active = True + assert not torch.equal(_forward(transformer, inputs), plain) + hooks.active = False + # Nothing of the active forward is left for the inactive one. + assert torch.equal(_forward(transformer, inputs), plain) + + +@torch.no_grad() +@pytest.mark.parametrize( + ("variant", "expected_hooks"), + # Two entry pre-hooks, then one post-hook per hinted block: the refiner pair (none for 1.0) and the main layers. + [ + (Z_IMAGE_FUN_CONTROL_1_0, 2 + 6), + (Z_IMAGE_FUN_CONTROL_2_1, 2 + 2 + 15), + (Z_IMAGE_FUN_CONTROL_2_1_LITE, 2 + 2 + 3), + ], + ids=lambda v: getattr(v, "name", str(v)), +) +def test_the_hooks_come_off_on_exit_and_when_the_body_raises( + transformer, inputs, context_4d, variant, expected_hooks +) -> None: + control = _control(variant) + context = control.pack(context_4d) + plain = _forward(transformer, inputs) + + with install_fun_control(transformer, control, context, scale=SCALE): + assert _hook_count(transformer) == expected_hooks + assert _hook_count(transformer) == 0 + + with pytest.raises(RuntimeError, match="mid-step"): + with install_fun_control(transformer, control, context, scale=SCALE): + _forward(transformer, inputs) + raise RuntimeError("mid-step") + assert _hook_count(transformer) == 0 + assert torch.equal(_forward(transformer, inputs), plain) + + +@torch.no_grad() +@pytest.mark.parametrize("variant", VARIANTS, ids=lambda v: v.name) +def test_the_chain_inputs_are_released_before_the_last_control_block(transformer, inputs, context_4d, variant) -> None: + """Each chain owns its input, and the stream the refiner stage hands on is held only until the main stage takes + it: by the last control block the embedded context, the handed-on stream and the main chain's joined input are all + freed. Any of them still alive is an image-length hidden state more than the hint reservation allows for.""" + control = _control(variant) + refs: dict[str, weakref.ref] = {} + alive_at_last_block: list[dict[str, bool]] = [] + # The block whose output the refiner stage hands on: 2.0 runs its refiner chain over the main control blocks. + config = _config(variant) + refined_by_main_blocks = config["add_control_noise_refiner"] and not config["add_control_noise_refiner_correctly"] + refiner_last = control.control_layers[-1] if refined_by_main_blocks else control.control_noise_refiner[-1] + + def remember(name: str, *, first_call_only: bool = False): + # Returns nothing: a forward hook's return value would replace the block's output. + def hook(_module, _args, output) -> None: + if not (first_call_only and name in refs): + refs[name] = weakref.ref(output) + + return hook + + def remember_joined(_module, args) -> None: + refs["joined"] = weakref.ref(args[0]) + + def check(*_args) -> None: + alive_at_last_block.append({name: ref() is not None for name, ref in refs.items()}) + + taps = [ + control.control_all_x_embedder["2-1"].register_forward_hook(remember("embedded")), + # The first call is the refiner stage's; 2.0 calls the same block again in the main stage. + refiner_last.register_forward_hook(remember("handed on", first_call_only=True)), + # The last call is the main stage's: block 0's projection takes the joined stream. + control.control_layers[0].before_proj.register_forward_pre_hook(remember_joined), + control.control_layers[-1].register_forward_pre_hook(check), + ] + try: + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE): + _forward(transformer, inputs) + finally: + for tap in taps: + tap.remove() + + assert alive_at_last_block[-1] == {"embedded": False, "handed on": False, "joined": False} + + +@torch.no_grad() +def test_a_main_stage_without_its_refiner_stage_is_refused(transformer, inputs, context_4d) -> None: + """The main chain starts from the control stream the refiner stage refined. A forward that reaches `layers[0]` + without having run `noise_refiner[0]` with the control active has none, and must not run on a stale one.""" + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE) as hooks: + hooks.active = False + # Switched on after the refiner stage's entry has passed, inactive, on its way. + tap = transformer.noise_refiner[0].register_forward_pre_hook(lambda *_args: setattr(hooks, "active", True)) + try: + with pytest.raises(RuntimeError, match="stage 1 was entered without"): + _forward(transformer, inputs) + finally: + tap.remove() + + +@torch.no_grad() +def test_a_control_context_of_another_size_is_refused(transformer, inputs) -> None: + """Encoded at another size, the control tokens do not line up with the image stream's.""" + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + other_size = control.pack(torch.randn(1, 33, LATENT_SIDE + 4, LATENT_SIDE)) + + with pytest.raises(ValueError, match="35 image tokens"): + with install_fun_control(transformer, control, other_size, scale=SCALE): + _forward(transformer, inputs) + assert _hook_count(transformer) == 0 + + +def test_a_control_for_another_transformer_is_refused_at_install(transformer, context_4d) -> None: + wider = ZImageFunControl(Z_IMAGE_FUN_CONTROL_2_1, hidden_size=128, num_attention_heads=4) + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + + with ( + pytest.raises(ValueError, match="128 wide"), + install_fun_control(transformer, wider, control.pack(context_4d), scale=1.0), + ): + pass + with ( + pytest.raises(ValueError, match="132"), + install_fun_control(transformer, control, control.pack(context_4d)[..., :64], scale=1.0), + ): + pass + assert _hook_count(transformer) == 0 + + +@torch.no_grad() +def test_the_control_blocks_take_the_regional_bias_on_the_positive_pass_only(transformer, inputs, context_4d) -> None: + """The control blocks take what the base block took: the regional attention bias the patched forward hands + `layers[0]` on the positive pass, nothing on the negative pass, and nothing in the refiner stage, which regional + prompting does not mask.""" + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + cap = inputs["cap_feats"][0] + negative = {**inputs, "cap_feats": [torch.randn(CAPTION_TOKENS, CAPTION_DIM)]} + side = IMAGE_TOKENS + CAPTION_TOKENS + restricting = torch.ones(side, side, dtype=torch.float16) + restricting[:IMAGE_TOKENS, IMAGE_TOKENS:] = 0 + restricting[IMAGE_TOKENS:, :IMAGE_TOKENS] = 0 + masks: dict[str, list] = {"refiner": [], "main": [], "base": []} + taps = [ + control.control_noise_refiner[0].register_forward_pre_hook( + lambda _m, _args, kwargs: masks["refiner"].append(kwargs["attn_mask"]), with_kwargs=True + ), + control.control_layers[0].register_forward_pre_hook( + lambda _m, _args, kwargs: masks["main"].append(kwargs["attn_mask"]), with_kwargs=True + ), + transformer.layers[0].register_forward_pre_hook(lambda _m, args: masks["base"].append(args[1])), + ] + try: + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE): + plain = _forward(transformer, inputs) + with patch_transformer_for_regional_prompting( + transformer, restricting, IMAGE_TOKENS, positive_cap_feats=cap + ): + positive = _forward(transformer, inputs) + _forward(transformer, negative) + finally: + for tap in taps: + tap.remove() + + assert not torch.allclose(positive, plain, atol=1e-3) + assert masks["refiner"] == [None, None, None] + assert masks["base"][0] is None and masks["base"][2] is None + assert masks["base"][1] is not None + assert torch.equal(masks["main"][1], masks["base"][1]) + assert masks["main"][0] is None and masks["main"][2] is None + + +@torch.no_grad() +def test_an_unrestricted_regional_mask_leaves_the_controlled_forward_as_it_was(transformer) -> None: + """With nothing to restrict and no padding the regional bias is all zeros, so the patched forward computes the + plain one's control. Unpadded on purpose: the bias also keeps real tokens from attending to padding rows, which the + plain forward does not.""" + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + torch.manual_seed(4) + side, caption_tokens = 16, 32 # 64 image tokens and a 32-token caption, both multiples of 32 + inputs = { + "x": [torch.randn(16, 1, side, side)], + "t": torch.tensor([0.4]), + "cap_feats": [torch.randn(caption_tokens, CAPTION_DIM)], + } + image_tokens = (side // 2) ** 2 + unrestricted = torch.ones(image_tokens + caption_tokens, image_tokens + caption_tokens, dtype=torch.float16) + + with install_fun_control(transformer, control, control.pack(torch.randn(1, 33, side, side)), scale=SCALE): + plain = _forward(transformer, inputs) + with patch_transformer_for_regional_prompting( + transformer, unrestricted, image_tokens, positive_cap_feats=inputs["cap_feats"][0] + ): + regional = _forward(transformer, inputs) + + torch.testing.assert_close(regional, plain, rtol=1e-4, atol=1e-5) + + +def test_pack_patches_the_context_as_the_transformer_patches_its_latents(transformer, context_4d) -> None: + for variant in (Z_IMAGE_FUN_CONTROL_1_0, Z_IMAGE_FUN_CONTROL_2_1): + control = ZImageFunControl(variant, hidden_size=HIDDEN, num_attention_heads=HEADS) + channels = _config(variant)["control_in_dim"] + expected, _, _ = transformer._patchify_image(context_4d[0, :channels].unsqueeze(1), 2, 1) + + assert torch.equal(control.pack(context_4d)[0], expected) + + +def _as_meta(keys: dict) -> dict[str, torch.Tensor]: + return {name: torch.empty(shape, dtype=torch.bfloat16, device="meta") for name, (shape, _dtype) in keys.items()} + + +@pytest.mark.parametrize( + ("variant", "keys"), + [ + (Z_IMAGE_FUN_CONTROL_1_0, keys_1_0), + (Z_IMAGE_FUN_CONTROL_2_1, keys_2_1), + (Z_IMAGE_FUN_CONTROL_2_1_LITE, keys_2_1_lite), + # The 2.0 file has the 2.1 keys; its refiner blocks were never trained and are not built. + ( + Z_IMAGE_FUN_CONTROL_2_0, + {name: value for name, value in keys_2_1.items() if not name.startswith("control_noise_refiner.")}, + ), + ], + ids=["1.0", "2.1", "2.1 Lite", "2.0"], +) +def test_each_variant_names_and_shapes_its_tensors_as_the_released_file_does(variant, keys) -> None: + """Against the headers of alibaba-pai's files, on the meta device so nothing of the 6.7 GB is allocated.""" + with accelerate.init_empty_weights(): + control = ZImageFunControl(variant, hidden_size=3840, num_attention_heads=30) + + assert {name: list(tensor.shape) for name, tensor in control.state_dict().items()} == { + name: shape for name, (shape, _dtype) in keys.items() + } + + +def test_the_1_0_header_reads_as_its_variant_at_full_width() -> None: + """1.0 is the one layout the loader can tell apart without reading a value, so its header goes through the whole + `from_state_dict` on the meta device.""" + control = ZImageFunControl.from_state_dict(_as_meta(keys_1_0), dtype=torch.bfloat16) + + assert control.variant == Z_IMAGE_FUN_CONTROL_1_0 + assert (control.hidden_size, control.control_in_dim, control.control_layers[0].attention.heads) == (3840, 64, 30) + + +def _state_dict(variant: ZImageFunControlVariant, seed: int = 1) -> dict[str, torch.Tensor]: + return {name: tensor.clone() for name, tensor in _control(variant, seed).state_dict().items()} + + +def _with_untrained_refiner(sd: dict) -> dict: + return { + name: torch.zeros_like(tensor) if name.startswith("control_noise_refiner.") and "_proj." in name else tensor + for name, tensor in sd.items() + } + + +@torch.no_grad() +@pytest.mark.parametrize( + ("released", "variant"), + [ + (lambda: _state_dict(Z_IMAGE_FUN_CONTROL_1_0), Z_IMAGE_FUN_CONTROL_1_0), + (lambda: _with_untrained_refiner(_state_dict(Z_IMAGE_FUN_CONTROL_2_1)), Z_IMAGE_FUN_CONTROL_2_0), + (lambda: _state_dict(Z_IMAGE_FUN_CONTROL_2_1), Z_IMAGE_FUN_CONTROL_2_1), + (lambda: _state_dict(Z_IMAGE_FUN_CONTROL_2_1_LITE), Z_IMAGE_FUN_CONTROL_2_1_LITE), + ], + ids=["1.0", "2.0", "2.1", "2.1 Lite"], +) +def test_from_state_dict_reads_the_variant_off_the_tensors(transformer, inputs, context_4d, released, variant) -> None: + """2.0 and 2.1 have the same keys; 2.0's refiner projections are all zero, and it runs its refiner chain over the + main control blocks. Each loaded control computes its variant's VideoX-Fun forward.""" + sd = released() + + loaded = ZImageFunControl.from_state_dict(sd, dtype=torch.float32, attention_head_dim=HEAD_DIM) + + assert loaded.variant == variant + assert not loaded.training + assert len(loaded.control_noise_refiner) == (0 if variant is Z_IMAGE_FUN_CONTROL_2_0 else 2) + with install_fun_control(transformer, loaded, loaded.pack(context_4d), scale=SCALE): + hooked = _forward(transformer, inputs) + assert torch.equal(hooked, _videox_forward(transformer, loaded, _config(variant), context_4d, SCALE, **inputs)) + + +def test_from_state_dict_casts_to_the_requested_dtype() -> None: + loaded = ZImageFunControl.from_state_dict( + _state_dict(Z_IMAGE_FUN_CONTROL_2_1_LITE), dtype=torch.float16, attention_head_dim=HEAD_DIM + ) + assert {parameter.dtype for parameter in loaded.parameters()} == {torch.float16} + + +@torch.no_grad() +def test_the_chain_runs_in_the_control_dtype_and_hands_hints_back_in_the_base_dtype( + transformer, inputs, context_4d +) -> None: + """The loader follows the device's bfloat16-safe dtype while the node runs bfloat16, so the control may meet + activations of another dtype. Checked with a float64 control against the float32 base: the same arithmetic in a + dtype every CPU has a GEMM for.""" + control = _control(Z_IMAGE_FUN_CONTROL_2_1) + with install_fun_control(transformer, control, control.pack(context_4d), scale=SCALE): + expected = _forward(transformer, inputs) + wider = ZImageFunControl.from_state_dict(control.state_dict(), dtype=torch.float64, attention_head_dim=HEAD_DIM) + + with install_fun_control(transformer, wider, wider.pack(context_4d), scale=SCALE): + hooked = _forward(transformer, inputs) + + assert hooked.dtype == torch.float32 + assert torch.allclose(hooked, expected, atol=1e-4, rtol=1e-4) + + +def _without(sd: dict, prefix: str) -> dict: + return {name: tensor for name, tensor in sd.items() if not name.startswith(prefix)} + + +@pytest.mark.parametrize( + ("rewrite", "message"), + [ + (lambda: _without(_state_dict(Z_IMAGE_FUN_CONTROL_2_1), "control_all_x_embedder."), "is missing"), + (lambda: _without(_state_dict(Z_IMAGE_FUN_CONTROL_2_1), "control_layers.14."), "14 control blocks"), + ( + lambda: _without(_state_dict(Z_IMAGE_FUN_CONTROL_2_1), "control_noise_refiner.1.after_proj."), + "incomplete refiner projections", + ), + ( + lambda: _with_untrained_refiner(_state_dict(Z_IMAGE_FUN_CONTROL_2_1_LITE)), + "3 control blocks and untrained", + ), + ( + lambda: { + **_state_dict(Z_IMAGE_FUN_CONTROL_1_0), + "control_all_x_embedder.2-1.weight": torch.zeros(HIDDEN, 132), + }, + "takes 132 features, with 6 control blocks", + ), + ( + lambda: _without(_state_dict(Z_IMAGE_FUN_CONTROL_2_1), "control_layers.3.feed_forward."), + "missing", + ), + ( + lambda: {**_state_dict(Z_IMAGE_FUN_CONTROL_2_1), "x_pad_token": torch.zeros(1, HIDDEN)}, + "unexpected", + ), + ], + ids=[ + "no-embedder", + "block-count", + "partial-projections", + "untrained-lite", + "1.0-with-33-channels", + "missing", + "stray", + ], +) +def test_from_state_dict_refuses_what_is_not_a_published_layout(rewrite, message) -> None: + with pytest.raises(ValueError, match=message): + ZImageFunControl.from_state_dict(rewrite(), dtype=torch.float32, attention_head_dim=HEAD_DIM) + + +def test_a_width_the_heads_do_not_divide_is_refused() -> None: + with pytest.raises(ValueError, match="128-wide attention heads"): + ZImageFunControl.from_state_dict(_state_dict(Z_IMAGE_FUN_CONTROL_2_1), dtype=torch.float32) diff --git a/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_1_0_keys.py b/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_1_0_keys.py new file mode 100644 index 00000000000..7911aaf8647 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_1_0_keys.py @@ -0,0 +1,152 @@ +"""Key layout of alibaba-pai's first Z-Image-Turbo Fun-ControlNet-Union (1.0). + +Captured on 2026-10-19 with `scripts/capture_state_dict_fixture.py`, +which reads the header only, from +`alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union/Z-Image-Turbo-Fun-Controlnet-Union.safetensors`. + +Subsetting rule: none -- every key in the file. Values are `(shape, dtype)`. + +136 bf16 tensors and no metadata: `control_all_x_embedder.2-1` takes the 16 control-latent channels packed 2x2 (64 +features), `control_layers.0..5` carry `after_proj` (and `before_proj` on block 0), and `control_noise_refiner.0..1` +are plain blocks without projections, which neither join the image stream nor add hints. VideoX-Fun's +`z_image_control.yaml` places the six hints after base layers 0, 5, 10, 15, 20 and 25. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "control_all_x_embedder.2-1.bias": ([3840], "BF16"), + "control_all_x_embedder.2-1.weight": ([3840, 64], "BF16"), + "control_layers.0.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.0.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.0.after_proj.bias": ([3840], "BF16"), + "control_layers.0.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.norm_k.weight": ([128], "BF16"), + "control_layers.0.attention.norm_q.weight": ([128], "BF16"), + "control_layers.0.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention_norm1.weight": ([3840], "BF16"), + "control_layers.0.attention_norm2.weight": ([3840], "BF16"), + "control_layers.0.before_proj.bias": ([3840], "BF16"), + "control_layers.0.before_proj.weight": ([3840, 3840], "BF16"), + "control_layers.0.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.0.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.0.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.0.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.0.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.1.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.1.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.1.after_proj.bias": ([3840], "BF16"), + "control_layers.1.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.norm_k.weight": ([128], "BF16"), + "control_layers.1.attention.norm_q.weight": ([128], "BF16"), + "control_layers.1.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention_norm1.weight": ([3840], "BF16"), + "control_layers.1.attention_norm2.weight": ([3840], "BF16"), + "control_layers.1.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.1.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.1.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.1.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.1.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.2.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.2.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.2.after_proj.bias": ([3840], "BF16"), + "control_layers.2.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.norm_k.weight": ([128], "BF16"), + "control_layers.2.attention.norm_q.weight": ([128], "BF16"), + "control_layers.2.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention_norm1.weight": ([3840], "BF16"), + "control_layers.2.attention_norm2.weight": ([3840], "BF16"), + "control_layers.2.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.2.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.2.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.2.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.2.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.3.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.3.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.3.after_proj.bias": ([3840], "BF16"), + "control_layers.3.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.norm_k.weight": ([128], "BF16"), + "control_layers.3.attention.norm_q.weight": ([128], "BF16"), + "control_layers.3.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention_norm1.weight": ([3840], "BF16"), + "control_layers.3.attention_norm2.weight": ([3840], "BF16"), + "control_layers.3.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.3.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.3.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.3.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.3.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.4.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.4.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.4.after_proj.bias": ([3840], "BF16"), + "control_layers.4.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.norm_k.weight": ([128], "BF16"), + "control_layers.4.attention.norm_q.weight": ([128], "BF16"), + "control_layers.4.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention_norm1.weight": ([3840], "BF16"), + "control_layers.4.attention_norm2.weight": ([3840], "BF16"), + "control_layers.4.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.4.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.4.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.4.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.4.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.5.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.5.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.5.after_proj.bias": ([3840], "BF16"), + "control_layers.5.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.norm_k.weight": ([128], "BF16"), + "control_layers.5.attention.norm_q.weight": ([128], "BF16"), + "control_layers.5.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention_norm1.weight": ([3840], "BF16"), + "control_layers.5.attention_norm2.weight": ([3840], "BF16"), + "control_layers.5.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.5.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.5.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.5.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.5.ffn_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.0.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_noise_refiner.0.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_noise_refiner.0.attention.norm_k.weight": ([128], "BF16"), + "control_noise_refiner.0.attention.norm_q.weight": ([128], "BF16"), + "control_noise_refiner.0.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.0.attention_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.0.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.0.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_noise_refiner.0.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.0.ffn_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.0.ffn_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.1.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_noise_refiner.1.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_noise_refiner.1.attention.norm_k.weight": ([128], "BF16"), + "control_noise_refiner.1.attention.norm_q.weight": ([128], "BF16"), + "control_noise_refiner.1.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.1.attention_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.1.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.1.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_noise_refiner.1.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.1.ffn_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.1.ffn_norm2.weight": ([3840], "BF16"), +} diff --git a/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_2_1_keys.py b/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_2_1_keys.py new file mode 100644 index 00000000000..fbd38d19b70 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_2_1_keys.py @@ -0,0 +1,312 @@ +"""Key layout of alibaba-pai's Z-Image-Turbo Fun-ControlNet-Union 2.1 (8-step build), the VACE-style control. + +Captured on 2026-10-19 with `scripts/capture_state_dict_fixture.py`, +which reads the header only, from +`alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.1/Z-Image-Turbo-Fun-Controlnet-Union-2.1-8steps.safetensors`. + +Subsetting rule: none -- every key in the file. Values are `(shape, dtype)`. + +295 bf16 tensors and no metadata: `control_all_x_embedder.2-1` takes 33 latent channels packed 2x2 (132 features) into +the 3840-wide stream, `control_layers.0..14` and `control_noise_refiner.0..1` are diffusers' `ZImageTransformerBlock` +under its own names, with `after_proj` on every block of both chains and `before_proj` on block 0 of each. No +`x_pad_token`: the control takes the base's. The Tile 2.1 builds and the 2.0 file have exactly these keys and shapes; +2.0 differs only in values, its refiner projections being all zero, which is how the loader tells it apart. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "control_all_x_embedder.2-1.bias": ([3840], "BF16"), + "control_all_x_embedder.2-1.weight": ([3840, 132], "BF16"), + "control_layers.0.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.0.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.0.after_proj.bias": ([3840], "BF16"), + "control_layers.0.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.norm_k.weight": ([128], "BF16"), + "control_layers.0.attention.norm_q.weight": ([128], "BF16"), + "control_layers.0.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention_norm1.weight": ([3840], "BF16"), + "control_layers.0.attention_norm2.weight": ([3840], "BF16"), + "control_layers.0.before_proj.bias": ([3840], "BF16"), + "control_layers.0.before_proj.weight": ([3840, 3840], "BF16"), + "control_layers.0.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.0.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.0.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.0.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.0.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.1.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.1.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.1.after_proj.bias": ([3840], "BF16"), + "control_layers.1.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.norm_k.weight": ([128], "BF16"), + "control_layers.1.attention.norm_q.weight": ([128], "BF16"), + "control_layers.1.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention_norm1.weight": ([3840], "BF16"), + "control_layers.1.attention_norm2.weight": ([3840], "BF16"), + "control_layers.1.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.1.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.1.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.1.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.1.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.10.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.10.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.10.after_proj.bias": ([3840], "BF16"), + "control_layers.10.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.10.attention.norm_k.weight": ([128], "BF16"), + "control_layers.10.attention.norm_q.weight": ([128], "BF16"), + "control_layers.10.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.10.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.10.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.10.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.10.attention_norm1.weight": ([3840], "BF16"), + "control_layers.10.attention_norm2.weight": ([3840], "BF16"), + "control_layers.10.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.10.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.10.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.10.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.10.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.11.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.11.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.11.after_proj.bias": ([3840], "BF16"), + "control_layers.11.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.11.attention.norm_k.weight": ([128], "BF16"), + "control_layers.11.attention.norm_q.weight": ([128], "BF16"), + "control_layers.11.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.11.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.11.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.11.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.11.attention_norm1.weight": ([3840], "BF16"), + "control_layers.11.attention_norm2.weight": ([3840], "BF16"), + "control_layers.11.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.11.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.11.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.11.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.11.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.12.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.12.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.12.after_proj.bias": ([3840], "BF16"), + "control_layers.12.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.12.attention.norm_k.weight": ([128], "BF16"), + "control_layers.12.attention.norm_q.weight": ([128], "BF16"), + "control_layers.12.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.12.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.12.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.12.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.12.attention_norm1.weight": ([3840], "BF16"), + "control_layers.12.attention_norm2.weight": ([3840], "BF16"), + "control_layers.12.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.12.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.12.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.12.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.12.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.13.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.13.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.13.after_proj.bias": ([3840], "BF16"), + "control_layers.13.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.13.attention.norm_k.weight": ([128], "BF16"), + "control_layers.13.attention.norm_q.weight": ([128], "BF16"), + "control_layers.13.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.13.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.13.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.13.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.13.attention_norm1.weight": ([3840], "BF16"), + "control_layers.13.attention_norm2.weight": ([3840], "BF16"), + "control_layers.13.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.13.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.13.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.13.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.13.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.14.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.14.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.14.after_proj.bias": ([3840], "BF16"), + "control_layers.14.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.14.attention.norm_k.weight": ([128], "BF16"), + "control_layers.14.attention.norm_q.weight": ([128], "BF16"), + "control_layers.14.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.14.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.14.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.14.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.14.attention_norm1.weight": ([3840], "BF16"), + "control_layers.14.attention_norm2.weight": ([3840], "BF16"), + "control_layers.14.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.14.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.14.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.14.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.14.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.2.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.2.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.2.after_proj.bias": ([3840], "BF16"), + "control_layers.2.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.norm_k.weight": ([128], "BF16"), + "control_layers.2.attention.norm_q.weight": ([128], "BF16"), + "control_layers.2.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention_norm1.weight": ([3840], "BF16"), + "control_layers.2.attention_norm2.weight": ([3840], "BF16"), + "control_layers.2.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.2.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.2.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.2.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.2.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.3.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.3.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.3.after_proj.bias": ([3840], "BF16"), + "control_layers.3.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.norm_k.weight": ([128], "BF16"), + "control_layers.3.attention.norm_q.weight": ([128], "BF16"), + "control_layers.3.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.3.attention_norm1.weight": ([3840], "BF16"), + "control_layers.3.attention_norm2.weight": ([3840], "BF16"), + "control_layers.3.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.3.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.3.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.3.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.3.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.4.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.4.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.4.after_proj.bias": ([3840], "BF16"), + "control_layers.4.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.norm_k.weight": ([128], "BF16"), + "control_layers.4.attention.norm_q.weight": ([128], "BF16"), + "control_layers.4.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.4.attention_norm1.weight": ([3840], "BF16"), + "control_layers.4.attention_norm2.weight": ([3840], "BF16"), + "control_layers.4.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.4.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.4.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.4.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.4.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.5.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.5.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.5.after_proj.bias": ([3840], "BF16"), + "control_layers.5.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.norm_k.weight": ([128], "BF16"), + "control_layers.5.attention.norm_q.weight": ([128], "BF16"), + "control_layers.5.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.5.attention_norm1.weight": ([3840], "BF16"), + "control_layers.5.attention_norm2.weight": ([3840], "BF16"), + "control_layers.5.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.5.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.5.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.5.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.5.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.6.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.6.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.6.after_proj.bias": ([3840], "BF16"), + "control_layers.6.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.6.attention.norm_k.weight": ([128], "BF16"), + "control_layers.6.attention.norm_q.weight": ([128], "BF16"), + "control_layers.6.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.6.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.6.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.6.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.6.attention_norm1.weight": ([3840], "BF16"), + "control_layers.6.attention_norm2.weight": ([3840], "BF16"), + "control_layers.6.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.6.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.6.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.6.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.6.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.7.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.7.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.7.after_proj.bias": ([3840], "BF16"), + "control_layers.7.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.7.attention.norm_k.weight": ([128], "BF16"), + "control_layers.7.attention.norm_q.weight": ([128], "BF16"), + "control_layers.7.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.7.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.7.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.7.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.7.attention_norm1.weight": ([3840], "BF16"), + "control_layers.7.attention_norm2.weight": ([3840], "BF16"), + "control_layers.7.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.7.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.7.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.7.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.7.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.8.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.8.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.8.after_proj.bias": ([3840], "BF16"), + "control_layers.8.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.8.attention.norm_k.weight": ([128], "BF16"), + "control_layers.8.attention.norm_q.weight": ([128], "BF16"), + "control_layers.8.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.8.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.8.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.8.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.8.attention_norm1.weight": ([3840], "BF16"), + "control_layers.8.attention_norm2.weight": ([3840], "BF16"), + "control_layers.8.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.8.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.8.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.8.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.8.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.9.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.9.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.9.after_proj.bias": ([3840], "BF16"), + "control_layers.9.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.9.attention.norm_k.weight": ([128], "BF16"), + "control_layers.9.attention.norm_q.weight": ([128], "BF16"), + "control_layers.9.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.9.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.9.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.9.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.9.attention_norm1.weight": ([3840], "BF16"), + "control_layers.9.attention_norm2.weight": ([3840], "BF16"), + "control_layers.9.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.9.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.9.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.9.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.9.ffn_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.0.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_noise_refiner.0.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_noise_refiner.0.after_proj.bias": ([3840], "BF16"), + "control_noise_refiner.0.after_proj.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.norm_k.weight": ([128], "BF16"), + "control_noise_refiner.0.attention.norm_q.weight": ([128], "BF16"), + "control_noise_refiner.0.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.0.attention_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.0.before_proj.bias": ([3840], "BF16"), + "control_noise_refiner.0.before_proj.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.0.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_noise_refiner.0.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.0.ffn_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.0.ffn_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.1.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_noise_refiner.1.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_noise_refiner.1.after_proj.bias": ([3840], "BF16"), + "control_noise_refiner.1.after_proj.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.norm_k.weight": ([128], "BF16"), + "control_noise_refiner.1.attention.norm_q.weight": ([128], "BF16"), + "control_noise_refiner.1.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.1.attention_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.1.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.1.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_noise_refiner.1.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.1.ffn_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.1.ffn_norm2.weight": ([3840], "BF16"), +} diff --git a/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_2_1_lite_keys.py b/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_2_1_lite_keys.py new file mode 100644 index 00000000000..8a92ee7f9c6 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/z_image_fun_controlnet_2_1_lite_keys.py @@ -0,0 +1,106 @@ +"""Key layout of alibaba-pai's Z-Image-Turbo Fun-ControlNet-Union 2.1 Lite (2602, 8-step build). + +Captured on 2026-10-19 with `scripts/capture_state_dict_fixture.py`, +which reads the header only, from +`alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.1/Z-Image-Turbo-Fun-Controlnet-Union-2.1-lite-2602-8steps.safetensors`. + +Subsetting rule: none -- every key in the file. Values are `(shape, dtype)`. + +91 bf16 tensors and no metadata: the 2.1 layout with three `control_layers` instead of fifteen, whose hints VideoX-Fun's +`z_image_control_2.1_lite.yaml` places after base layers 0, 10 and 20. The block count is what the loader reads the +variant from. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "control_all_x_embedder.2-1.bias": ([3840], "BF16"), + "control_all_x_embedder.2-1.weight": ([3840, 132], "BF16"), + "control_layers.0.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.0.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.0.after_proj.bias": ([3840], "BF16"), + "control_layers.0.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.norm_k.weight": ([128], "BF16"), + "control_layers.0.attention.norm_q.weight": ([128], "BF16"), + "control_layers.0.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.0.attention_norm1.weight": ([3840], "BF16"), + "control_layers.0.attention_norm2.weight": ([3840], "BF16"), + "control_layers.0.before_proj.bias": ([3840], "BF16"), + "control_layers.0.before_proj.weight": ([3840, 3840], "BF16"), + "control_layers.0.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.0.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.0.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.0.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.0.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.1.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.1.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.1.after_proj.bias": ([3840], "BF16"), + "control_layers.1.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.norm_k.weight": ([128], "BF16"), + "control_layers.1.attention.norm_q.weight": ([128], "BF16"), + "control_layers.1.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.1.attention_norm1.weight": ([3840], "BF16"), + "control_layers.1.attention_norm2.weight": ([3840], "BF16"), + "control_layers.1.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.1.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.1.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.1.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.1.ffn_norm2.weight": ([3840], "BF16"), + "control_layers.2.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_layers.2.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_layers.2.after_proj.bias": ([3840], "BF16"), + "control_layers.2.after_proj.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.norm_k.weight": ([128], "BF16"), + "control_layers.2.attention.norm_q.weight": ([128], "BF16"), + "control_layers.2.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_layers.2.attention_norm1.weight": ([3840], "BF16"), + "control_layers.2.attention_norm2.weight": ([3840], "BF16"), + "control_layers.2.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_layers.2.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_layers.2.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_layers.2.ffn_norm1.weight": ([3840], "BF16"), + "control_layers.2.ffn_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.0.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_noise_refiner.0.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_noise_refiner.0.after_proj.bias": ([3840], "BF16"), + "control_noise_refiner.0.after_proj.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.norm_k.weight": ([128], "BF16"), + "control_noise_refiner.0.attention.norm_q.weight": ([128], "BF16"), + "control_noise_refiner.0.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.attention_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.0.attention_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.0.before_proj.bias": ([3840], "BF16"), + "control_noise_refiner.0.before_proj.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.0.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.0.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_noise_refiner.0.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.0.ffn_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.0.ffn_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.1.adaLN_modulation.0.bias": ([15360], "BF16"), + "control_noise_refiner.1.adaLN_modulation.0.weight": ([15360, 256], "BF16"), + "control_noise_refiner.1.after_proj.bias": ([3840], "BF16"), + "control_noise_refiner.1.after_proj.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.norm_k.weight": ([128], "BF16"), + "control_noise_refiner.1.attention.norm_q.weight": ([128], "BF16"), + "control_noise_refiner.1.attention.to_k.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_out.0.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_q.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention.to_v.weight": ([3840, 3840], "BF16"), + "control_noise_refiner.1.attention_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.1.attention_norm2.weight": ([3840], "BF16"), + "control_noise_refiner.1.feed_forward.w1.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.1.feed_forward.w2.weight": ([3840, 10240], "BF16"), + "control_noise_refiner.1.feed_forward.w3.weight": ([10240, 3840], "BF16"), + "control_noise_refiner.1.ffn_norm1.weight": ([3840], "BF16"), + "control_noise_refiner.1.ffn_norm2.weight": ([3840], "BF16"), +} diff --git a/tests/backend/model_manager/load/test_fp8_capability.py b/tests/backend/model_manager/load/test_fp8_capability.py index d1c15f02704..7b2587fbe5c 100644 --- a/tests/backend/model_manager/load/test_fp8_capability.py +++ b/tests/backend/model_manager/load/test_fp8_capability.py @@ -281,7 +281,7 @@ def test_a_loader_that_implements_the_cast_is_supported(self) -> None: def test_a_declared_loader_is_refused_with_its_own_reason(self) -> None: verdict = fp8_storage_verdict(BaseModelType.ZImage, ModelType.ControlNet, ModelFormat.Checkpoint) assert verdict.supported is False - assert verdict.reason is not None and "never cast" in verdict.reason + assert verdict.reason is not None and "Union and Tile 2.1" in verdict.reason def test_a_quantized_format_is_refused_although_its_loader_declares_nothing(self) -> None: """The rule the loaders do not have to repeat. `WanGGUFCheckpointModel` declares nothing and diff --git a/tests/backend/model_manager/load/test_fun_controlnet_loader.py b/tests/backend/model_manager/load/test_fun_controlnet_loader.py index 4734546f2e9..75929bf395a 100644 --- a/tests/backend/model_manager/load/test_fun_controlnet_loader.py +++ b/tests/backend/model_manager/load/test_fun_controlnet_loader.py @@ -2,7 +2,7 @@ tensors, the dtype they cast to, and the quantized files they refuse rather than load unscaled. Like the Z-Image ControlNet loader it builds under `init_empty_weights` and assigns the file's tensors, so an fp8 -weight with its scale dropped would become the parameter as it is (`test_z_image_control_adapter_loader.py`). No +weight with its scale dropped would become the parameter as it is (`test_z_image_fun_controlnet_loader.py`). No quantized build of this adapter is published (alibaba-pai ships two dense bf16 files), so refusing costs nothing. """ diff --git a/tests/backend/model_manager/load/test_load_default_fp8.py b/tests/backend/model_manager/load/test_load_default_fp8.py index 0ac15ed4ca9..d033f715ec0 100644 --- a/tests/backend/model_manager/load/test_load_default_fp8.py +++ b/tests/backend/model_manager/load/test_load_default_fp8.py @@ -588,11 +588,11 @@ def test_should_use_fp8_allows_z_image(): def test_should_use_fp8_refuses_a_loader_that_declared_it_does_not_implement_the_cast(): - """A Z-Image ControlNet asks for fp8 and its loader never casts, so the gate must say no. + """A Z-Image ControlNet asks for fp8 and its loader does not implement the cast, so the gate must say no. Not the same question as the one above: Z-Image *main* implements the cast. Both are `base=z-image`, which is why the answer cannot live on the architecture -- it lives on the - loader key, and `ZImageControlCheckpointModel` declares `Unimplemented` there. + loader key, and `ZImageFunControlNetLoader` declares `Unimplemented` there. The device probe must not run for it either: an inert request is no reason to touch the GPU. """ diff --git a/tests/backend/model_manager/load/test_z_image_control_adapter_loader.py b/tests/backend/model_manager/load/test_z_image_control_adapter_loader.py deleted file mode 100644 index 80ee958e982..00000000000 --- a/tests/backend/model_manager/load/test_z_image_control_adapter_loader.py +++ /dev/null @@ -1,89 +0,0 @@ -"""What the Z-Image ControlNet loader does with a quantized checkpoint. - -It handles no quantization side channel at all, and it is the worst-shaped of the loaders that do -not: it builds the adapter under `init_empty_weights` and then calls -`load_state_dict(sd, assign=True, strict=False)`. `assign=True` means the checkpoint's tensors -*become* the parameters, so an fp8 weight stays fp8 -- there is no cast anywhere to widen it -- and -`strict=False` sends the orphaned `weight_scale` to `unexpected_keys`, where `log_unexpected_keys` -reports it at DEBUG and drops it. The result is an adapter whose every quantized weight is off by -`1/weight_scale` -- orders of magnitude -- and nothing says so at the default log level. - -No quantized Z-Image ControlNet is published today (checked 2026-09-18: `alibaba-pai`'s ten builds -are all dense, and `Comfy-Org/z_image_turbo` ships none), so refusing costs nothing that works. -""" - -from pathlib import Path - -import pytest -import safetensors.torch -import torch - -from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_ZImage_Config -from invokeai.backend.model_manager.load.model_loaders import z_image -from invokeai.backend.model_manager.load.model_loaders.z_image import ZImageControlCheckpointModel -from tests.fixtures.loader_seams import Seam, prepare -from tests.fixtures.quantized_payloads import comfy_quant_marker, quantize_convrot, quantize_scaled_fp8 - -SEAM = Seam( - loader=ZImageControlCheckpointModel, - module=z_image, - entry="_load_control_adapter", - # The loader imports `load_file` inside the method, so the name it resolves is the package's. - load_file_host=safetensors.torch, - casts_fp8_storage=False, -) - - -def _dense_adapter() -> dict[str, torch.Tensor]: - """The smallest state dict the loader's geometry probe accepts: one control block, two refiner - layers, and the embedder it reads `control_in_dim` off.""" - return { - "control_layers.0.weight": torch.ones(8, 8), - "control_noise_refiner.0.weight": torch.ones(8, 8), - "control_noise_refiner.1.weight": torch.ones(8, 8), - "control_all_x_embedder.2-1.weight": torch.ones(3840, 64), - } - - -def _driver(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, state_dict: dict[str, torch.Tensor]): - checkpoint = tmp_path / "z_image_controlnet.safetensors" - checkpoint.touch() - config = ControlNet_Checkpoint_ZImage_Config.model_construct(path=str(checkpoint), name="z-image-control") - return prepare(SEAM, monkeypatch, state_dict=state_dict), config - - -def test_a_scaled_fp8_checkpoint_is_refused_rather_than_loaded_unscaled(monkeypatch, tmp_path) -> None: - """`assign=True` would make the raw e4m3 codes the adapter's weights, and the scale that says - what they mean is dropped as an unexpected key.""" - payload = quantize_scaled_fp8(torch.randn(8, 8)) - state_dict = _dense_adapter() - state_dict["control_layers.0.weight"] = payload.codes - state_dict["control_layers.0.weight_scale"] = payload.scale - run, config = _driver(monkeypatch, tmp_path, state_dict) - - with pytest.raises(ValueError, match="quantization side channel"): - run.load(config) - - -def test_an_int8_checkpoint_is_refused_by_its_marker(monkeypatch, tmp_path) -> None: - """int8 is not even a dtype the adapter could run: `assign=True` would install an int8 tensor - where a float parameter belongs, and the un-rotation the scheme needs has no home here.""" - codes, scale, _ = quantize_convrot(torch.randn(8, 256)) - state_dict = _dense_adapter() - state_dict["control_layers.0.weight"] = codes - state_dict["control_layers.0.weight_scale"] = scale - state_dict["control_layers.0.comfy_quant"] = comfy_quant_marker({"format": "int8_tensorwise", "convrot": True}) - run, config = _driver(monkeypatch, tmp_path, state_dict) - - with pytest.raises(ValueError, match="quantization side channel"): - run.load(config) - - -def test_a_dense_checkpoint_still_loads(monkeypatch, tmp_path) -> None: - """The other half: the guard must not cost the builds that exist. A refusal that fires on every - checkpoint is not a guard, it is an outage.""" - run, config = _driver(monkeypatch, tmp_path, _dense_adapter()) - - model = run.load(config) - - assert model.__class__.__name__ == "ZImageControlAdapter" diff --git a/tests/backend/model_manager/load/test_z_image_fun_controlnet_loader.py b/tests/backend/model_manager/load/test_z_image_fun_controlnet_loader.py new file mode 100644 index 00000000000..b92cb1c71c4 --- /dev/null +++ b/tests/backend/model_manager/load/test_z_image_fun_controlnet_loader.py @@ -0,0 +1,118 @@ +"""What the Z-Image Fun-ControlNet loader makes of a checkpoint: the layout it reads off the tensors, the dtype it casts +to, and the quantized files it refuses rather than load unscaled. + +It builds under `init_empty_weights` and assigns the file's tensors, so an fp8 weight whose scale was dropped would +become the parameter as it is, off by `1/weight_scale`. No quantized Z-Image ControlNet is published (checked +2026-09-18: alibaba-pai's builds are all dense bf16, and `Comfy-Org/z_image_turbo` ships none), so refusing costs +nothing that works. +""" + +from pathlib import Path + +import pytest +import safetensors.torch +import torch + +from invokeai.backend.fun_controlnet.z_image import ( + Z_IMAGE_FUN_CONTROL_2_0, + Z_IMAGE_FUN_CONTROL_2_1, + Z_IMAGE_FUN_CONTROL_2_1_LITE, + ZImageFunControl, +) +from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_ZImage_Config +from invokeai.backend.model_manager.load.model_loaders import z_image +from invokeai.backend.model_manager.load.model_loaders.z_image import ZImageFunControlNetLoader +from tests.fixtures.loader_seams import Seam, prepare +from tests.fixtures.quantized_payloads import comfy_quant_marker, quantize_convrot, quantize_scaled_fp8 + +# One 128-wide head: the narrowest control the loader's head rule accepts. +HIDDEN = 128 +QUANTIZED_LAYER = "control_layers.0.attention.to_q" + +SEAM = Seam( + loader=ZImageFunControlNetLoader, + module=z_image, + entry="_load_model", + # The loader imports `load_file` inside the method, so the name it resolves is the package's. + load_file_host=safetensors.torch, + compute_dtype=torch.bfloat16, + # The loader takes its dtype from the device and never reads `self._torch_dtype`. + patches_device=True, + sets_torch_dtype=False, + casts_fp8_storage=False, +) + + +def _checkpoint(variant=Z_IMAGE_FUN_CONTROL_2_1_LITE) -> dict[str, torch.Tensor]: + """A float32 file with the variant's real layout; the loader has to cast it.""" + torch.manual_seed(0) + control = ZImageFunControl(variant, hidden_size=HIDDEN, num_attention_heads=1) + return {name: tensor.detach().clone() for name, tensor in control.state_dict().items()} + + +def _driver(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, state_dict: dict[str, torch.Tensor]): + checkpoint = tmp_path / "z_image_controlnet.safetensors" + checkpoint.touch() + config = ControlNet_Checkpoint_ZImage_Config.model_construct(path=str(checkpoint), name="z-image-control") + return prepare(SEAM, monkeypatch, state_dict=state_dict), config + + +def test_a_dense_checkpoint_loads_as_its_layout_in_the_compute_dtype(monkeypatch, tmp_path) -> None: + run, config = _driver(monkeypatch, tmp_path, _checkpoint()) + + model = run.load(config) + + assert isinstance(model, ZImageFunControl) + assert model.variant == Z_IMAGE_FUN_CONTROL_2_1_LITE + assert (model.hidden_size, model.control_in_dim, model.control_layers[0].attention.heads) == (HIDDEN, 132, 1) + assert {parameter.dtype for parameter in model.parameters()} == {torch.bfloat16} + assert not model.training + + +def test_the_2_0_layout_loads_without_its_untrained_refiner(monkeypatch, tmp_path) -> None: + """The 2.0 file has the 2.1 keys, with the refiner projections all zero; its refiner chain runs over the main + control blocks, so the refiner blocks are dropped rather than kept resident.""" + state_dict = { + name: torch.zeros_like(tensor) if name.startswith("control_noise_refiner.") and "_proj." in name else tensor + for name, tensor in _checkpoint(Z_IMAGE_FUN_CONTROL_2_1).items() + } + run, config = _driver(monkeypatch, tmp_path, state_dict) + + model = run.load(config) + + assert model.variant == Z_IMAGE_FUN_CONTROL_2_0 + assert len(model.control_noise_refiner) == 0 + assert not any(name.startswith("control_noise_refiner.") for name in model.state_dict()) + + +def test_a_scaled_fp8_checkpoint_is_refused_rather_than_loaded_unscaled(monkeypatch, tmp_path) -> None: + payload = quantize_scaled_fp8(torch.randn(HIDDEN, HIDDEN)) + state_dict = _checkpoint() + state_dict[f"{QUANTIZED_LAYER}.weight"] = payload.codes + state_dict[f"{QUANTIZED_LAYER}.weight_scale"] = payload.scale + run, config = _driver(monkeypatch, tmp_path, state_dict) + + with pytest.raises(ValueError, match="quantization side channel"): + run.load(config) + + +def test_an_int8_checkpoint_is_refused_by_its_marker(monkeypatch, tmp_path) -> None: + """int8 is not even a dtype the control could run: the tensor would be assigned where a float parameter belongs, + and the un-rotation the scheme needs has no home here.""" + codes, scale, _ = quantize_convrot(torch.randn(HIDDEN, 256)) + state_dict = _checkpoint() + state_dict[f"{QUANTIZED_LAYER}.weight"] = codes + state_dict[f"{QUANTIZED_LAYER}.weight_scale"] = scale + state_dict[f"{QUANTIZED_LAYER}.comfy_quant"] = comfy_quant_marker({"format": "int8_tensorwise", "convrot": True}) + run, config = _driver(monkeypatch, tmp_path, state_dict) + + with pytest.raises(ValueError, match="quantization side channel"): + run.load(config) + + +def test_a_layout_that_was_never_published_is_refused_by_name(monkeypatch, tmp_path) -> None: + state_dict = {name: tensor for name, tensor in _checkpoint().items() if not name.startswith("control_layers.2.")} + run, config = _driver(monkeypatch, tmp_path, state_dict) + + with pytest.raises(ValueError, match="2 control blocks"): + run.load(config) diff --git a/tests/backend/z_image/test_z_image_control_adapter.py b/tests/backend/z_image/test_z_image_control_adapter.py deleted file mode 100644 index 557f3fbeba4..00000000000 --- a/tests/backend/z_image/test_z_image_control_adapter.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Z-Image control blocks hand back each hint on its own instead of stacking it onto the running state. - -Upstream (VideoX-Fun) returns ``torch.stack([hint_0, ..., hint_k, state])`` from every block and unbinds it again in -the next. Each stack is a fresh allocation larger than any the caching allocator has freed, so a 15-block adapter -reserved 136 hint-sized tensors at 1024px -- 4.5 GiB -- where separate hints need 19. The rewrite must not change a -single hint. -""" - -import torch -from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock - -from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter -from invokeai.backend.z_image.z_image_controlnet_extension import ( - ZImageControlNetExtension, - z_image_forward_with_control, -) - -LATENT_SIDE = 16 - - -def _upstream_hints(adapter: ZImageControlAdapter, c, x, attn_mask, freqs_cis, adaln_input) -> tuple[torch.Tensor, ...]: - """The stacked formulation, as VideoX-Fun's `ZImageControlTransformerBlock.forward` and its caller wrote it.""" - for layer in adapter.control_layers: - if layer.block_id == 0: - c = layer.before_proj(c) + x - all_c = [] - else: - all_c = list(torch.unbind(c)) - c = all_c.pop(-1) - c = ZImageTransformerBlock.forward(layer, c, attn_mask=attn_mask, freqs_cis=freqs_cis, adaln_input=adaln_input) - all_c += [layer.after_proj(c), c] - c = torch.stack(all_c) - return torch.unbind(c)[:-1] - - -@torch.no_grad() -def test_the_control_forward_computes_the_stacked_formulations_hints(tiny_z_image_transformer) -> None: - """Driven through `z_image_forward_with_control`, the path the denoise node takes, with the padding mask it builds - and a caption that needs padding. The first control block's inputs are captured and replayed through the upstream - formulation.""" - model = tiny_z_image_transformer - adapter = ZImageControlAdapter( - num_control_blocks=3, - control_in_dim=16, - dim=model.config.dim, - n_refiner_layers=1, - n_heads=model.config.n_heads, - n_kv_heads=model.config.n_kv_heads, - ).eval() - for param in adapter.parameters(): - torch.nn.init.normal_(param, std=0.1) - extension = ZImageControlNetExtension( - control_adapter=adapter, control_cond=torch.randn(16, 1, LATENT_SIDE, LATENT_SIDE) - ) - - first_block_inputs = {} - adapter.control_layers[0].register_forward_pre_hook( - lambda module, args, kwargs: first_block_inputs.update(c=args[0], **kwargs), with_kwargs=True - ) - computed = [] - compute_hints = extension.compute_hints - - def recording_compute_hints(**kwargs): - computed.append(compute_hints(**kwargs)) - return computed[-1] - - extension.compute_hints = recording_compute_hints - - z_image_forward_with_control( - transformer=model, - x=[torch.randn(16, 1, LATENT_SIDE, LATENT_SIDE)], - t=torch.full((1,), 0.5), - cap_feats=[torch.randn(20, model.config.cap_feat_dim)], - control_extension=extension, - ) - - assert first_block_inputs["attn_mask"] is not None - expected = _upstream_hints(adapter, **first_block_inputs) - (hints,) = computed - assert len(hints) == len(expected) == 3 - for hint, reference in zip(hints, expected, strict=True): - assert torch.equal(hint, reference) diff --git a/tests/backend/z_image/test_z_image_controlnet_extension.py b/tests/backend/z_image/test_z_image_controlnet_extension.py deleted file mode 100644 index 3be00f091c4..00000000000 --- a/tests/backend/z_image/test_z_image_controlnet_extension.py +++ /dev/null @@ -1,24 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import patch - -import torch - -from invokeai.backend.z_image.z_image_controlnet_extension import ZImageControlNetExtension - - -def test_init_logs_control_adapter_diagnostics_without_stdout(capsys): - control_adapter = SimpleNamespace( - control_layers=[ - SimpleNamespace( - after_proj=SimpleNamespace(weight=torch.zeros(1)), - ) - ] - ) - - with patch("invokeai.backend.z_image.z_image_controlnet_extension.logger") as logger_mock: - ZImageControlNetExtension(control_adapter=control_adapter, control_cond=torch.zeros(1)) - - captured = capsys.readouterr() - assert captured.out == "" - assert logger_mock.debug.call_count == 3 - logger_mock.warning.assert_called_once()