From 12c801c423f4ecb3296addc9afed042b0c3a788a Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 7 Jun 2026 11:42:18 +0200 Subject: [PATCH 1/4] fix(z-image): repair regional guidance forward after diffusers refactor Z-Image Regional Guidance crashed with "split_with_sizes expects split_sizes to sum exactly to 162 ... but got split_sizes=[160]". The regional-prompting patch was a hand-copied snapshot of an outdated ZImageTransformer2DModel.forward. The installed diffusers version changed _pad_with_ids so caption pos_ids are now longer than the caption feature tensor, while the stale patch split RoPE embeddings by feature lengths instead of pos_ids lengths. Rewrite create_regional_forward to delegate to the model's own helpers (patchify_and_embed, _prepare_sequence, _build_unified_sequence) and only override the main-layer attention mask to inject the regional mask. This keeps the patch in sync with upstream diffusers and stops re-implementing the drift-prone patchify/RoPE/padding logic. --- .../z_image/z_image_transformer_patch.py | 225 +++++++----------- 1 file changed, 80 insertions(+), 145 deletions(-) diff --git a/invokeai/backend/z_image/z_image_transformer_patch.py b/invokeai/backend/z_image/z_image_transformer_patch.py index 6e42e678558..fcb61899e8c 100644 --- a/invokeai/backend/z_image/z_image_transformer_patch.py +++ b/invokeai/backend/z_image/z_image_transformer_patch.py @@ -4,7 +4,6 @@ from typing import Callable, List, Optional, Tuple import torch -from torch.nn.utils.rnn import pad_sequence def create_regional_forward( @@ -14,15 +13,21 @@ def create_regional_forward( ) -> Callable: """Create a modified forward function that uses a regional attention mask. - The regional attention mask replaces the internally computed padding mask, - allowing for regional prompting where different image regions attend to - different text prompts. + The regional attention mask replaces the internally computed padding mask on the + main transformer layers (alternating with the plain padding mask), allowing for + regional prompting where different image regions attend to different text prompts. + + This delegates to the model's own helper methods (``patchify_and_embed``, + ``_prepare_sequence``, ``_build_unified_sequence``) so it stays in sync with the + upstream diffusers ``ZImageTransformer2DModel.forward`` implementation. Only the + main-layer attention mask is overridden. Args: - original_forward: The original forward method of ZImageTransformer2DModel. - regional_attn_mask: Attention mask of shape (seq_len, seq_len) where - seq_len = img_seq_len + txt_seq_len. - img_seq_len: Number of image tokens in the sequence. + original_forward: The original forward method of ZImageTransformer2DModel + (kept for signature compatibility; not used directly). + regional_attn_mask: Boolean attention mask of shape (seq_len, seq_len) where + seq_len = img_seq_len + txt_seq_len, ordered [img, txt]. + img_seq_len: Number of (unpadded) image tokens in the sequence. Returns: A modified forward function with regional attention support. @@ -38,160 +43,90 @@ def regional_forward( ) -> Tuple[List[torch.Tensor], dict]: """Modified forward with regional attention mask injection. - This is based on the original ZImageTransformer2DModel.forward but - replaces the padding-based attention mask with a regional attention mask. + Mirrors the basic (non-omni) path of ZImageTransformer2DModel.forward but + injects a regional attention mask into the main transformer layers. """ assert patch_size in self.all_patch_size assert f_patch_size in self.all_f_patch_size - bsz = len(x) device = x[0].device - t_scaled = t * self.t_scale - t_emb = self.t_embedder(t_scaled) - SEQ_MULTI_OF = 32 # From diffusers transformer_z_image.py + # Single adaLN embedding for all tokens (basic mode). + adaln_input = self.t_embedder(t * self.t_scale).type_as(x[0]) - # Patchify and embed (reusing the original method) + # Patchify & embed (basic mode: single image per batch item). ( x, cap_feats, x_size, x_pos_ids, cap_pos_ids, - x_inner_pad_mask, - cap_inner_pad_mask, + x_pad_mask, + cap_pad_mask, ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) - # x embed & refine - 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_cat = torch.cat(x, dim=0) - x_cat = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_cat) - - adaln_input = t_emb.type_as(x_cat) - x_cat[torch.cat(x_inner_pad_mask)] = self.x_pad_token - x_list = list(x_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)) - - x_padded = pad_sequence(x_list, batch_first=True, padding_value=0.0) - x_freqs_cis_padded = 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 - - # Process through noise_refiner - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer in self.noise_refiner: - x_padded = self._gradient_checkpointing_func( - layer, x_padded, x_attn_mask, x_freqs_cis_padded, adaln_input - ) - else: - for layer in self.noise_refiner: - x_padded = layer(x_padded, x_attn_mask, x_freqs_cis_padded, adaln_input) - - # cap embed & refine - 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_cat = torch.cat(cap_feats, dim=0) - cap_cat = self.cap_embedder(cap_cat) - cap_cat[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token - cap_list = list(cap_cat.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_padded = pad_sequence(cap_list, batch_first=True, padding_value=0.0) - cap_freqs_cis_padded = 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 - - # Process through context_refiner - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer in self.context_refiner: - cap_padded = self._gradient_checkpointing_func(layer, cap_padded, cap_attn_mask, cap_freqs_cis_padded) - else: - for layer in self.context_refiner: - cap_padded = layer(cap_padded, cap_attn_mask, cap_freqs_cis_padded) - - # Unified sequence: [img_tokens, txt_tokens] - 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_padded[i][:x_len], cap_freqs_cis_padded[i][:cap_len]])) - - unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens, strict=False)] - assert unified_item_seqlens == [len(_) for _ in unified] - unified_max_item_seqlen = max(unified_item_seqlens) - - unified_padded = pad_sequence(unified, batch_first=True, padding_value=0.0) - unified_freqs_cis_padded = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) + # X embed & refine. + x_seqlens = [len(xi) for xi in x] + x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](torch.cat(x, dim=0)) + x, x_freqs, x_mask, _, _ = self._prepare_sequence( + list(x.split(x_seqlens, dim=0)), x_pos_ids, x_pad_mask, self.x_pad_token, None, device + ) + for layer in self.noise_refiner: + x = layer(x, x_mask, x_freqs, adaln_input, None, None, None) + + # Cap embed & refine. + cap_seqlens = [len(ci) for ci in cap_feats] + cap_feats = self.cap_embedder(torch.cat(cap_feats, dim=0)) + cap_feats, cap_freqs, cap_mask, _, _ = self._prepare_sequence( + list(cap_feats.split(cap_seqlens, dim=0)), cap_pos_ids, cap_pad_mask, self.cap_pad_token, None, device + ) + for layer in self.context_refiner: + cap_feats = layer(cap_feats, cap_mask, cap_freqs) + + # Unified sequence: basic mode order [x, cap]. + unified, unified_freqs, unified_mask, _ = self._build_unified_sequence( + x, + x_freqs, + x_seqlens, + None, + cap_feats, + cap_freqs, + cap_seqlens, + None, + None, + None, + None, + None, + False, # omni_mode + device, + ) + + bsz = unified.shape[0] + unified_seqlen = unified.shape[1] # --- REGIONAL ATTENTION MASK INJECTION --- - # Instead of using the padding mask, we use the regional attention mask - # The regional mask is (seq_len, seq_len), we need to expand it to (batch, seq_len, seq_len) - # and then add the batch dimension for broadcasting: (batch, 1, seq_len, seq_len) - - # Expand regional mask to match the actual sequence length (may include padding) - if regional_attn_mask.shape[0] != unified_max_item_seqlen: - # Pad the regional mask to match unified sequence length - padded_regional_mask = torch.zeros( - (unified_max_item_seqlen, unified_max_item_seqlen), - dtype=regional_attn_mask.dtype, - device=device, - ) - mask_size = min(regional_attn_mask.shape[0], unified_max_item_seqlen) - padded_regional_mask[:mask_size, :mask_size] = regional_attn_mask[:mask_size, :mask_size] - else: - padded_regional_mask = regional_attn_mask.to(device) - - # Convert boolean mask to additive float mask for attention - # True (attend) -> 0.0, False (block) -> -inf - # This is required because the attention backend expects additive masks for 4D inputs - # Use bfloat16 to match the transformer's query dtype - float_mask = torch.zeros_like(padded_regional_mask, dtype=torch.bfloat16) - float_mask[~padded_regional_mask] = float("-inf") - - # Expand to (batch, 1, seq_len, seq_len) for attention - unified_attn_mask = float_mask.unsqueeze(0).unsqueeze(0).expand(bsz, 1, -1, -1) - - # Process through main layers with regional attention mask - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer_idx, layer in enumerate(self.layers): - # Alternate between regional mask and full attention - if layer_idx % 2 == 0: - unified_padded = self._gradient_checkpointing_func( - layer, unified_padded, unified_attn_mask, unified_freqs_cis_padded, adaln_input - ) - else: - # Use padding mask only for odd layers (allows global coherence) - padding_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - padding_mask[i, :seq_len] = 1 - unified_padded = self._gradient_checkpointing_func( - layer, unified_padded, padding_mask, unified_freqs_cis_padded, adaln_input - ) - else: - for layer_idx, layer in enumerate(self.layers): - # Alternate between regional mask and full attention - if layer_idx % 2 == 0: - unified_padded = layer(unified_padded, unified_attn_mask, unified_freqs_cis_padded, adaln_input) - else: - # Use padding mask only for odd layers (allows global coherence) - padding_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - padding_mask[i, :seq_len] = 1 - unified_padded = layer(unified_padded, padding_mask, unified_freqs_cis_padded, adaln_input) - - # Final layer - unified_out = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified_padded, adaln_input) - unified_list = list(unified_out.unbind(dim=0)) - x_out = self.unpatchify(unified_list, x_size, patch_size, f_patch_size) + # The regional mask is (S, S) with S = img_seq_len + txt_seq_len, ordered [img, txt]. + # The unified sequence may be longer due to per-item padding to SEQ_MULTI_OF, so we + # place the regional mask in the top-left block and block (-inf) everything else, which + # also masks out padding tokens. + regional = regional_attn_mask.to(device=device, dtype=torch.bool) + mask_size = min(regional.shape[0], unified_seqlen) + bool_mask = torch.zeros((unified_seqlen, unified_seqlen), dtype=torch.bool, device=device) + bool_mask[:mask_size, :mask_size] = regional[:mask_size, :mask_size] + + # Convert boolean mask to additive float mask: True (attend) -> 0.0, False (block) -> -inf. + float_mask = torch.zeros((unified_seqlen, unified_seqlen), dtype=unified.dtype, device=device) + float_mask[~bool_mask] = float("-inf") + regional_4d_mask = float_mask.unsqueeze(0).unsqueeze(0).expand(bsz, 1, -1, -1) + + # Main transformer layers: alternate regional mask (even) with plain padding mask (odd). + for layer_idx, layer in enumerate(self.layers): + attn_mask = regional_4d_mask if layer_idx % 2 == 0 else unified_mask + unified = layer(unified, attn_mask, unified_freqs, adaln_input, None, None, None) + + # Final layer + unpatchify. + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, c=adaln_input) + x_out = self.unpatchify(list(unified.unbind(dim=0)), x_size, patch_size, f_patch_size) return x_out, {} From 486f2d6201cea6c048f28508cfbe78db2f2ba96e Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 7 Jun 2026 12:23:18 +0200 Subject: [PATCH 2/4] fix(z-image): repair & realign regional guidance after diffusers refactor Z-Image Regional Guidance crashed with "split_with_sizes expects split_sizes to sum exactly to 162 ... but got split_sizes=[160]". The regional-prompting patch was a hand-copied snapshot of an outdated ZImageTransformer2DModel.forward; the installed diffusers version changed _pad_with_ids so caption pos_ids are longer than the caption feature tensor, while the stale patch split RoPE embeddings by feature lengths instead of pos_ids lengths. Rewrite create_regional_forward to delegate to the model's own helpers (patchify_and_embed, _prepare_sequence, _build_unified_sequence) so it stays in sync with upstream diffusers, and only override the main-layer attention mask. Also fix two reasons regional guidance had no visible effect: - Mask alignment: the unified sequence pads the image and caption blocks individually to a multiple of 32, so the real layout is [img_real | img_pad | txt_real | txt_pad]. Scatter the four regional sub-blocks into their padding-aware positions instead of assuming a contiguous top-left block (which only matched square 1024x1024). - CFG pass: the patched forward also runs for the negative prompt; only apply the regional mask to passes whose caption length matches the positive prompt, otherwise fall back to the plain padding mask. --- .../z_image/z_image_transformer_patch.py | 71 +++++++++++++++---- 1 file changed, 58 insertions(+), 13 deletions(-) diff --git a/invokeai/backend/z_image/z_image_transformer_patch.py b/invokeai/backend/z_image/z_image_transformer_patch.py index fcb61899e8c..ac430229503 100644 --- a/invokeai/backend/z_image/z_image_transformer_patch.py +++ b/invokeai/backend/z_image/z_image_transformer_patch.py @@ -105,23 +105,68 @@ def regional_forward( unified_seqlen = unified.shape[1] # --- REGIONAL ATTENTION MASK INJECTION --- - # The regional mask is (S, S) with S = img_seq_len + txt_seq_len, ordered [img, txt]. - # The unified sequence may be longer due to per-item padding to SEQ_MULTI_OF, so we - # place the regional mask in the top-left block and block (-inf) everything else, which - # also masks out padding tokens. + # The regional mask is (S, S) with S = img_seq_len + txt_seq_len, ordered [img, txt], + # using the *unpadded* image and text token counts. In the unified sequence, however, + # both the image block and the caption block are individually padded to a multiple of + # SEQ_MULTI_OF, so the real layout per item is: + # [ img_real | img_pad | txt_real | txt_pad ] + # We therefore scatter the four regional sub-blocks (img-img, img-txt, txt-img, txt-txt) + # into their padding-aware positions instead of assuming a contiguous top-left block. + # + # The patched forward also runs for the negative/CFG pass (a different prompt with a + # different text length). The regional mask was built for the positive prompt only, so + # we apply it only to items whose layout matches the positive prompt and fall back to + # the plain padding mask otherwise. regional = regional_attn_mask.to(device=device, dtype=torch.bool) - mask_size = min(regional.shape[0], unified_seqlen) - bool_mask = torch.zeros((unified_seqlen, unified_seqlen), dtype=torch.bool, device=device) - bool_mask[:mask_size, :mask_size] = regional[:mask_size, :mask_size] - - # Convert boolean mask to additive float mask: True (attend) -> 0.0, False (block) -> -inf. - float_mask = torch.zeros((unified_seqlen, unified_seqlen), dtype=unified.dtype, device=device) - float_mask[~bool_mask] = float("-inf") - regional_4d_mask = float_mask.unsqueeze(0).unsqueeze(0).expand(bsz, 1, -1, -1) + txt_seq_len = regional.shape[0] - img_seq_len + + # Build a per-item additive float mask. Start from the plain padding mask (0 where a + # token is valid, -inf where it is padding) so non-matching items behave normally. + neg_inf = torch.finfo(unified.dtype).min + float_mask = torch.where( + unified_mask.bool().unsqueeze(1).unsqueeze(1), # (bsz, 1, 1, S) + torch.zeros((), dtype=unified.dtype, device=device), + torch.full((), neg_inf, dtype=unified.dtype, device=device), + ).expand(bsz, 1, unified_seqlen, unified_seqlen).clone() + + applied_regional = [False] * bsz + for i in range(bsz): + x_len = x_seqlens[i] + cap_len = cap_seqlens[i] + # The caption block is padded to a multiple of SEQ_MULTI_OF (=32). The positive + # prompt the regional mask was built for has exactly this padded caption length; + # any other pass (e.g. the negative/CFG prompt) is skipped so it runs normally. + SEQ_MULTI_OF = 32 + expected_cap_len = txt_seq_len + ((-txt_seq_len) % SEQ_MULTI_OF) + if ( + txt_seq_len <= 0 + or img_seq_len > x_len + or cap_len != expected_cap_len + or x_len + cap_len > unified_seqlen + ): + continue + applied_regional[i] = True + + ii, it = slice(0, img_seq_len), slice(img_seq_len, img_seq_len + txt_seq_len) + ui = slice(0, img_seq_len) # real image positions in unified item + ut = slice(x_len, x_len + txt_seq_len) # real text positions in unified item + + # Reset the masked region so only regional rules apply to real img/txt tokens; their + # rows start fully blocked and we open the allowed sub-blocks below. + float_mask[i, 0, ui, :] = neg_inf + float_mask[i, 0, ut, :] = neg_inf + + zero = torch.zeros((), dtype=unified.dtype, device=device) + float_mask[i, 0, ui, ui] = torch.where(regional[ii, ii], zero, neg_inf) # img -> img + float_mask[i, 0, ui, ut] = torch.where(regional[ii, it], zero, neg_inf) # img -> txt + float_mask[i, 0, ut, ui] = torch.where(regional[it, ii], zero, neg_inf) # txt -> img + float_mask[i, 0, ut, ut] = torch.where(regional[it, it], zero, neg_inf) # txt -> txt # Main transformer layers: alternate regional mask (even) with plain padding mask (odd). + # If no item matched the positive layout, skip regional injection entirely. + use_regional = any(applied_regional) for layer_idx, layer in enumerate(self.layers): - attn_mask = regional_4d_mask if layer_idx % 2 == 0 else unified_mask + attn_mask = float_mask if (use_regional and layer_idx % 2 == 0) else unified_mask unified = layer(unified, attn_mask, unified_freqs, adaln_input, None, None, None) # Final layer + unpatchify. From 23181700dbf90f685a0d755ddf3932fef97efb1c Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 7 Jun 2026 12:27:32 +0200 Subject: [PATCH 3/4] Chore Ruff + Typegen --- .../backend/z_image/z_image_transformer_patch.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/invokeai/backend/z_image/z_image_transformer_patch.py b/invokeai/backend/z_image/z_image_transformer_patch.py index ac430229503..c5b649b7b33 100644 --- a/invokeai/backend/z_image/z_image_transformer_patch.py +++ b/invokeai/backend/z_image/z_image_transformer_patch.py @@ -123,11 +123,15 @@ def regional_forward( # Build a per-item additive float mask. Start from the plain padding mask (0 where a # token is valid, -inf where it is padding) so non-matching items behave normally. neg_inf = torch.finfo(unified.dtype).min - float_mask = torch.where( - unified_mask.bool().unsqueeze(1).unsqueeze(1), # (bsz, 1, 1, S) - torch.zeros((), dtype=unified.dtype, device=device), - torch.full((), neg_inf, dtype=unified.dtype, device=device), - ).expand(bsz, 1, unified_seqlen, unified_seqlen).clone() + float_mask = ( + torch.where( + unified_mask.bool().unsqueeze(1).unsqueeze(1), # (bsz, 1, 1, S) + torch.zeros((), dtype=unified.dtype, device=device), + torch.full((), neg_inf, dtype=unified.dtype, device=device), + ) + .expand(bsz, 1, unified_seqlen, unified_seqlen) + .clone() + ) applied_regional = [False] * bsz for i in range(bsz): From 9423a2a041319b2dfb03337e6285974bfed6fe05 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 25 Jun 2026 07:51:00 +0200 Subject: [PATCH 4/4] fix(z-image): use identity to gate regional mask onto the positive pass The regional attention patch ran for both the conditioned and negative/CFG forward passes and distinguished them by comparing the padded caption length against the positive prompt's expected length. Two short prompts that round up to the same multiple of 32 collided, so the positive regional mask could be injected into the unconditional prediction and silently corrupt CFG. Discriminate the conditioned pass by tensor identity (cap_feats is the exact positive_cap_feats the mask was built for) instead of a length heuristic, so the positive and negative passes can never be confused. The context manager now requires positive_cap_feats whenever a regional mask is provided, turning the previously inferred invariant into an enforced one rather than a silent no-op. Also build the (bsz, 1, S, S) float mask lazily: compute applied_regional from cheap scalar checks first and skip materializing/cloning the full mask on passes that never match (every negative pass), avoiding a ~33 MB bf16 clone per call. --- invokeai/app/invocations/z_image_denoise.py | 1 + .../z_image/z_image_transformer_patch.py | 122 +++++++++++------- 2 files changed, 73 insertions(+), 50 deletions(-) diff --git a/invokeai/app/invocations/z_image_denoise.py b/invokeai/app/invocations/z_image_denoise.py index c1e864ea179..576d10ac9a1 100644 --- a/invokeai/app/invocations/z_image_denoise.py +++ b/invokeai/app/invocations/z_image_denoise.py @@ -558,6 +558,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: transformer=transformer, regional_attn_mask=regional_extension.regional_attn_mask, img_seq_len=img_seq_len, + positive_cap_feats=pos_prompt_embeds, ) ) diff --git a/invokeai/backend/z_image/z_image_transformer_patch.py b/invokeai/backend/z_image/z_image_transformer_patch.py index c5b649b7b33..a6707fb7f50 100644 --- a/invokeai/backend/z_image/z_image_transformer_patch.py +++ b/invokeai/backend/z_image/z_image_transformer_patch.py @@ -10,6 +10,7 @@ def create_regional_forward( original_forward: Callable, regional_attn_mask: torch.Tensor, img_seq_len: int, + positive_cap_feats: torch.Tensor, ) -> Callable: """Create a modified forward function that uses a regional attention mask. @@ -28,6 +29,13 @@ def create_regional_forward( regional_attn_mask: Boolean attention mask of shape (seq_len, seq_len) where seq_len = img_seq_len + txt_seq_len, ordered [img, txt]. img_seq_len: Number of (unpadded) image tokens in the sequence. + positive_cap_feats: The exact caption-embedding tensor the regional mask was + built for (the conditioned/positive pass). The regional mask is applied only + to forward calls whose ``cap_feats`` is this same object; the negative/CFG + pass supplies a different tensor and is left to run with the plain padding + mask. Identity is used instead of a token-length heuristic so the positive + and negative passes can never be confused even when their padded lengths + coincide. Returns: A modified forward function with regional attention support. @@ -51,6 +59,13 @@ def regional_forward( device = x[0].device + # Identify which caption inputs belong to the conditioned (positive) pass the regional + # mask was built for. Capture this before patchify_and_embed reassigns ``cap_feats``. + # The negative/CFG pass supplies a different tensor, so object identity distinguishes the + # passes regardless of token length (avoids the positive mask leaking into the uncond + # prediction when prompt lengths happen to pad to the same multiple). + is_positive_pass = [ci is positive_cap_feats for ci in cap_feats] + # Single adaLN embedding for all tokens (basic mode). adaln_input = self.t_embedder(t * self.t_scale).type_as(x[0]) @@ -113,62 +128,62 @@ def regional_forward( # We therefore scatter the four regional sub-blocks (img-img, img-txt, txt-img, txt-txt) # into their padding-aware positions instead of assuming a contiguous top-left block. # - # The patched forward also runs for the negative/CFG pass (a different prompt with a - # different text length). The regional mask was built for the positive prompt only, so - # we apply it only to items whose layout matches the positive prompt and fall back to - # the plain padding mask otherwise. + # The patched forward also runs for the negative/CFG pass (a different prompt). The + # regional mask was built for the positive prompt only, so we apply it only to the + # conditioned items and fall back to the plain padding mask otherwise. regional = regional_attn_mask.to(device=device, dtype=torch.bool) txt_seq_len = regional.shape[0] - img_seq_len - # Build a per-item additive float mask. Start from the plain padding mask (0 where a - # token is valid, -inf where it is padding) so non-matching items behave normally. - neg_inf = torch.finfo(unified.dtype).min - float_mask = ( - torch.where( - unified_mask.bool().unsqueeze(1).unsqueeze(1), # (bsz, 1, 1, S) - torch.zeros((), dtype=unified.dtype, device=device), - torch.full((), neg_inf, dtype=unified.dtype, device=device), - ) - .expand(bsz, 1, unified_seqlen, unified_seqlen) - .clone() - ) + # Decide per item whether the regional mask applies, using only cheap scalar checks, so + # that on passes that never match (e.g. every negative/CFG pass) we avoid materializing + # the (bsz, 1, S, S) float mask at all. + applied_regional = [ + is_positive_pass[i] + and txt_seq_len > 0 + and img_seq_len <= x_seqlens[i] + and x_seqlens[i] + cap_seqlens[i] <= unified_seqlen + for i in range(bsz) + ] - applied_regional = [False] * bsz - for i in range(bsz): - x_len = x_seqlens[i] - cap_len = cap_seqlens[i] - # The caption block is padded to a multiple of SEQ_MULTI_OF (=32). The positive - # prompt the regional mask was built for has exactly this padded caption length; - # any other pass (e.g. the negative/CFG prompt) is skipped so it runs normally. - SEQ_MULTI_OF = 32 - expected_cap_len = txt_seq_len + ((-txt_seq_len) % SEQ_MULTI_OF) - if ( - txt_seq_len <= 0 - or img_seq_len > x_len - or cap_len != expected_cap_len - or x_len + cap_len > unified_seqlen - ): - continue - applied_regional[i] = True - - ii, it = slice(0, img_seq_len), slice(img_seq_len, img_seq_len + txt_seq_len) - ui = slice(0, img_seq_len) # real image positions in unified item - ut = slice(x_len, x_len + txt_seq_len) # real text positions in unified item - - # Reset the masked region so only regional rules apply to real img/txt tokens; their - # rows start fully blocked and we open the allowed sub-blocks below. - float_mask[i, 0, ui, :] = neg_inf - float_mask[i, 0, ut, :] = neg_inf + # Main transformer layers: alternate regional mask (even) with plain padding mask (odd). + # If no item matched the positive pass, skip regional injection entirely. + use_regional = any(applied_regional) + float_mask = None + if use_regional: + # Build a per-item additive float mask. Start from the plain padding mask (0 where a + # token is valid, -inf where it is padding) so non-matching items behave normally. + neg_inf = torch.finfo(unified.dtype).min zero = torch.zeros((), dtype=unified.dtype, device=device) - float_mask[i, 0, ui, ui] = torch.where(regional[ii, ii], zero, neg_inf) # img -> img - float_mask[i, 0, ui, ut] = torch.where(regional[ii, it], zero, neg_inf) # img -> txt - float_mask[i, 0, ut, ui] = torch.where(regional[it, ii], zero, neg_inf) # txt -> img - float_mask[i, 0, ut, ut] = torch.where(regional[it, it], zero, neg_inf) # txt -> txt + float_mask = ( + torch.where( + unified_mask.bool().unsqueeze(1).unsqueeze(1), # (bsz, 1, 1, S) + zero, + torch.full((), neg_inf, dtype=unified.dtype, device=device), + ) + .expand(bsz, 1, unified_seqlen, unified_seqlen) + .clone() + ) + + for i in range(bsz): + if not applied_regional[i]: + continue + x_len = x_seqlens[i] + + ii, it = slice(0, img_seq_len), slice(img_seq_len, img_seq_len + txt_seq_len) + ui = slice(0, img_seq_len) # real image positions in unified item + ut = slice(x_len, x_len + txt_seq_len) # real text positions in unified item + + # Reset the masked region so only regional rules apply to real img/txt tokens; + # their rows start fully blocked and we open the allowed sub-blocks below. + float_mask[i, 0, ui, :] = neg_inf + float_mask[i, 0, ut, :] = neg_inf + + float_mask[i, 0, ui, ui] = torch.where(regional[ii, ii], zero, neg_inf) # img -> img + float_mask[i, 0, ui, ut] = torch.where(regional[ii, it], zero, neg_inf) # img -> txt + float_mask[i, 0, ut, ui] = torch.where(regional[it, ii], zero, neg_inf) # txt -> img + float_mask[i, 0, ut, ut] = torch.where(regional[it, it], zero, neg_inf) # txt -> txt - # Main transformer layers: alternate regional mask (even) with plain padding mask (odd). - # If no item matched the positive layout, skip regional injection entirely. - use_regional = any(applied_regional) for layer_idx, layer in enumerate(self.layers): attn_mask = float_mask if (use_regional and layer_idx % 2 == 0) else unified_mask unified = layer(unified, attn_mask, unified_freqs, adaln_input, None, None, None) @@ -187,6 +202,7 @@ def patch_transformer_for_regional_prompting( transformer, regional_attn_mask: Optional[torch.Tensor], img_seq_len: int, + positive_cap_feats: Optional[torch.Tensor] = None, ): """Context manager to temporarily patch the transformer for regional prompting. @@ -195,6 +211,9 @@ def patch_transformer_for_regional_prompting( regional_attn_mask: Regional attention mask of shape (seq_len, seq_len). If None, the transformer is not patched. img_seq_len: Number of image tokens. + positive_cap_feats: The caption-embedding tensor the regional mask was built for. + Required when ``regional_attn_mask`` is provided; the mask is applied only to + forward calls whose ``cap_feats`` is this exact object (the conditioned pass). Yields: The (possibly patched) transformer. @@ -204,11 +223,14 @@ def patch_transformer_for_regional_prompting( yield transformer return + if positive_cap_feats is None: + raise ValueError("positive_cap_feats is required when regional_attn_mask is provided") + # Store original forward original_forward = transformer.forward # Create and bind the regional forward - regional_fwd = create_regional_forward(original_forward, regional_attn_mask, img_seq_len) + regional_fwd = create_regional_forward(original_forward, regional_attn_mask, img_seq_len, positive_cap_feats) transformer.forward = lambda *args, **kwargs: regional_fwd(transformer, *args, **kwargs) try: