modified: src/flux/model.py

src/flux/math.py

@@ -111,12 +111,6 @@ 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)
 
-    '''
-    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
@@ -135,7 +129,7 @@ def scaled_dot_product_attention(query, key, value, txt_shape, img_shape, cur_st
         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] 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()}")

src/flux/model.py

@@ -90,17 +90,13 @@ class Flux(nn.Module):
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
 
-        print(f"img_{cur_step}:{img}")
-        print(f"img_ids_{cur_step}:{img_ids}")
-        print(f"txt_{cur_step}:{txt}")
-        print(f"txt_ids_{cur_step}:{txt_ids}")
         # --- 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("Model img_in weight sample:", self.img_in.weight[0, :10]) # Print first 10 elements of first row
-        print("Model img_in bias sample:", self.img_in.bias[:10]) # Print first 10 elements o
+        #print("Model img_in weight sample:", self.img_in.weight[0, :10]) # Print first 10 elements of first row
+        #print("Model img_in bias sample:", self.img_in.bias[:10]) # Print first 10 elements o
         # --- 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
         if img.device != weight_device:
@@ -120,8 +116,8 @@ class Flux(nn.Module):
         img = self.img_in(img)
         vec = self.time_in(timestep_embedding(timesteps, 256))
 
-        print(f"self.img_in(img)_{cur_step}:{img}")
-        print(f"self.time_in(timestep_embedding(timesteps, 256))_{cur_step}:{vec}")
+        #print(f"self.img_in(img)_{cur_step}:{img}")
+        #print(f"self.time_in(timestep_embedding(timesteps, 256))_{cur_step}:{vec}")
 
         if self.params.guidance_embed:
             if guidance is None: