modified: src/flux/model.py

src/flux/model.py

@@ -90,6 +90,10 @@ 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)"
@@ -114,6 +118,10 @@ class Flux(nn.Module):
         # running on sequences img
         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}")
+
         if self.params.guidance_embed:
             if guidance is None:
                 raise ValueError("Didn't get guidance strength for guidance distilled model.")

src/flux/modules/layers.py

@@ -158,11 +158,6 @@ class DoubleStreamBlock(nn.Module):
         self.cur_block = cur_block
 
     def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, cur_step: int, info) -> tuple[Tensor, Tensor]:
-        
-        print(f"img_{cur_step}:{img}")
-        print(f"txt_{cur_step}:{txt}")
-        print(f"vec_{cur_step}:{vec}")
-        print(f"pe_{cur_step}:{pe}")
 
         img_mod1, img_mod2 = self.img_mod(vec)
         txt_mod1, txt_mod2 = self.txt_mod(vec)
@@ -175,18 +170,12 @@ class DoubleStreamBlock(nn.Module):
 
         img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
 
-        print(f"img_modulated_{cur_step}:{img_modulated}")
-        print(f"img_qkv_{cur_step}:{img_qkv}")
-        print(f"img_q_{cur_step}:{img_q}")
-        print(f"img_k_{cur_step}:{img_k}")
         # prepare txt for attention
         txt_modulated = self.txt_norm1(txt)
         txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
         txt_qkv = self.txt_attn.qkv(txt_modulated)
         txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
         txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
-        print(f"txt_q_{cur_step}:{txt_q}")
-        print(f"txt_k_{cur_step}:{txt_k}")
         # run actual attention
         q = torch.cat((txt_q, img_q), dim=2) #[8, 24, 512, 128] + [8, 24, 900, 128] -> [8, 24, 1412, 128]
         k = torch.cat((txt_k, img_k), dim=2)