Cleanup Teacache params (#386)

fastvideo/v1/configs/sample/base.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Union
 
 from fastvideo.v1.logger import init_logger
@@ -35,6 +35,9 @@ class SamplingParam:
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
 
+    # TeaCache parameters
+    enable_teacache: bool = False
+
     # Misc
     save_video: bool = True
     return_frames: bool = False
@@ -75,10 +78,3 @@ class SamplingParam:
 @dataclass
 class CacheParams:
     cache_type: str = "none"
-
-
-@dataclass
-class TeaCacheParams(CacheParams):
-    cache_type: str = "teacache"
-    teacache_thresh: float = 0.0
-    coefficients: List[float] = field(default_factory=list)

fastvideo/v1/configs/sample/hunyuan.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 
-from fastvideo.v1.configs.sample.base import SamplingParam, TeaCacheParams
+from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.v1.configs.sample.teacache import TeaCacheParams
 
 
 @dataclass
@@ -14,7 +15,6 @@ class HunyuanSamplingParam(SamplingParam):
 
     guidance_scale: float = 1.0
 
-    enable_teacache: bool = False
     teacache_params: TeaCacheParams = field(
         default_factory=lambda: TeaCacheParams(
             teacache_thresh=0.15,

fastvideo/v1/configs/sample/teacache.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.sample.base import CacheParams
+
+
+@dataclass
+class TeaCacheParams(CacheParams):
+    cache_type: str = "teacache"
+    teacache_thresh: float = 0.0
+    coefficients: list[float] = field(default_factory=list)
+
+
+@dataclass
+class WanTeaCacheParams(CacheParams):
+    # Unfortunately, TeaCache is very different for Wan than other models
+    cache_type: str = "teacache"
+    teacache_thresh: float = 0.0
+    use_ret_steps: bool = True
+    ret_steps_coeffs: list[float] = field(default_factory=list)
+    non_ret_steps_coeffs: list[float] = field(default_factory=list)
+
+    @property
+    def coefficients(self) -> list[float]:
+        if self.use_ret_steps:
+            return self.ret_steps_coeffs
+        else:
+            return self.non_ret_steps_coeffs
+
+    @property
+    def ret_steps(self) -> int:
+        if self.use_ret_steps:
+            return 5 * 2
+        else:
+            return 1 * 2
+
+    def get_cutoff_steps(self, num_inference_steps: int) -> int:
+        if self.use_ret_steps:
+            return num_inference_steps * 2
+        else:
+            return num_inference_steps * 2 - 2

fastvideo/v1/configs/sample/wan.py

@@ -1,36 +1,7 @@
 from dataclasses import dataclass, field
 
-from fastvideo.v1.configs.sample.base import CacheParams, SamplingParam
-
-
-@dataclass
-class WanTeaCacheParams(CacheParams):
-    # Unfortunately, TeaCache is very different for Wan than other models
-    cache_type: str = "teacache"
-    teacache_thresh: float = 0.0
-    use_ret_steps: bool = True
-    ret_steps_coeffs: list[float] = field(default_factory=list)
-    non_ret_steps_coeffs: list[float] = field(default_factory=list)
-
-    @property
-    def coefficients(self) -> list[float]:
-        if self.use_ret_steps:
-            return self.ret_steps_coeffs
-        else:
-            return self.non_ret_steps_coeffs
-
-    @property
-    def ret_steps(self) -> int:
-        if self.use_ret_steps:
-            return 5 * 2
-        else:
-            return 1 * 2
-
-    def get_cutoff_steps(self, num_inference_steps: int) -> int:
-        if self.use_ret_steps:
-            return num_inference_steps * 2
-        else:
-            return num_inference_steps * 2 - 2
+from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.v1.configs.sample.teacache import WanTeaCacheParams
 
 
 @dataclass
@@ -46,7 +17,6 @@ class WanT2V_1_3B_SamplingParam(SamplingParam):
     negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
     num_inference_steps: int = 50
 
-    enable_teacache: bool = False
     teacache_params: WanTeaCacheParams = field(
         default_factory=lambda: WanTeaCacheParams(
             teacache_thresh=0.08,
@@ -73,7 +43,6 @@ class WanT2V_14B_SamplingParam(SamplingParam):
     negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
     num_inference_steps: int = 50
 
-    enable_teacache: bool = False
     teacache_params: WanTeaCacheParams = field(
         default_factory=lambda: WanTeaCacheParams(
             teacache_thresh=0.20,

fastvideo/v1/entrypoints/video_generator.py

@@ -99,11 +99,6 @@ class VideoGenerator:
             **config_args)
         fastvideo_args.check_fastvideo_args()
 
-        # Temporary hacks for tea cache
-        if fastvideo_args.cache_strategy == "teacache":
-            assert fastvideo_args.dit_config.cache_config is not None
-            fastvideo_args.dit_config.cache_config.enable_teacache = True
-
         return cls.from_fastvideo_args(fastvideo_args)
 
     @classmethod

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 
 from fastvideo.v1.attention import DistributedAttention, LocalAttention
 from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
-from fastvideo.v1.configs.sample.base import TeaCacheParams
+from fastvideo.v1.configs.sample.teacache import TeaCacheParams
 from fastvideo.v1.distributed.parallel_state import (
     get_sequence_model_parallel_world_size)
 from fastvideo.v1.forward_context import get_forward_context

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -12,8 +12,8 @@ from typing import Any, Dict, List, Optional, Union
 
 import torch
 
-from fastvideo.v1.configs.sample.base import TeaCacheParams
-from fastvideo.v1.configs.sample.wan import WanTeaCacheParams
+from fastvideo.v1.configs.sample.teacache import (TeaCacheParams,
+                                                  WanTeaCacheParams)
 
 
 @dataclass