Add attn masks.

node.py

@@ -427,7 +427,7 @@ class PreviewSubjectMask:
         
         # Stack all fg_masks into a single tensor
         if fg_masks:
-            # Stack the masks to create a tensor with shape [B, H, W]
+            # Stack masks to create a tensor with shape [B, H, W]
             combined_mask = torch.cat(fg_masks, dim=0)
         else:
             # If no fg_masks found, return empty tensor
@@ -887,7 +887,7 @@ class GetCrossSim:
         cross_sim["argmax_indices"] = torch.cat(all_argmax_indices).cpu()
         cross_sim["debug_mask"] = torch.cat(all_cross_sim).cpu()
         cross_sim["id_mask"] = id_mask
-        cross_sim["cur_mask"] = curr_mask
+        cross_sim["curr_mask"] = curr_mask
         return (cross_sim,)
 
 class CrossSimMaskExtract:
@@ -990,15 +990,21 @@ class GenConsistent:
                     sampling_step = sampled_step + 1
                 
                 text_len = q.shape[2] - w*h
+                original_seq_len = q.shape[2]
                 # Apply rotary positional encoding for Flux models
                 if pe is not None:
                     q, k = flux_math.apply_rope(q, k, pe)
+                
+                # Store original sequence length before K,V concatenation
+                original_k_len = k.shape[2]
+                
                 # Store k, v, and x in a dict with layer id as key
                 if sampling_step >= kv_start_step and sampling_step < kv_end_step:
                     # For all other steps, we may apply the key-value expansion algorithm
                     # Get cross-sim data for this layer
                     id_fg_inds = cross_sim.get("id_fg_inds", None)
                     curr_fg_inds = cross_sim.get("curr_fg_inds", None)
+                    curr_mask = cross_sim.get("curr_mask", None)
                     
                     # Process all subjects, not just the last one (index = len(subject) - 1)
                     for subj_idx in range(len(subject)):
@@ -1039,8 +1045,48 @@ class GenConsistent:
                                     k = torch.cat([k, saved_key_fg], dim=2)
                                     v = torch.cat([v, saved_value_fg], dim=2)
                 
+                # Create attention mask for concatenated K
+                attn_mask = None
+                if k.shape[2] > original_k_len and curr_mask is not None:
+                    # We have concatenated keys, need to create attention mask
+                    B, H, L, D = q.shape
+                    extended_k_len = k.shape[2]
+                    
+                    # Initialize attention mask with 1s (allow all attention initially)
+                    attn_mask = torch.ones(B, H, L, extended_k_len, dtype=torch.bool, device=q.device)
+                    
+                    # For concatenated keys (indices >= original_k_len), 
+                    # only allow queries from valid positions in curr_mask to attend
+                    concat_start_idx = original_k_len
+                    concat_end_idx = extended_k_len
+                    
+                    # Process each subject index
+                    for subj_idx in range(len(subject)):
+                        if subj_idx < curr_mask.shape[0]:
+                            subj_mask = curr_mask[subj_idx]  # flattened mask [H*W]
+                            
+                            # Create a mask for vision queries (after text_len) using curr_mask
+                            # Vision queries start at text_len and correspond to spatial positions
+                            vision_query_mask = torch.zeros(L, dtype=torch.bool, device=q.device)
+                            if text_len < L and subj_mask.numel() == (L - text_len):
+                                # Map curr_mask to vision query positions
+                                vision_query_mask[text_len:] = subj_mask
+                            
+                            # Apply the mask: only vision queries with valid positions can attend to concatenated keys
+                            # Text queries (before text_len) cannot attend to concatenated keys
+                            for b in range(B):
+                                # For each head
+                                for h_idx in range(H):
+                                    # Disallow text queries from attending to concatenated keys
+                                    attn_mask[b, h_idx, :text_len, concat_start_idx:] = False
+                                    
+                                    # Allow only valid vision queries to attend to concatenated keys
+                                    if subj_mask.numel() == (L - text_len):
+                                        valid_vision_queries = vision_query_mask[text_len:]
+                                        attn_mask[b, h_idx, text_len:, concat_start_idx:] = valid_vision_queries.unsqueeze(1)
+                
                 x = F.scaled_dot_product_attention(
-                    q, k, v, attn_mask=None
+                    q, k, v, attn_mask=attn_mask
                 )
                 # Reshape back to the expected output format
                 x = rearrange(x, "B H L D -> B L (H D)")