From 9f1e7c19d283a1ee11749b4fd74e45c9a9bb2b1d Mon Sep 17 00:00:00 2001
From: Shao Duan <46639847+shaoxiongduan@users.noreply.github.com>
Date: Wed, 5 Aug 2026 16:35:09 +0800
Subject: [PATCH] [bugfix] Wan I2V: CLIP image conditioning silently dropped
 when passed as a tensor during training (#1673)

---
 fastvideo/models/dits/causal_wanvideo.py    | 16 ++++++----------
 fastvideo/models/dits/dreamx_world.py       |  8 +++-----
 fastvideo/models/dits/lingbotworld/model.py |  8 +++-----
 fastvideo/models/dits/wanvideo.py           |  8 +++-----
 4 files changed, 15 insertions(+), 25 deletions(-)

diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py
index 88eced7ce..015615668 100644
--- a/fastvideo/models/dits/causal_wanvideo.py
+++ b/fastvideo/models/dits/causal_wanvideo.py
@@ -588,11 +588,9 @@ class CausalWanTransformer3DModel(BaseDiT):
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (encoder_hidden_states_image[0]
+                                           if len(encoder_hidden_states_image) > 0 else None)
 
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
@@ -702,11 +700,9 @@ class CausalWanTransformer3DModel(BaseDiT):
         teacher_forcing = clean_x is not None
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (encoder_hidden_states_image[0]
+                                           if len(encoder_hidden_states_image) > 0 else None)
 
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
diff --git a/fastvideo/models/dits/dreamx_world.py b/fastvideo/models/dits/dreamx_world.py
index 4a24d9550..e80164d6b 100644
--- a/fastvideo/models/dits/dreamx_world.py
+++ b/fastvideo/models/dits/dreamx_world.py
@@ -416,11 +416,9 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel):
         if encoder_hidden_states is not None and not isinstance(
                 encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (encoder_hidden_states_image[0]
+                                           if len(encoder_hidden_states_image) > 0 else None)
 
         batch_size, _, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
diff --git a/fastvideo/models/dits/lingbotworld/model.py b/fastvideo/models/dits/lingbotworld/model.py
index c15014015..56c51e972 100644
--- a/fastvideo/models/dits/lingbotworld/model.py
+++ b/fastvideo/models/dits/lingbotworld/model.py
@@ -311,11 +311,9 @@ class LingBotWorldTransformer3DModel(BaseDiT):
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (encoder_hidden_states_image[0]
+                                           if len(encoder_hidden_states_image) > 0 else None)
 
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py
index 6b925c5f0..6a153c2ee 100644
--- a/fastvideo/models/dits/wanvideo.py
+++ b/fastvideo/models/dits/wanvideo.py
@@ -694,11 +694,9 @@ class WanTransformer3DModel(BaseDiT):
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
+        if isinstance(encoder_hidden_states_image, list):
+            encoder_hidden_states_image = (encoder_hidden_states_image[0]
+                                           if len(encoder_hidden_states_image) > 0 else None)
 
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size