From 3f3462d7ce8b633be8a408356207ffff6f95037d Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 6 May 2025 15:52:05 -0700 Subject: [PATCH] Cleanup Teacache params (#386) --- fastvideo/v1/configs/sample/base.py | 12 ++---- fastvideo/v1/configs/sample/hunyuan.py | 4 +- fastvideo/v1/configs/sample/teacache.py | 40 +++++++++++++++++++ fastvideo/v1/configs/sample/wan.py | 35 +--------------- fastvideo/v1/entrypoints/video_generator.py | 5 --- fastvideo/v1/models/dits/hunyuanvideo.py | 2 +- fastvideo/v1/pipelines/pipeline_batch_info.py | 4 +- 7 files changed, 51 insertions(+), 51 deletions(-) create mode 100644 fastvideo/v1/configs/sample/teacache.py diff --git a/fastvideo/v1/configs/sample/base.py b/fastvideo/v1/configs/sample/base.py index 83a1f5468..fedd5b341 100644 --- a/fastvideo/v1/configs/sample/base.py +++ b/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) diff --git a/fastvideo/v1/configs/sample/hunyuan.py b/fastvideo/v1/configs/sample/hunyuan.py index eb5ca0fba..cc21c8738 100644 --- a/fastvideo/v1/configs/sample/hunyuan.py +++ b/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, diff --git a/fastvideo/v1/configs/sample/teacache.py b/fastvideo/v1/configs/sample/teacache.py new file mode 100644 index 000000000..faa894536 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/fastvideo/v1/configs/sample/wan.py b/fastvideo/v1/configs/sample/wan.py index 12dc17f19..f03da5015 100644 --- a/fastvideo/v1/configs/sample/wan.py +++ b/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, diff --git a/fastvideo/v1/entrypoints/video_generator.py b/fastvideo/v1/entrypoints/video_generator.py index 11551da42..7f0cab131 100644 --- a/fastvideo/v1/entrypoints/video_generator.py +++ b/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 diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index b598eda56..4c99302c9 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/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 diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py index 97bb4d5f5..045a729dc 100644 --- a/fastvideo/v1/pipelines/pipeline_batch_info.py +++ b/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