From ca0974d5f0a122c220b6b2bafb80017262d9444e Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 29 Apr 2025 01:46:39 +0300
Subject: [PATCH] cleanup

---
 wanvideo/modules/model.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index 453f298..15151f3 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -341,7 +341,7 @@ class WanT2VCrossAttention(WanSelfAttention):
 
         # output
         x = x.flatten(2)
-        
+
         # FantasyTalking audio attention
         if audio_proj is not None:
             if len(audio_proj.shape) == 4:
@@ -388,25 +388,24 @@ class WanI2VCrossAttention(WanSelfAttention):
             context(Tensor): Shape [B, L2, C]
             context_lens(Tensor): Shape [B]
         """
-        #context_img = context[:, :clip_embed.shape[1]]
-        #context = context[:, clip_embed.shape[1]:]
         b, n, d = x.size(0), self.num_heads, self.head_dim
 
         # compute query, key, value
         q = self.norm_q(self.q(x)).view(b, -1, n, d)
         k = self.norm_k(self.k(context)).view(b, -1, n, d)
         v = self.v(context).view(b, -1, n, d)
+           
+        # text attention
+        x = attention(q, k, v, k_lens=context_lens, attention_mode=self.attention_mode)
+        x = x.flatten(2)
+
+        #img attention
         if clip_embed is not None:
             k_img = self.norm_k_img(self.k_img(clip_embed)).view(b, -1, n, d)
             v_img = self.v_img(clip_embed).view(b, -1, n, d)
             img_x = attention(q, k_img, v_img, k_lens=None, attention_mode=self.attention_mode)
-        # compute attention
-        x = attention(q, k, v, k_lens=context_lens, attention_mode=self.attention_mode)
-
-        # output
-        x = x.flatten(2)
-        if clip_embed is not None:
             img_x = img_x.flatten(2)
+
             x = x + img_x
 
         # FantasyTalking audio attention