Wan code small cleanup.

comfyanonymous · Feb 27, 2025 · f4dac8a · f4dac8a
1 parent b07f116
commit f4dac8a
Showing 1 changed file with 4 additions and 9 deletions.
diff --git a/comfy/ldm/wan/model.py b/comfy/ldm/wan/model.py
@@ -212,14 +212,10 @@ def forward(
 
         x = x + y * e[2]
 
-        # cross-attention & ffn function
-        def cross_attn_ffn(x, context, e):
-            x = x + self.cross_attn(self.norm3(x), context)
-            y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])
-            x = x + y * e[5]
-            return x
-
-        x = cross_attn_ffn(x, context, e)
+        # cross-attention & ffn
+        x = x + self.cross_attn(self.norm3(x), context)
+        y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])
+        x = x + y * e[5]
         return x
 
 
@@ -442,7 +438,6 @@ def forward_orig(
         # unpatchify
         x = self.unpatchify(x, grid_sizes)
         return x
-        # return [u.float() for u in x]
 
     def forward(self, x, timestep, context, clip_fea=None, **kwargs):
         bs, c, t, h, w = x.shape