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