Save curr_mask

node.py

@@ -886,7 +886,8 @@ class GetCrossSim:
         cross_sim["max_sim"] = torch.cat(all_max_sim).cpu()
         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
         return (cross_sim,)
 
 class CrossSimMaskExtract:
@@ -969,7 +970,7 @@ class GenConsistent:
             device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
             attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
             fg_bg_mask = None
-            def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
+            def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, _: Tensor) -> Tensor:
                 """
                 Implementation using PyTorch's scaled_dot_product_attention with custom masking.
                 
@@ -1039,7 +1040,7 @@ class GenConsistent:
                                     v = torch.cat([v, saved_value_fg], dim=2)
                 
                 x = F.scaled_dot_product_attention(
-                    q, k, v, attn_mask=mask
+                    q, k, v, attn_mask=None
                 )
                 # Reshape back to the expected output format
                 x = rearrange(x, "B H L D -> B L (H D)")