modified: src/flux/modules/layers.py

src/flux/modules/layers.py

@@ -159,6 +159,11 @@ class DoubleStreamBlock(nn.Module):
 
     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)
 
@@ -170,21 +175,24 @@ 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)
         v = torch.cat((txt_v, img_v), dim=2)
         # import pdb;pdb.set_trace()
-        print(f"q_{cur_step}:{q}")
-        print(f"k_{cur_step}:{k}")
-        print(f"v_{cur_step}:{v}")
+        
 
         # if using adaptive attention guidance during samping
         if not info['inverse'] and 'attn_guidance' in info['editing_strategy']: