# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. ##### Enjoy this spagheti VRAM optimizations done by DeepBeepMeep ! # I am sure you are a nice person and as you copy this code, you will give me officially proper credits: # Please link to https://github.com/deepbeepmeep/Wan2GP and @deepbeepmeep on twitter import math from typing import List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence from diffusers.models.normalization import RMSNorm from shared.attention import pay_attention class ModuleWrapper: def __init__(self, module): self.module = module def __call__(self): return self.module def apply_rotary_emb_inplace(x_list: list, freqs_cis: torch.Tensor) -> torch.Tensor: x_in = x_list.pop() dtype = x_in.dtype x = x_in.float().reshape(*x_in.shape[:-1], -1, 2) # [B, S, H, D//2, 2] x_in = None # freqs_cis shape: [B, S, rope_dim, 2] -> cos/sin: [B, S, 1, rope_dim] to broadcast over heads cos = freqs_cis[..., 0].unsqueeze(2) # [B, S, 1, rope_dim] sin = freqs_cis[..., 1].unsqueeze(2) # [B, S, 1, rope_dim] x0, x1 = x[..., 0], x[..., 1] # [B, S, H, D//2] x0_orig = x0.clone() x0.mul_(cos).addcmul_(x1, sin, value=-1) x1.mul_(cos).addcmul_(x0_orig, sin) return x.flatten(3).to(dtype) ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 class TimestepEmbedder(nn.Module): def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): super().__init__() if mid_size is None: mid_size = out_size self.mlp = nn.Sequential( nn.Linear( frequency_embedding_size, mid_size, bias=True, ), nn.SiLU(), nn.Linear( mid_size, out_size, bias=True, ), ) self.frequency_embedding_size = frequency_embedding_size @staticmethod def timestep_embedding(t, dim, max_period=10000): with torch.amp.autocast("cuda", enabled=False): half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half ) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) return embedding def forward(self, t): t_freq = self.timestep_embedding(t, self.frequency_embedding_size) weight_dtype = self.mlp[0].weight.dtype if weight_dtype.is_floating_point: t_freq = t_freq.to(weight_dtype) t_emb = self.mlp(t_freq) return t_emb class FeedForward(nn.Module): def __init__(self, dim: int, hidden_dim: int): super().__init__() self.w1 = nn.Linear(dim, hidden_dim, bias=False) self.w2 = nn.Linear(hidden_dim, dim, bias=False) self.w3 = nn.Linear(dim, hidden_dim, bias=False) def _forward_silu_gating(self, x1, x3): return F.silu(x1) * x3 def forward(self, x): return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x))) class Attention(nn.Module): """Simple attention module to match weight key structure.""" def __init__(self, dim: int, n_heads: int, qk_norm: bool): super().__init__() self.n_heads = n_heads self.head_dim = dim // n_heads self.to_q = nn.Linear(dim, dim, bias=False) self.to_k = nn.Linear(dim, dim, bias=False) self.to_v = nn.Linear(dim, dim, bias=False) # Sequential to match weight keys: to_out.0.weight self.to_out = nn.Sequential(nn.Linear(dim, dim, bias=False)) self.norm_q = RMSNorm(self.head_dim, eps=1e-5) if qk_norm else None self.norm_k = RMSNorm(self.head_dim, eps=1e-5) if qk_norm else None def forward(self, h_list: list, freqs_cis: torch.Tensor, NAG= None) -> torch.Tensor: """Compute self-attention with RoPE. h_list is cleared to free memory early.""" h = h_list.pop() query = self.to_q(h) key = self.to_k(h) value = self.to_v(h); del h # Reshape to [batch, seq, heads, head_dim] query = query.unflatten(-1, (self.n_heads, -1)) key = key.unflatten(-1, (self.n_heads, -1)) value = value.unflatten(-1, (self.n_heads, -1)) # Apply QK normalization if self.norm_q is not None: query = self.norm_q(query) if self.norm_k is not None: key = self.norm_k(key) if freqs_cis is not None: q_list = [query]; del query query = apply_rotary_emb_inplace(q_list, freqs_cis) k_list = [key]; del key key = apply_rotary_emb_inplace(k_list, freqs_cis) dtype = query.dtype # NAG for joint-attention models: transformer duplicates batch into [pos, neg] halves. if NAG is not None: nag_scale = NAG["scale"] nag_alpha = NAG["alpha"] nag_tau = NAG["tau"] cap_embed_len = NAG["cap_embed_len"] x_list = [query[:, :-cap_embed_len], key[:, :-cap_embed_len], value[:, :-cap_embed_len] ] x_pos = pay_attention(x_list).flatten(2, 3).to(dtype) query[:, -2 *cap_embed_len:-cap_embed_len] = query[:, -cap_embed_len:] key[:, -2 *cap_embed_len:-cap_embed_len] = key[:, -cap_embed_len:] value[:, -2 *cap_embed_len:-cap_embed_len] = value[:, -cap_embed_len:] x_list = [query[:, :-cap_embed_len], key[:, :-cap_embed_len], value[:, :-cap_embed_len] ] del query, key, value x_neg = pay_attention(x_list).flatten(2, 3).to(dtype) x_neg_tail = x_neg[:, -cap_embed_len:].clone() x_guidance = x_neg x_guidance.mul_(1 - nag_scale) x_guidance.add_(x_pos, alpha=nag_scale) norm_positive = torch.norm(x_pos, p=1, dim=-1, keepdim=True) norm_guidance = torch.norm(x_guidance, p=1, dim=-1, keepdim=True) scale = norm_guidance / norm_positive scale = torch.nan_to_num(scale, 10) factor = (1 / (norm_guidance + 1e-7) * norm_positive * nag_tau).to(x_guidance.dtype) x_guidance = torch.where(scale > nag_tau, x_guidance * factor, x_guidance).to(dtype) del norm_positive, norm_guidance, scale, factor x_guidance.mul_(nag_alpha) x_guidance.add_(x_pos, alpha=(1 - nag_alpha)) x_pos = None out = torch.cat([x_guidance, x_neg_tail], dim=1) x_pos = x_neg = x_guidance = None else: x_list = [query, key, value] del query, key, value out = pay_attention(x_list).flatten(2, 3).to(dtype); return self.to_out(out) class ZImageTransformerBlock(nn.Module): def __init__( self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int, # kept for API compatibility norm_eps: float, qk_norm: bool, modulation=True, ): super().__init__() self.dim = dim self.n_heads = n_heads self.head_dim = dim // n_heads # Attention module (named 'attention' to match weight keys) self.attention = Attention(dim, n_heads, qk_norm) # SwiGLU FFN: hidden_dim = dim * 8/3 ≈ 2.67x expansion self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) self.ffn_mult = 8 / 3 # For chunking: matches SwiGLU expansion factor self.layer_id = layer_id self.attention_norm1 = RMSNorm(dim, eps=norm_eps) self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) self.attention_norm2 = RMSNorm(dim, eps=norm_eps) self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) self.modulation = modulation if modulation: self.adaLN_modulation = nn.Sequential( nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True), ) def _apply_ffn_chunked(self, ffn_in: torch.Tensor) -> None: _, seq_len, dim = ffn_in.shape ffn_in_flat = ffn_in.reshape(-1, dim) chunk_size = max(int(seq_len // self.ffn_mult), 1) for ffn_chunk in torch.split(ffn_in_flat, chunk_size): ffn_chunk[...] = self.feed_forward(ffn_chunk) def forward( self, x: torch.Tensor, attn_mask: torch.Tensor, freqs_cis: torch.Tensor, adaln_input: Optional[torch.Tensor] = None, NAG = None, ): if self.modulation: scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2) # In-place modulation for attention block scale_msa.add_(1.0) normed = self.attention_norm1(x) normed.mul_(scale_msa) attn_out = self.attention_norm2(self.attention([normed], freqs_cis, NAG=NAG)) attn_out.mul_(gate_msa.tanh_()) x.add_(attn_out); attn_out = None # In-place modulation for FFN block (chunked) scale_mlp.add_(1.0) normed = self.ffn_norm1(x) normed.mul_(scale_mlp) self._apply_ffn_chunked(normed) ffn_out = self.ffn_norm2(normed); normed = None ffn_out.mul_(gate_mlp.tanh_()) x.add_(ffn_out); ffn_out = None else: x.add_(self.attention_norm2(self.attention([self.attention_norm1(x)], freqs_cis, NAG=NAG))) normed = self.ffn_norm1(x) self._apply_ffn_chunked(normed) x.add_(self.ffn_norm2(normed)); normed = None return x class ZImageControlTransformerBlock(ZImageTransformerBlock): """Control block that processes control signals and produces skip connections.""" def __init__( self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int, norm_eps: float, qk_norm: bool, modulation=True, block_id=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(self.dim, self.dim) nn.init.zeros_(self.before_proj.weight) nn.init.zeros_(self.before_proj.bias) self.after_proj = nn.Linear(self.dim, self.dim) nn.init.zeros_(self.after_proj.weight) nn.init.zeros_(self.after_proj.bias) def forward(self, hints, x, **kwargs): # behold dbm magic ! c = hints[0] hints[0] = None if self.block_id == 0: c = self.before_proj(c) bz = x.shape[0] if bz > c.shape[0]: c = c.repeat(bz, 1, 1 ) c += x c = super().forward(c, **kwargs) c_skip = self.after_proj(c) hints[0] = c return c_skip class BaseZImageTransformerBlock(ZImageTransformerBlock): """Base block that can optionally apply control hints.""" def __init__( self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int, norm_eps: float, qk_norm: bool, modulation=True, block_id=None ): 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, hints=None, hints_kwargs=None, context_scale=1.0, **kwargs): hints_processed = None if self.block_id is not None and hints is not None: hints_processed = self.control()(hints, **hints_kwargs) hidden_states = super().forward(hidden_states, **kwargs) if hints_processed is not None: hidden_states[:, :hints_processed.shape[1]].add_(hints_processed, alpha=context_scale) return hidden_states class FinalLayer(nn.Module): def __init__(self, hidden_size, out_channels): super().__init__() self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.linear = nn.Linear(hidden_size, out_channels, bias=True) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), ) def forward(self, x, c): scale = self.adaLN_modulation(c) scale.add_(1.0) x = self.norm_final(x) x.mul_(scale.unsqueeze(1)) return self.linear(x) class RopeEmbedder: def __init__( self, theta: float = 256.0, axes_dims: List[int] = (16, 56, 56), axes_lens: List[int] = (64, 128, 128), ): self.theta = theta self.axes_dims = axes_dims self.axes_lens = axes_lens assert len(axes_dims) == len(axes_lens), "axes_dims and axes_lens must have the same length" self.freqs_cis = None @staticmethod def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): """Precompute cos/sin frequencies for RoPE (real arithmetic, no complex).""" with torch.device("cpu"): freqs_cis = [] for d, e in zip(dim, end): freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64) / d)) timestep = torch.arange(e, dtype=torch.float64) freqs = torch.outer(timestep, freqs).float() # Stack cos/sin in last dim: [end, dim//2, 2] freqs_cis.append(torch.stack([freqs.cos(), freqs.sin()], dim=-1)) return freqs_cis def __call__(self, ids: torch.Tensor): assert ids.ndim == 2 assert ids.shape[-1] == len(self.axes_dims) device = ids.device if self.freqs_cis is None: self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta) self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] else: # Ensure freqs_cis are on the same device as ids if self.freqs_cis[0].device != device: self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] result = [] for i in range(len(self.axes_dims)): index = ids[:, i] result.append(self.freqs_cis[i][index]) # Cat on dim=-2 (D//2 dimension) since format is [S, D//2, 2] return torch.cat(result, dim=-2) class ZImageTransformer2DModel(nn.Module): def preprocess_loras(self, model_type, sd): first = next(iter(sd), None) if first is None: return sd if ".default." not in first and ".lora." not in first: return sd new_sd = {} for k, v in sd.items(): if ".default." in k: k = k.replace(".default.", ".") if ".lora." in k: k = k.replace(".lora.", ".lora_") new_sd[k] = v return new_sd def __init__( self, # Control-specific parameters (optional) control_layers_places=None, control_refiner_layers_places=None, control_in_dim=None, add_control_noise_refiner=False, enable_control=False, use_separate_control_refiner=False, # Base model parameters all_patch_size=(2,), all_f_patch_size=(1,), in_channels=16, dim=3840, n_layers=30, n_refiner_layers=2, n_heads=30, n_kv_heads=30, norm_eps=1e-5, qk_norm=True, cap_feat_dim=2560, siglip_feat_dim=None, rope_theta=256.0, t_scale=1000.0, axes_dims=[32, 48, 48], axes_lens=[1024, 512, 512], ) -> None: super().__init__() self.in_channels = in_channels self.out_channels = in_channels self.all_patch_size = all_patch_size self.all_f_patch_size = all_f_patch_size self.dim = dim self.n_heads = n_heads self.rope_theta = rope_theta self.t_scale = t_scale # Control-specific attributes self.control_in_dim = control_in_dim if control_in_dim is not None else in_channels self.control_layers_places = [i for i in range(0, n_layers, 2)] if control_layers_places is None else control_layers_places self.control_refiner_layers_places = [i for i in range(0, n_refiner_layers)] if control_refiner_layers_places is None else control_refiner_layers_places self.add_control_noise_refiner = add_control_noise_refiner self.enable_control = enable_control self.use_separate_control_refiner = use_separate_control_refiner # Track whether the refiner uses its own control blocks (v2.1 fix) self._control_noise_uses_dedicated_layers = False assert 0 in self.control_layers_places self.control_layers_mapping = {i: n for n, i in enumerate(self.control_layers_places)} self.control_refiner_layers_mapping = {i: n for n, i in enumerate(self.control_refiner_layers_places)} assert len(all_patch_size) == len(all_f_patch_size) all_x_embedder = {} all_final_layer = {} for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)): x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True) all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels) all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer self.all_x_embedder = nn.ModuleDict(all_x_embedder) self.all_final_layer = nn.ModuleDict(all_final_layer) # Noise refiner - use control version if enable_control and add_control_noise_refiner if enable_control and add_control_noise_refiner: self.noise_refiner = nn.ModuleList( [ BaseZImageTransformerBlock( 1000 + layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=True, block_id=self.control_refiner_layers_mapping[layer_id] if layer_id in self.control_refiner_layers_places else None ) for layer_id in range(n_refiner_layers) ] ) else: self.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) ] ) self.context_refiner = nn.ModuleList( [ ZImageTransformerBlock( layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=False, ) for layer_id in range(n_refiner_layers) ] ) self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024) self.cap_embedder = nn.Sequential( RMSNorm(cap_feat_dim, eps=norm_eps), nn.Linear(cap_feat_dim, dim, bias=True), ) self.x_pad_token = nn.Parameter(torch.empty((1, dim))) self.cap_pad_token = nn.Parameter(torch.empty((1, dim))) # Main layers - use control version if enable_control if enable_control: 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) ] ) else: self.layers = nn.ModuleList( [ ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm) for layer_id in range(n_layers) ] ) head_dim = dim // n_heads assert head_dim == sum(axes_dims) self.axes_dims = axes_dims self.axes_lens = axes_lens self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens) # Control-specific layers (only created when enable_control=True) if enable_control: # Control blocks self.control_layers = nn.ModuleList( [ ZImageControlTransformerBlock( i, dim, n_heads, n_kv_heads, norm_eps, qk_norm, block_id=i ) for i in self.control_layers_places ] ) # Control patch embeddings control_all_x_embedder = {} for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)): x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * self.control_in_dim, dim, bias=True) control_all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder self.control_all_x_embedder = nn.ModuleDict(control_all_x_embedder) # Control noise refiner (for v2 control) if add_control_noise_refiner: self.control_noise_refiner = nn.ModuleList( [ ZImageControlTransformerBlock( 1000 + layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=True, block_id=layer_id ) for layer_id in range(n_refiner_layers) ] ) else: # For v1 control 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) ] ) self.adapt_control_model() def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size) -> List[torch.Tensor]: pH = pW = patch_size pF = f_patch_size bsz = len(x) assert len(size) == bsz x_out_list = [] for i in range(bsz): F, H, W = size[i] ori_len = (F // pF) * (H // pH) * (W // pW) # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)" x_out_list.append( x[i][:ori_len] .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels) .permute(6, 0, 3, 1, 4, 2, 5) .reshape(self.out_channels, F, H, W) ) return torch.stack(x_out_list) @staticmethod def create_coordinate_grid(size, start=None, device=None): if start is None: start = (0 for _ in size) axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)] grids = torch.meshgrid(axes, indexing="ij") return torch.stack(grids, dim=-1) def patchify( self, all_image: List[torch.Tensor], patch_size: int, f_patch_size: int, cap_padding_len: int, ): 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 = [] for i, image in enumerate(all_image): ### Process 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 image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) # "c f pf h ph w pw -> (f h w) (pf ph pw c)" 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 image_ori_pos_ids = self.create_coordinate_grid( size=(F_tokens, H_tokens, W_tokens), start=(cap_padding_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) # pad 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 patchify_and_embed( self, all_image: List[torch.Tensor], all_cap_feats: List[torch.Tensor], patch_size: int, f_patch_size: int, ): 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 = [] all_cap_pos_ids = [] all_cap_pad_mask = [] all_cap_feats_out = [] for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)): ### Process Caption cap_ori_len = len(cap_feat) cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF # padded position ids cap_padded_pos_ids = self.create_coordinate_grid( size=(cap_ori_len + cap_padding_len, 1, 1), start=(1, 0, 0), device=device, ).flatten(0, 2) all_cap_pos_ids.append(cap_padded_pos_ids) # pad mask all_cap_pad_mask.append( torch.cat( [ torch.zeros((cap_ori_len,), dtype=torch.bool, device=device), torch.ones((cap_padding_len,), dtype=torch.bool, device=device), ], dim=0, ) ) # padded feature cap_padded_feat = torch.cat( [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], dim=0, ) all_cap_feats_out.append(cap_padded_feat) ### Process 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 image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) # "c f pf h ph w pw -> (f h w) (pf ph pw c)" 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 image_ori_pos_ids = self.create_coordinate_grid( size=(F_tokens, H_tokens, W_tokens), start=(cap_ori_len + cap_padding_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) # pad 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_cap_feats_out, all_image_size, all_image_pos_ids, all_cap_pos_ids, all_image_pad_mask, all_cap_pad_mask, ) @property def has_control(self) -> bool: """Returns True if the model has control layers enabled.""" return self.enable_control def adapt_control_model(self): """Move control blocks to be submodules of their corresponding main layers.""" if not self.enable_control or not hasattr(self, 'control_layers'): return # Assume we will fall back to legacy behavior unless we wire dedicated control refiner blocks. self._control_noise_uses_dedicated_layers = False modules_dict = {k: m for k, m in self.named_modules()} for model_layer, control_idx in self.control_layers_mapping.items(): control_module = modules_dict[f"control_layers.{control_idx}"] target = modules_dict[f"layers.{model_layer}"] setattr(target, "control", ModuleWrapper(control_module)) for model_layer, control_idx in self.control_refiner_layers_mapping.items(): control_module = None if ( self.add_control_noise_refiner and self.use_separate_control_refiner and hasattr(self, "control_noise_refiner") ): noise_key = f"control_noise_refiner.{control_idx}" control_module = modules_dict.get(noise_key, None) if control_module is not None: self._control_noise_uses_dedicated_layers = True if control_module is None: control_module = modules_dict.get(f"control_layers.{control_idx}") if control_module is None: continue target = modules_dict[f"noise_refiner.{model_layer}"] setattr(target, "control", ModuleWrapper(control_module)) def prepare_forward_control_1_0( self, x, cap_feats, control_context, kwargs, t=None, patch_size=2, f_patch_size=1, ): """Control v1.0 processing (without noise refiner control).""" # embeddings bsz = len(control_context) device = control_context[0].device ( control_context, x_size, x_pos_ids, x_inner_pad_mask, ) = self.patchify(control_context, patch_size, f_patch_size, cap_feats[0].size(0)) # control_context embed & refine x_item_seqlens = [len(_) for _ in control_context] assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) x_max_item_seqlen = max(x_item_seqlens) control_context = torch.cat(control_context, dim=0) control_context = self.control_all_x_embedder[f"{patch_size}-{f_patch_size}"](control_context) # Match t_embedder output dtype to control_context for layerwise casting compatibility adaln_input = t.type_as(control_context) control_context[torch.cat(x_inner_pad_mask)] = self.x_pad_token control_context = list(control_context.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 = pad_sequence(control_context, 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 for layer in self.control_noise_refiner: control_context = layer(control_context, x_attn_mask, x_freqs_cis, adaln_input) # unified cap_item_seqlens = [len(_) for _ in cap_feats] 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[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 new_kwargs = dict(x=x) new_kwargs.update(kwargs) return new_kwargs, c def prepare_forward_control_2_0_refiner( self, x, cap_feats, control_context, kwargs, t=None, patch_size=2, f_patch_size=1, ): """Control v2.0 refiner processing.""" # embeddings bsz = len(control_context) device = control_context[0].device ( control_context, control_context_size, control_context_pos_ids, control_context_inner_pad_mask, ) = self.patchify(control_context, patch_size, f_patch_size, cap_feats[0].size(0)) # control_context embed & refine control_context_item_seqlens = [len(_) for _ in control_context] assert all(_ % SEQ_MULTI_OF == 0 for _ in control_context_item_seqlens) control_context_max_item_seqlen = max(control_context_item_seqlens) control_context = torch.cat(control_context, dim=0) control_context = self.control_all_x_embedder[f"{patch_size}-{f_patch_size}"](control_context) # Match t_embedder output dtype to control_context for layerwise casting compatibility adaln_input = t.type_as(control_context) control_context[torch.cat(control_context_inner_pad_mask)] = self.x_pad_token control_context = list(control_context.split(control_context_item_seqlens, dim=0)) control_context_freqs_cis = list(self.rope_embedder(torch.cat(control_context_pos_ids, dim=0)).split(control_context_item_seqlens, dim=0)) control_context = pad_sequence(control_context, batch_first=True, padding_value=0.0) control_context_freqs_cis = pad_sequence(control_context_freqs_cis, batch_first=True, padding_value=0.0) control_context_attn_mask = torch.zeros((bsz, control_context_max_item_seqlen), dtype=torch.bool, device=device) for i, seq_len in enumerate(control_context_item_seqlens): control_context_attn_mask[i, :seq_len] = 1 c = control_context new_kwargs = dict( x=x, attn_mask=control_context_attn_mask, freqs_cis=control_context_freqs_cis, adaln_input=adaln_input, ) new_kwargs.update(kwargs) return new_kwargs, c, control_context_item_seqlens def prepare_forward_control_2_0_layers( self, x, cap_feats, control_context, control_context_item_seqlens, kwargs, ): """Control v2.0 layers processing.""" control_context_len = control_context_item_seqlens cap_len = cap_feats.shape[1] control_context_unified= torch.cat([control_context[:control_context_len], cap_feats[:cap_len]], dim=1) c = pad_sequence(control_context_unified, batch_first=True, padding_value=0.0) new_kwargs = dict(x=x) new_kwargs.update(kwargs) return new_kwargs, c def forward( self, x_list: List[torch.Tensor], t, cap_feats_list: List[torch.Tensor], patch_size=2, f_patch_size=1, control_context_list=None, control_context_scale=1.0, target_timestep=None, callback=None, pipeline=None, NAG =None, ): """Forward pass with list-based processing (outer loop over layers, inner loop over samples).""" assert patch_size in self.all_patch_size assert f_patch_size in self.all_f_patch_size num_noise_samples = len(x_list) assert len(cap_feats_list) == num_noise_samples, "cap_feats_list must match x_list length" device = x_list[0].device t_high = t.to(dtype=torch.float32) t_emb = self.t_embedder(t_high.abs() * self.t_scale) if target_timestep is not None: target_t_high = target_timestep.to(dtype=torch.float32) delta_t = t_high - target_t_high delta_t_abs = delta_t.abs() t_emb_2 = self.t_embedder((target_t_high - t_high) * self.t_scale) t_emb = t_emb + t_emb_2 * delta_t_abs.unsqueeze(1) t = t_emb # Patchify and embed each (x, cap_feats) pair x_embedder = self.all_x_embedder[f"{patch_size}-{f_patch_size}"] embedded_x_list = [] cap_embedded_list = [] cap_ori_len_ref = None cap_pad_mask_ref = None per_sample_kwargs = [] for i,(x,cap_feats) in enumerate(zip(x_list, cap_feats_list)): bsz = x.shape[0] ( x_patches, cap_out, x_i_size, x_pos_ids, cap_pos_ids, x_inner_pad_mask, cap_inner_pad_mask) = self.patchify_and_embed(x, [cap_feats]* bsz, patch_size, f_patch_size) # Store x_size from first sample (all should be same shape) if i==0: x_size = x_i_size cap_ori_len_ref = len(cap_feats) cap_pad_mask_ref = cap_inner_pad_mask[0] x_seqlen = len(x_patches[0]) cap_seqlen = len(cap_out[0]) x_freqs_cis = self.rope_embedder(torch.cat(x_pos_ids[:1], dim=0)).unsqueeze(0) cap_freqs_cis = self.rope_embedder(torch.cat(cap_pos_ids[:1], dim=0)).unsqueeze(0) # Embed x # x_patches = x_embedded = x_embedder(torch.stack(x_patches)) x_embedded[torch.stack(x_inner_pad_mask)] = self.x_pad_token embedded_x_list.append(x_embedded) # Embed cap_feats cap_embedded = self.cap_embedder(torch.stack(cap_out)) cap_embedded[torch.stack(cap_inner_pad_mask)] = self.cap_pad_token cap_embedded_list.append(cap_embedded) per_sample_kwargs.append( { "x_seqlen": x_seqlen, "cap_seqlen": cap_seqlen, "x_freqs": x_freqs_cis, "cap_freqs": cap_freqs_cis, "x_attn_mask": torch.ones((1, x_seqlen), dtype=torch.bool, device=device), "cap_attn_mask": torch.ones((1, cap_seqlen), dtype=torch.bool, device=device), } ) x = cap_feats = None # Match t_embedder output dtype adaln_input = t.type_as(embedded_x_list[0])[:1] # Same timestep for all samples # Control processing - compute hints with batch dimension [num_samples, seq, dim] refiner_hints_tuple_list = [] # List of tensors with batch dim, one per layer ctrl_ctx_tensor_list = [] # control context for v2.0 layers ctrl_seqlens_list = [] refiner_hints_kwargs_list= [] any_control = control_context_list is not None and self.enable_control control_V2 = any_control and self.add_control_noise_refiner if control_V2: # Control v2.0 refiner - process all samples together (no duplication) # Stack all embedded_x into batch: [num_samples, seq, dim] for i,(x_batch,cap_embedded_ref, control_ctx_input) in enumerate(zip(embedded_x_list, cap_embedded_list, control_context_list)): kwargs = dict( attn_mask=per_sample_kwargs[i]["x_attn_mask"], freqs_cis=per_sample_kwargs[i]["x_freqs"], adaln_input=adaln_input, ) # Pass all samples at once - hints will have shape [num_samples, seq, dim] each #ctrl_ctx_tensor, refiner_hints_kwargs, refiner_hints_tuple, ctrl_seqlens = self.prepare_forward_control_2_0_refiner( x_batch, [cap_embedded_ref] , [control_ctx_input.squeeze(0)] , kwargs = kwargs, t=adaln_input, patch_size=patch_size, f_patch_size=f_patch_size ) refiner_hints_kwargs_list.append(refiner_hints_kwargs) refiner_hints_tuple_list.append([refiner_hints_tuple]) # ctrl_ctx_tensor_list.append(ctrl_ctx_tensor) ctrl_seqlens_list.append(ctrl_seqlens[0]) x_batch = cap_embedded_ref = control_ctx_input = None # Noise refiner for layer in self.noise_refiner: for i, x_i in enumerate(embedded_x_list): kwargs = dict( attn_mask=per_sample_kwargs[i]["x_attn_mask"], freqs_cis=per_sample_kwargs[i]["x_freqs"], adaln_input=adaln_input, ) if control_V2: # v2 control # kwargs["hints"] = refiner_hints_tuple_list[i] kwargs["hints"] = refiner_hints_tuple_list[i] kwargs["hints_kwargs"] = refiner_hints_kwargs_list[i] kwargs["context_scale"] = control_context_scale embedded_x_list[i] = layer(x_i, **kwargs) x_i = None if control_V2: # finish processing v2 control hints control_layer_offset = 0 if getattr(self, "_control_noise_uses_dedicated_layers", False) else len(self.control_refiner_layers_places) control_layers_seq = self.control_layers[control_layer_offset:] for hints, hints_kwargs in zip(refiner_hints_tuple_list, refiner_hints_kwargs_list): for layer in control_layers_seq: layer(hints, **hints_kwargs) ctrl_ctx_tensor_list.append(hints[0]) hints = hints_kwargs = None # Context refiner # NAG: prepare negative caption embedding once (it is static w.r.t. timestep). NAG_index = -1 nag_enabled = NAG is not None if nag_enabled: nag_kwargs = per_sample_kwargs[0] neg_feats = NAG["neg_feats"] NAG_index = 0 if len(neg_feats) < cap_ori_len_ref: pad_len = cap_ori_len_ref - len(neg_feats) neg_feats = torch.cat([neg_feats, neg_feats[-1:].repeat(pad_len, 1)], dim=0) elif len(neg_feats) > cap_ori_len_ref: neg_feats = neg_feats[:cap_ori_len_ref] pad_len = nag_kwargs["cap_seqlen"] - len(neg_feats) if pad_len > 0: neg_feats = torch.cat([neg_feats, neg_feats[-1:].repeat(pad_len, 1)], dim=0) neg_cap_embedded = self.cap_embedder(neg_feats.unsqueeze(0)).type_as(cap_embedded_list[0]) neg_cap_embedded[cap_pad_mask_ref.unsqueeze(0)] = self.cap_pad_token for layer in self.context_refiner: neg_cap_embedded = layer( neg_cap_embedded, attn_mask=nag_kwargs["cap_attn_mask"], freqs_cis=nag_kwargs["cap_freqs"], ) NAG["cap_embed_len"] = neg_cap_embedded.shape[1] neg_feats = None for layer in self.context_refiner: for i, cap_i in enumerate(cap_embedded_list): cap_embedded_list[i] = layer( cap_i, attn_mask=per_sample_kwargs[i]["cap_attn_mask"], freqs_cis=per_sample_kwargs[i]["cap_freqs"], ) cap_i = None # Create unified (x + cap) for each sample unified_list = [] unified_freqs_list = [] unified_attn_masks = [] control_seqlens = [] for i,(embedded_x, cap_embedded) in enumerate(zip(embedded_x_list, cap_embedded_list)): is_nag = nag_enabled and i == NAG_index unified_list.append( torch.cat( [embedded_x, cap_embedded] + ([neg_cap_embedded.expand(len(cap_embedded), -1, -1)] if is_nag else []), dim=1, ) ) embedded_x_list[i] = None neg_cap_embedded = None for i in range(num_noise_samples): is_nag = nag_enabled and i == NAG_index x_freqs_i = per_sample_kwargs[i]["x_freqs"] cap_freqs_i = per_sample_kwargs[i]["cap_freqs"] unified_freqs_i = torch.cat([x_freqs_i, cap_freqs_i] + ([cap_freqs_i] if is_nag else []), dim=1) unified_freqs_list.append(unified_freqs_i) control_seqlen = per_sample_kwargs[i]["x_seqlen"] + per_sample_kwargs[i]["cap_seqlen"] control_seqlens.append(control_seqlen) unified_seqlen = control_seqlen + (per_sample_kwargs[i]["cap_seqlen"] if is_nag else 0) unified_attn_masks.append(torch.ones((1, unified_seqlen), dtype=torch.bool, device=device)) hints_list = [] hints_kwargs_list = [] # Compute control hints for main layers if any_control: cap_embedded_ref = cap_embedded_list[0] adaln_input_expanded = adaln_input.expand(num_noise_samples, -1) if control_V2: # Control v2.0 for i, (unified_batch, ctrl_ctx_tensor, ctrl_seqlens) in enumerate(zip(unified_list, ctrl_ctx_tensor_list, ctrl_seqlens_list)): control_kwargs = dict( attn_mask=unified_attn_masks[i][:, :control_seqlens[i]], freqs_cis=unified_freqs_list[i][:, :control_seqlens[i]], adaln_input=adaln_input_expanded, ) hints_kwargs, hints_tuple = self.prepare_forward_control_2_0_layers( unified_batch[:, :control_seqlens[i]], cap_embedded_ref, ctrl_ctx_tensor, ctrl_seqlens, control_kwargs ) hints_kwargs_list.append(hints_kwargs) hints_list.append([hints_tuple]) unified_batch = ctrl_ctx_tensor = ctrl_seqlens = None else: # Control v1.0 for i, (unified_batch, ctrl_ctx_tensor) in enumerate(zip(unified_list, control_context_list)): control_kwargs = dict( attn_mask=unified_attn_masks[i][:, :control_seqlens[i]], freqs_cis=unified_freqs_list[i][:, :control_seqlens[i]], adaln_input=adaln_input_expanded, ) hints_kwargs, hints_tuple = self.prepare_forward_control_1_0( unified_batch[:, :control_seqlens[i]], cap_embedded_ref, ctrl_ctx_tensor, control_kwargs, t=adaln_input, patch_size=patch_size, f_patch_size=f_patch_size, ) hints_kwargs_list.append(hints_kwargs) hints_list.append([hints_tuple]) unified_batch = ctrl_ctx_tensor = None if len(hints_list) == 0: hints_list = [None] * len(x_list) hints_kwargs_list = [None] * len(x_list) # Main layers for layer in self.layers: if callback is not None: callback(-1, None, False, True) if pipeline is not None and getattr(pipeline, "_interrupt", False): return None for i, (unified_i, hints, hints_kwargs) in enumerate(zip(unified_list, hints_list, hints_kwargs_list)): kwargs = {} if hints is not None: kwargs.update(dict(hints= hints, hints_kwargs = hints_kwargs, context_scale = control_context_scale)) if NAG_index == i: kwargs["NAG"]= NAG unified_list[i] = layer( unified_i, attn_mask=unified_attn_masks[i], freqs_cis=unified_freqs_list[i], adaln_input=adaln_input, **kwargs, ) unified_i = hints = kwargs = hints_kwargs = None # Final layer and unpatchify output_list = [] final_layer = self.all_final_layer[f"{patch_size}-{f_patch_size}"] for i in range(num_noise_samples): final_out = final_layer(unified_list[i], adaln_input) final_out = final_out[:, :per_sample_kwargs[i]["x_seqlen"]] unpatchified = self.unpatchify(final_out, x_size, patch_size, f_patch_size) output_list.append(unpatchified) final_out = None return output_list