Add masking logic.

node.py

@@ -442,6 +442,7 @@ class GenSubject:
         B,L,H,W = latent["samples"].shape
         subject = [{"k": [], "v":[]}]
         def prepare_callback(model, steps, x0_output_dict=None):
+            nonlocal W, H
             sampled_step = -1
             sampling_step = -1
             layer = 0
@@ -449,13 +450,16 @@ class GenSubject:
             w = W // model.model.diffusion_model.patch_size
             h = H // model.model.diffusion_model.patch_size
             def get_fg_attn_weights(text_len, image_start, q, k):
+                nonlocal w, h, fg_condition_mask
                 attn_weights = torch.matmul(q[:, :, image_start:], k[:, :, :text_len].transpose(-2, -1)).to(torch.float16)
                 attn_weights = attn_weights.softmax(-1).sum(1)
                 if fg_condition_mask is not None:
                     fg_condition_mask = rearrange(fg_condition_mask, "s -> 1 1 s")
-                    attn_weights = (attn_weights * fg_condition_mask).mean(1)
-                mask = rearrange(attn_weights, "b (h w) -> b h w", h=h, w=w)
-                return mask
+                    attn_weights = (attn_weights * fg_condition_mask).mean(-1)
+                else:
+                    attn_weights = attn_weights.mean(-1)
+                attn_weights = rearrange(attn_weights, "b (h w) -> b h w", h=h, w=w)
+                return attn_weights
             def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
                 """
                 Implementation using PyTorch's scaled_dot_product_attention with custom masking.
@@ -470,13 +474,15 @@ class GenSubject:
                 Returns:
                     Output tensor of shape [B, L, H*D]
                 """
-                nonlocal sampling_step, sampled_step, layer
+                nonlocal sampling_step, sampled_step, layer, subject
                 if sampling_step != sampled_step + 1:
                     layer = 0
                     sampling_step = sampled_step + 1
+                
                 if sampling_step >= kv_start_step and sampling_step < kv_end_step:
-                    subject[0]["k"] = k
-                    subject[0]["v"] = v
+                    print("save kv at step ", sampling_step, " layer ", layer)
+                    subject[0]["k"].append(k.cpu())
+                    subject[0]["v"].append(v.cpu())
                 # Apply rotary positional encoding for Flux models
                 if pe is not None:
                     q, k = flux_math.apply_rope(q, k, pe)
@@ -488,15 +494,15 @@ class GenSubject:
 
                 
                 if sampling_step == save_mask_step:
-                    subject[0]["front_mask"] = get_fg_attn_weights(text_len, text_len, q, k)
+                    subject[0]["front_mask"] = get_fg_attn_weights(text_len, text_len, q, k).cpu()
 
-                h = F.scaled_dot_product_attention(
+                x = F.scaled_dot_product_attention(
                     q, k, v, attn_mask=mask
                 )
                 # Reshape back to the expected output format
-                h = rearrange(h, "B H L D -> B L (H D)")
+                x = rearrange(x, "B H L D -> B L (H D)")
                 layer += 1
-                return h
+                return x
 
             # Create the attention override function
             def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
@@ -522,6 +528,6 @@ class GenSubject:
                 origin_callback(step, x0, x, total_steps)
             return callback_change_sampled_step
         latent_preview.prepare_callback = prepare_callback
-        return (latent,)
+        return (latent,subject,)