# 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, ) -> 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 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 for layer in self.control_layers: c = layer( c, x=unified_hidden_states, attn_mask=attn_mask, freqs_cis=freqs_cis, adaln_input=adaln_input, ) hints = torch.unbind(c)[:-1] return hints