40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
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 |