modified: src/flux/math.py

src/flux/math.py

@@ -111,6 +111,17 @@ def scaled_dot_product_attention(query, key, value, txt_shape, img_shape, cur_st
     attn_weight = torch.softmax(attn_weight, dim=-1)
     attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
 
+    if info['inverse']:
+        print(f"[DEBUG_Inverse] query:\n{query}")
+        print(f"[DEBUG_Inverse] key :\n{key}")
+        print(f"[DEBUG_Inverse] value :\n{value}")
+        print(f"[DEBUG_Inverse] attn_weight:\n{attn_weight}")
+    else:
+        print(f"[DEBUG_Sample] query:\n{query}")
+        print(f"[DEBUG_Sample] key :\n{key}")
+        print(f"[DEBUG_Sample] value :\n{value}")
+        print(f"[DEBUG_Sample] attn_weight values:\n{attn_weight}")
+
     if not info['inverse']:
         # GENERATE MASK
         txt_img_cross = attn_weight[:, :, -img_shape:, :txt_shape]  # lower left part
@@ -130,10 +141,11 @@ def scaled_dot_product_attention(query, key, value, txt_shape, img_shape, cur_st
         save_image(mask_img.unsqueeze(0), save_path)
         # Debug information
         print(f"[DEBUG] cur_step: {cur_step}, cur_block: {cur_block}")
-        print(f"[DEBUG] norm_heatmap values:\n{norm_heatmap}")
-        print(f"[DEBUG] mask_img (before thresholding) stats: min={mask_img.min().item()}, max={mask_img.max().item()}, mean={mask_img.mean().item()}")
-        print(f"[DEBUG] thre value: {thre}")
-        print(f"[DEBUG] mask_img (before thresholding) values:\n{mask_img}")
+
+        #print(f"[DEBUG] norm_heatmap values:\n{norm_heatmap}")
+        #print(f"[DEBUG] mask_img (before thresholding) stats: min={mask_img.min().item()}, max={mask_img.max().item()}, mean={mask_img.mean().item()}")
+        #print(f"[DEBUG] thre value: {thre}")
+        #print(f"[DEBUG] mask_img (before thresholding) values:\n{mask_img}")
 
         mask_img[mask_img >= thre] = 1
         mask_img[mask_img < thre] = 0