From 05c009b59f089cd69a93c42c802c5bb65c4475cc Mon Sep 17 00:00:00 2001
From: JerryZhou54 <zhouw.jerry2017@outlook.com>
Date: Mon, 22 Sep 2025 22:47:45 +0000
Subject: [PATCH] Remove extra logger statement

---
 fastvideo/models/dits/causal_wanvideo.py |  9 ++-------
 fastvideo/models/dits/wanvideo.py        | 15 ++++++++-------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py
index e854bf35a..3d05e4b98 100644
--- a/fastvideo/models/dits/causal_wanvideo.py
+++ b/fastvideo/models/dits/causal_wanvideo.py
@@ -147,8 +147,8 @@ class CausalWanSelfAttention(nn.Module):
                 # Assign new keys/values directly up to current_end
                 local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
                 local_start_index = local_end_index - num_new_tokens
-                # kv_cache["k"] = kv_cache["k"].detach()
-                # kv_cache["v"] = kv_cache["v"].detach()
+                kv_cache["k"] = kv_cache["k"].detach()
+                kv_cache["v"] = kv_cache["v"].detach()
                 # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None)
                 kv_cache["k"][:, local_start_index:local_end_index] = roped_key
                 kv_cache["v"][:, local_start_index:local_end_index] = v
@@ -262,14 +262,9 @@ class CausalWanTransformerBlock(nn.Module):
         # *_msa.shape: [batch_size, num_frames, 1, inner_dim]
         # assert shift_msa.dtype == torch.float32
 
-        # logger.info("temb sum: %s, dtype: %s", temb.float().sum().item(), temb.dtype)
-        # logger.info("scale_msa sum: %s, dtype: %s", scale_msa.float().sum().item(), scale_msa.dtype)
-        # logger.info("shift_msa sum: %s, dtype: %s", shift_msa.float().sum().item(), shift_msa.dtype)
-
         # 1. Self-attention
         norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
                         (1 + scale_msa) + shift_msa).flatten(1, 2)
-        # logger.info("norm_hidden_states sum: %s, shape: %s", norm_hidden_states.float().sum().item(), norm_hidden_states.shape)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py
index 3262a9a1e..e403bf500 100644
--- a/fastvideo/models/dits/wanvideo.py
+++ b/fastvideo/models/dits/wanvideo.py
@@ -283,16 +283,16 @@ class WanTransformerBlock(nn.Module):
         # 2. Cross-attention
         if added_kv_proj_dim is not None:
             # I2V
-            self.attn2 = WanI2VCrossAttention(dim, 
-                                              num_heads, 
-                                              qk_norm=qk_norm, 
+            self.attn2 = WanI2VCrossAttention(dim,
+                                              num_heads,
+                                              qk_norm=qk_norm,
                                               eps=eps)
 
         else:
             # T2V
-            self.attn2 = WanT2VCrossAttention(dim, 
-                                              num_heads, 
-                                              qk_norm=qk_norm, 
+            self.attn2 = WanT2VCrossAttention(dim,
+                                              num_heads,
+                                              qk_norm=qk_norm,
                                               eps=eps)
 
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
@@ -835,4 +835,5 @@ class WanTransformer3DModel(CachableDiT):
         if self.is_even:
             return hidden_states + self.previous_residual_even
         else:
-            return hidden_states + self.previous_residual_odd
\ No newline at end of file
+            return hidden_states + self.previous_residual_odd
+            
\ No newline at end of file