modified: src/flux/math.py

src/flux/math.py

@@ -128,10 +128,18 @@ def scaled_dot_product_attention(query, key, value, txt_shape, img_shape, cur_st
         save_path = f'heatmap/step_{cur_step}_layer_{cur_block}_token{token_index}.png'
         load_path = [f'heatmap/step_{cur_step-1}_layer_{i}_token{token_index}.png' for i in layer]        #save_image(mask_img.unsqueeze(0), save_path)
         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}")
 
         mask_img[mask_img >= thre] = 1
         mask_img[mask_img < thre] = 0
         #save_image(mask_img.unsqueeze(0), save_path)
+        print(f"[DEBUG] mask_img (after thresholding) unique values: {mask_img.unique().tolist()}")
+        print(f"[DEBUG] mask_img (after thresholding) values:\n{mask_img}")
 
         mask_tensor = torch.zeros_like(mask_img)  # Set mask_tensor as a zero tensor
         if cur_step > 3:

src/flux/model.py

@@ -93,8 +93,8 @@ class Flux(nn.Module):
         # --- CRITICAL DEBUG: Check the device of self.img_in's parameters ---
         weight_device = self.img_in.weight.device
         bias_device = self.img_in.bias.device if self.img_in.bias is not None else "N/A (None)"
-        print(f"self.img_in.weight device: {weight_device}")
-        print(f"self.img_in.bias device: {bias_device}")
+        #print(f"self.img_in.weight device: {weight_device}")
+        #print(f"self.img_in.bias device: {bias_device}")
 
         # --- FIX: Explicitly move img to the device of img_in's weight if they differ ---
         # This is the core fix if the mismatch is here
@@ -109,7 +109,7 @@ class Flux(nn.Module):
             # Add similar checks/moves for txt_ids, y, timesteps, guidance if needed,
             # but based on your previous debug, they were on cuda:0.
         
-        print("--- End of Critical Debug ---")
+        #print("--- End of Critical Debug ---")
 
         # running on sequences img
         img = self.img_in(img)