Fix Teacache (#38)

* support teacache

* fix teacache

---------

Co-authored-by: Steven <[email protected]>
Co-authored-by: Howe <[email protected]>

Author: 3 people

skyreels_v2_infer/modules/transformer.py

@@ -759,7 +759,10 @@ def forward(self, x, t, context, clip_fea=None, y=None, fps=None):
                     for block in self.blocks:
                         x = block(x, **kwargs)
                     self.previous_residual_odd = x - ori_x
-        
+
+            self.cnt += 1
+            if self.cnt >= self.num_steps:
+                self.cnt = 0
         else:
             for block in self.blocks:
                 x = block(x, **kwargs)
@@ -768,9 +771,7 @@ def forward(self, x, t, context, clip_fea=None, y=None, fps=None):
 
         # unpatchify
         x = self.unpatchify(x, grid_sizes)
-        self.cnt += 1
-        if self.cnt >= self.num_steps:
-            self.cnt = 0
+
         return x.float()
 
     def unpatchify(self, x, grid_sizes):