Update models.py

2025-05-09 15:05:40 +03:00
parent 9b321d09c5
commit 27e541b0a5
1 changed files with 6 additions and 0 deletions
@@ -1166,6 +1166,12 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
                    self.teacache_skipped_steps_uncond += 1
                else:
                    self.teacache_skipped_steps_cond += 1
+                
+                # Verify tensor dimensions match before adding
+                if img.shape == previous_residual.shape:
+                    img = img + previous_residual.to(img.device)
+                else:
+                    should_calc = True # Force recalculation if dimensions don't match

            if should_calc:
                ori_img = img.clone()