From 1d0516a2a97dd99d2f4ebc5eed2b0c8e5b48edf7 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 4 Nov 2025 10:26:37 +0200
Subject: [PATCH] Avoid graph break for LongCat

---
 wanvideo/modules/model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index e089733..3358fc3 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -672,7 +672,7 @@ class WanT2VCrossAttention(WanSelfAttention):
 
         if is_longcat:
             if num_cond_latents is not None and num_cond_latents > 0:
-                num_cond_latents_thw = num_cond_latents * (s // grid_sizes[0][0])
+                num_cond_latents_thw = num_cond_latents * (s // num_latent_frames)
                 x = x[:, num_cond_latents_thw:]
             q = self.norm_q(self.q(x).view(b, -1, n, d))
         else:
@@ -1185,7 +1185,7 @@ class WanAttentionBlock(nn.Module):
             full_v = torch.cat([v, v_ip], dim=1)
             y = self.self_attn.forward(q, full_k, full_v, seq_lens)
         elif is_longcat and num_cond_latents is not None and num_cond_latents > 0:
-            num_cond_latents_thw = num_cond_latents * (N // grid_sizes[0][0])
+            num_cond_latents_thw = num_cond_latents * (N // num_latent_frames)
             # process the condition tokens
             x_cond = self.self_attn.forward(
                 q[:, :num_cond_latents_thw].contiguous(),