diff --git a/comfyui/examples/FastHunyuan-diffusers.json b/comfyui/examples/FastHunyuan-diffusers.json index db4a71d13..645d96b89 100644 --- a/comfyui/examples/FastHunyuan-diffusers.json +++ b/comfyui/examples/FastHunyuan-diffusers.json @@ -443,7 +443,6 @@ 1025, "fixed", 24, - -99999, -99999 ], "auto_widget_states": { @@ -491,11 +490,6 @@ "isAuto": true, "value": -99999, "cachedValue": "X://insert/path/here.mp4" - }, - "enable_teacache": { - "isAuto": true, - "value": -99999, - "cachedValue": true } } }, @@ -642,4 +636,4 @@ "VHS_KeepIntermediate": true }, "version": 0.4 -} \ No newline at end of file +} diff --git a/comfyui/examples/Wan2.1-I2V-14B-480P-Diffusers.json b/comfyui/examples/Wan2.1-I2V-14B-480P-Diffusers.json index 1cb01d7e9..61c15d4a3 100644 --- a/comfyui/examples/Wan2.1-I2V-14B-480P-Diffusers.json +++ b/comfyui/examples/Wan2.1-I2V-14B-480P-Diffusers.json @@ -353,8 +353,7 @@ 1024, "fixed", 24, - "X://insert/path/here.mp4", - true + "X://insert/path/here.mp4" ], "auto_widget_states": { "height": { @@ -401,11 +400,6 @@ "isAuto": true, "value": "X://insert/path/here.mp4", "cachedValue": "X://insert/path/here.mp4" - }, - "enable_teacache": { - "isAuto": true, - "value": true, - "cachedValue": true } } }, @@ -694,4 +688,4 @@ "VHS_KeepIntermediate": true }, "version": 0.4 -} \ No newline at end of file +} diff --git a/comfyui/video_generator/inference_args.py b/comfyui/video_generator/inference_args.py index 2ba46157e..1cf1e4d3c 100644 --- a/comfyui/video_generator/inference_args.py +++ b/comfyui/video_generator/inference_args.py @@ -31,9 +31,6 @@ class InferenceArgs: "image_path": ("STRING", { "default": "X://insert/path/here.mp4" }), - "enable_teacache": ([True, False], { - "default": False - }), } } @@ -57,7 +54,6 @@ class InferenceArgs: seed, fps, image_path, - enable_teacache, ): raw_args = { "height": height, @@ -69,7 +65,6 @@ class InferenceArgs: "seed": seed, "fps": fps, "image_path": image_path, - "enable_teacache": enable_teacache, } # Filter out keys where value is -99999, handling different types properly diff --git a/comfyui/web/js/FastVideo.core.js b/comfyui/web/js/FastVideo.core.js index 87f4009d7..53afb22ad 100644 --- a/comfyui/web/js/FastVideo.core.js +++ b/comfyui/web/js/FastVideo.core.js @@ -552,7 +552,7 @@ app.registerExtension({ ] const floatWidgetNames = ["embedded_cfg_scale", "guidance_scale"] const comboWidgetNames = ["vae_tiling", "vae_precision", "vae_sp", "text_encoder_precision", "precision", - "load_encoder", "load_decoder", "use_tiling", "use_temporal_tiling", "use_parallel_tiling", "dit_cpu_offload", "enable_teacache" + "load_encoder", "load_decoder", "use_tiling", "use_temporal_tiling", "use_parallel_tiling", "dit_cpu_offload" ] const stringWidgetNames = ["prefix", "quant_config", "lora_config", "image_path"] diff --git a/docs/index.md b/docs/index.md index 6a161f067..0dc484048 100644 --- a/docs/index.md +++ b/docs/index.md @@ -27,7 +27,6 @@ FastVideo has the following features: - State-of-the-art performance optimizations for inference - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507) - - [TeaCache](https://arxiv.org/pdf/2411.19108) - [Sage Attention](https://arxiv.org/abs/2410.02367) - E2E post-training support - Data preprocessing pipeline for video data diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 1c794386e..417f41d46 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -12,9 +12,6 @@ This page describes the various options for speeding up generation times in Fast - [Sage Attention](#sage-attention) - [Sage Attention 3](#sage-attention-3) -- Caching Techniques - - [TeaCache](#teacache) - ## Attention Backends ### Available Backends @@ -117,63 +114,24 @@ These backends are model-specific and require the corresponding kernels and dependencies. Use the support matrix and model examples to confirm compatibility before enabling them. -## Teacache - -TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo. - -### What is TeaCache? - -See the official [TeaCache](https://github.com/ali-vilab/TeaCache) repo and their paper for more details. - -### How to Enable TeaCache - -Enabling TeaCache is straightforward - simply add the `enable_teacache=True` parameter to your `generate_video()` call: - -```python -# ... previous code -generator.generate_video( - prompt="Your prompt here", - sampling_param=params, - enable_teacache=True -) -# more code ... -``` - -### Complete Example - -At the bottom is a complete example of using TeaCache for faster video generation. You can run it using the following command: - -```bash -python examples/inference/optimizations/teacache_example.py -``` - -### Advanced Configuration - -While TeaCache works well with default settings, you can fine-tune its behavior by adjusting the threshold value: - -1. Lower threshold values (e.g., 0.1) will result in more skipped calculations and faster generation with slightly more potential for quality degradation -2. Higher threshold values (e.g., 0.15-0.23) will skip fewer calculations but maintain quality closer to the original - -Note that the optimal threshold depends on your specific model and content. - ## Benchmarking different optimizations -To benchmark the performance improvement, try generating the same video with and without TeaCache enabled and compare the generation times: +To benchmark backend performance, generate the same prompt with the same seed and compare end-to-end generation times: ```python -# Without TeaCache -start_time = time.perf_counter() -generator.generate_video(prompt="Your prompt", enable_teacache=False) -standard_time = time.perf_counter() - start_time +import os +import time -# With TeaCache -start_time = time.perf_counter() -generator.generate_video(prompt="Your prompt", enable_teacache=True) -teacache_time = time.perf_counter() - start_time - -print(f"Standard generation: {standard_time:.2f} seconds") -print(f"TeaCache generation: {teacache_time:.2f} seconds") -print(f"Speedup: {standard_time/teacache_time:.2f}x") +for backend in ["TORCH_SDPA", "FLASH_ATTN", "SAGE_ATTN"]: + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend + generator = VideoGenerator.from_pretrained("your-model-id") + start_time = time.perf_counter() + generator.generate_video( + prompt="Your prompt", + seed=1024, + ) + elapsed = time.perf_counter() - start_time + print(f"{backend}: {elapsed:.2f}s") ``` -Note: If you want to benchmark different attention backends, you'll need to reinstantiate `VideoGenerator`. +Note: reinstantiate `VideoGenerator` after changing `FASTVIDEO_ATTENTION_BACKEND`. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 298cabf7f..f373137cc 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -46,26 +46,26 @@ pipeline initialization and sampling. } -| Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | VSA | BSA | -|------------|---------------------|-------------|----------|-------------------|-----------|-----|-----| -| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | -| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | -| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | -| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P
720P | ❌ | ❌ | ✅ | ⭕ | ⭕ | -| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P
720P | ❌ | ❌ | ✅ | ⭕ | ⭕ | -| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p
544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ | -| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p
544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ | -| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ | ⭕ | -| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅* | ✅ | ⭕ | ⭕ | -| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ | ⭕ | -| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ | -| TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P
720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| LongCat T2V 13.6B | See note** | 480P
720P | ❌ | ❌ | ❌ | ⭕ | ✅ | -| Matrix Game 2.0 Base | `FastVideo/Matrix-Game-2.0-Base-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| Matrix Game 2.0 GTA | `FastVideo/Matrix-Game-2.0-GTA-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | -| Matrix Game 2.0 TempleRun | `FastVideo/Matrix-Game-2.0-TempleRun-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | +| Model Name | HuggingFace Model ID | Resolutions | Sliding Tile Attn | Sage Attn | VSA | BSA | +|------------|---------------------|-------------|-------------------|-----------|-----|-----| +| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ✅ | ⭕ | +| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | +| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ✅ | ⭕ | ⭕ | +| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P
720P | ❌ | ✅ | ⭕ | ⭕ | +| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P
720P | ❌ | ✅ | ⭕ | ⭕ | +| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p
544px960p | ✅ | ✅ | ⭕ | ⭕ | +| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p
544px960p | ✅ | ✅ | ⭕ | ⭕ | +| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅* | ✅ | ⭕ | ⭕ | +| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅* | ✅ | ⭕ | ⭕ | +| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅* | ✅ | ⭕ | ⭕ | +| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ⭕ | ⭕ | +| TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | +| TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ | +| TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P
720P | ⭕ | ⭕ | ⭕ | ⭕ | +| LongCat T2V 13.6B | See note** | 480P
720P | ❌ | ❌ | ⭕ | ✅ | +| Matrix Game 2.0 Base | `FastVideo/Matrix-Game-2.0-Base-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | +| Matrix Game 2.0 GTA | `FastVideo/Matrix-Game-2.0-GTA-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | +| Matrix Game 2.0 TempleRun | `FastVideo/Matrix-Game-2.0-TempleRun-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. diff --git a/examples/inference/optimizations/README.md b/examples/inference/optimizations/README.md index 66ea69318..ee644a90b 100644 --- a/examples/inference/optimizations/README.md +++ b/examples/inference/optimizations/README.md @@ -3,7 +3,3 @@ ```bash python examples/inference/optimizations/attention_example.py ``` - -```bash -python examples/inference/optimizations/teacache_example.py -``` diff --git a/examples/inference/optimizations/teacache_example.py b/examples/inference/optimizations/teacache_example.py deleted file mode 100644 index 77e0582a6..000000000 --- a/examples/inference/optimizations/teacache_example.py +++ /dev/null @@ -1,47 +0,0 @@ -import time - -from fastvideo import VideoGenerator, SamplingParam - - -def main(): - start_time = time.perf_counter() - - gen = VideoGenerator.from_pretrained( - model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - num_gpus=1, - dit_cpu_offload=False, - vae_cpu_offload=False, - text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" - ) - load_time = time.perf_counter() - start_time - print(f"Model loading time: {load_time:.2f} seconds") - - gen_start_time = time.perf_counter() - - params = SamplingParam.from_pretrained( - model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - ) - # this controls the threshold for the tea cache - params.teacache_params.teacache_thresh = 0.08 - gen.generate_video( - prompt= - "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", - sampling_param=params, - height=480, - width=832, - num_frames=61, # 85 ,77 - num_inference_steps=50, - enable_teacache=True, - seed=1024, - output_path="example_outputs/") - - generation_time = time.perf_counter() - gen_start_time - print(f"Video generation time: {generation_time:.2f} seconds") - - total_time = time.perf_counter() - start_time - print(f"Total execution time: {total_time:.2f} seconds") - - -if __name__ == "__main__": - main() diff --git a/fastvideo/configs/models/base.py b/fastvideo/configs/models/base.py index 39d00c178..df235dc1b 100644 --- a/fastvideo/configs/models/base.py +++ b/fastvideo/configs/models/base.py @@ -24,7 +24,7 @@ class ModelConfig: arch_config: ArchConfig = field(default_factory=ArchConfig) # FastVideo-specific parameters here - # i.e. STA, quantization, teacache + # i.e. STA, quantization def __getattr__(self, name): # Only called if 'name' is not found in ModelConfig directly diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py index 8656d1428..2728899fb 100644 --- a/fastvideo/configs/sample/base.py +++ b/fastvideo/configs/sample/base.py @@ -73,9 +73,6 @@ class SamplingParam: boundary_ratio: float | None = None sigmas: list[float] | None = None - # TeaCache parameters - enable_teacache: bool = False - # Misc save_video: bool = True return_frames: bool = False diff --git a/fastvideo/configs/sample/hunyuan.py b/fastvideo/configs/sample/hunyuan.py index 49d01fbfe..5441829af 100644 --- a/fastvideo/configs/sample/hunyuan.py +++ b/fastvideo/configs/sample/hunyuan.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass, field +from dataclasses import dataclass from fastvideo.configs.sample.base import SamplingParam -from fastvideo.configs.sample.teacache import TeaCacheParams @dataclass @@ -16,14 +15,6 @@ class HunyuanSamplingParam(SamplingParam): guidance_scale: float = 1.0 - teacache_params: TeaCacheParams = field( - default_factory=lambda: TeaCacheParams( - teacache_thresh=0.15, - coefficients=[ - 7.33226126e+02, -4.01131952e+02, 6.75869174e+01, - -3.14987800e+00, 9.61237896e-02 - ])) - @dataclass class FastHunyuanSamplingParam(HunyuanSamplingParam): diff --git a/fastvideo/configs/sample/hunyuangamecraft.py b/fastvideo/configs/sample/hunyuangamecraft.py index 897c5cf8b..853745322 100644 --- a/fastvideo/configs/sample/hunyuangamecraft.py +++ b/fastvideo/configs/sample/hunyuangamecraft.py @@ -5,11 +5,10 @@ Sampling parameters for HunyuanGameCraft video generation. GameCraft generates game-like videos with camera/action control. Default parameters are based on the official implementation. """ -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any from fastvideo.configs.sample.base import SamplingParam -from fastvideo.configs.sample.teacache import TeaCacheParams @dataclass @@ -64,16 +63,6 @@ class HunyuanGameCraftSamplingParam(SamplingParam): # Number of conditioning frames (for autoregressive) - maps to num_cond_frames num_cond_frames: int = 0 - # TeaCache parameters (if enabled) - teacache_params: TeaCacheParams = field( - default_factory=lambda: TeaCacheParams( - teacache_thresh=0.15, - coefficients=[ - 7.33226126e+02, -4.01131952e+02, 6.75869174e+01, - -3.14987800e+00, 9.61237896e-02 - ], - )) - def __post_init__(self) -> None: super().__post_init__() # Validate action lists diff --git a/fastvideo/configs/sample/teacache.py b/fastvideo/configs/sample/teacache.py deleted file mode 100644 index 060353074..000000000 --- a/fastvideo/configs/sample/teacache.py +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass, field - -from fastvideo.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/configs/sample/wan.py b/fastvideo/configs/sample/wan.py index 85f76a7ac..a96cf0d29 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass, field +from dataclasses import dataclass from fastvideo.configs.sample.base import SamplingParam -from fastvideo.configs.sample.teacache import WanTeaCacheParams @dataclass @@ -18,18 +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 - teacache_params: WanTeaCacheParams = field( - default_factory=lambda: WanTeaCacheParams( - teacache_thresh=0.08, - ret_steps_coeffs=[ - -5.21862437e+04, 9.23041404e+03, -5.28275948e+02, - 1.36987616e+01, -4.99875664e-02 - ], - non_ret_steps_coeffs=[ - 2.39676752e+03, -1.31110545e+03, 2.01331979e+02, - -8.29855975e+00, 1.37887774e-01 - ])) - @dataclass class WanT2V_14B_SamplingParam(SamplingParam): @@ -44,19 +31,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 - teacache_params: WanTeaCacheParams = field( - default_factory=lambda: WanTeaCacheParams( - teacache_thresh=0.20, - use_ret_steps=False, - ret_steps_coeffs=[ - -3.03318725e+05, 4.90537029e+04, -2.65530556e+03, - 5.87365115e+01, -3.15583525e-01 - ], - non_ret_steps_coeffs=[ - -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, - -13.02252404 - ])) - @dataclass class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParam): @@ -64,18 +38,6 @@ class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParam): guidance_scale: float = 5.0 num_inference_steps: int = 40 - teacache_params: WanTeaCacheParams = field( - default_factory=lambda: WanTeaCacheParams( - teacache_thresh=0.26, - ret_steps_coeffs=[ - -3.03318725e+05, 4.90537029e+04, -2.65530556e+03, - 5.87365115e+01, -3.15583525e-01 - ], - non_ret_steps_coeffs=[ - -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, - -13.02252404 - ])) - @dataclass class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam): @@ -83,18 +45,6 @@ class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam): guidance_scale: float = 5.0 num_inference_steps: int = 40 - teacache_params: WanTeaCacheParams = field( - default_factory=lambda: WanTeaCacheParams( - teacache_thresh=0.3, - ret_steps_coeffs=[ - -3.03318725e+05, 4.90537029e+04, -2.65530556e+03, - 5.87365115e+01, -3.15583525e-01 - ], - non_ret_steps_coeffs=[ - -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, - -13.02252404 - ])) - @dataclass class FastWanT2V480P_SamplingParam(WanT2V_1_3B_SamplingParam): @@ -129,8 +79,6 @@ class Wan2_1_Fun_1_3B_Control_SamplingParam(SamplingParam): height: int = 832 width: int = 480 guidance_scale: float = 6.0 - teacache_params: WanTeaCacheParams = field( - default_factory=lambda: WanTeaCacheParams(teacache_thresh=0.1, )) # ============================================= diff --git a/fastvideo/models/dits/base.py b/fastvideo/models/dits/base.py index d60812a45..5a9b1d276 100644 --- a/fastvideo/models/dits/base.py +++ b/fastvideo/models/dits/base.py @@ -68,60 +68,3 @@ class BaseDiT(nn.Module, ABC): @property def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]: return self._supported_attention_backends - - -class CachableDiT(BaseDiT): - """ - An intermediate base class that adds TeaCache optimization functionality to DiT models. - TeaCache accelerates inference by selectively skipping redundant computation when consecutive - diffusion steps are similar enough. - """ - # These are required class attributes that should be overridden by concrete implementations - _fsdp_shard_conditions = [] - param_names_mapping = {} - reverse_param_names_mapping = {} - lora_param_names_mapping: dict = {} - # Ensure these instance attributes are properly defined in subclasses - hidden_size: int - num_attention_heads: int - num_channels_latents: int - # always supports torch_sdpa - _supported_attention_backends: tuple[ - AttentionBackendEnum, ...] = DiTConfig()._supported_attention_backends - - def __init__(self, config: DiTConfig, **kwargs) -> None: - super().__init__(config, **kwargs) - - self.cnt = 0 - self.teacache_thresh = 0 - self.coefficients: list[float] = [] - - # NOTE(will): Only wan2.1 needs these, so we are hardcoding it here - if self.config.prefix == "wan": - self.use_ret_steps = self.config.cache_config.use_ret_steps - self.is_even = False - self.previous_residual_even: torch.Tensor | None = None - self.previous_residual_odd: torch.Tensor | None = None - self.accumulated_rel_l1_distance_even = 0 - self.accumulated_rel_l1_distance_odd = 0 - self.should_calc_even = True - self.should_calc_odd = True - else: - self.accumulated_rel_l1_distance = 0 - self.previous_modulated_input = None - self.previous_resiual = None - self.previous_e0_even: torch.Tensor | None = None - self.previous_e0_odd: torch.Tensor | None = None - - def maybe_cache_states(self, hidden_states: torch.Tensor, - original_hidden_states: torch.Tensor) -> None: - pass - - def should_skip_forward_for_cached_states(self, - **kwargs: dict[str, Any]) -> bool: - return False - - def retrieve_cached_states(self, - hidden_states: torch.Tensor) -> torch.Tensor: - raise NotImplementedError( - "maybe_retrieve_cached_states is not implemented") diff --git a/fastvideo/models/dits/cosmos.py b/fastvideo/models/dits/cosmos.py index 4028307ec..e4a703c56 100644 --- a/fastvideo/models/dits/cosmos.py +++ b/fastvideo/models/dits/cosmos.py @@ -7,16 +7,12 @@ import numpy as np import torch import torch.nn as nn -from fastvideo.attention import DistributedAttention, LocalAttention from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig -from fastvideo.forward_context import get_forward_context from fastvideo.layers.layernorm import RMSNorm -from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import apply_rotary_emb from fastvideo.layers.visual_embedding import Timesteps from fastvideo.models.dits.base import BaseDiT -from fastvideo.platforms import AttentionBackendEnum class CosmosPatchEmbed(nn.Module): @@ -605,18 +601,6 @@ class CosmosTransformer3DModel(BaseDiT): bias=False) self.gradient_checkpointing = False - - # For TeaCache - self.previous_e0_even = None - self.previous_e0_odd = None - self.previous_residual_even = None - self.previous_residual_odd = None - self.is_even = True - self.should_calc_even = True - self.should_calc_odd = True - self.accumulated_rel_l1_distance_even = 0 - self.accumulated_rel_l1_distance_odd = 0 - self.cnt = 0 self.__post_init__() def forward(self, @@ -628,10 +612,6 @@ class CosmosTransformer3DModel(BaseDiT): condition_mask: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None, **kwargs) -> torch.Tensor: - forward_batch = get_forward_context().forward_batch - enable_teacache = forward_batch is not None and forward_batch.enable_teacache - - orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] diff --git a/fastvideo/models/dits/hunyuangamecraft.py b/fastvideo/models/dits/hunyuangamecraft.py index 8b71c1c1c..f993a55e6 100644 --- a/fastvideo/models/dits/hunyuangamecraft.py +++ b/fastvideo/models/dits/hunyuangamecraft.py @@ -20,7 +20,7 @@ from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify -from fastvideo.models.dits.base import CachableDiT +from fastvideo.models.dits.base import BaseDiT from fastvideo.models.dits.hunyuanvideo import ( MMDoubleStreamBlock, MMSingleStreamBlock, @@ -178,7 +178,7 @@ class CameraNet(nn.Module): return camera_states * self.scale -class HunyuanGameCraftTransformer3DModel(CachableDiT): +class HunyuanGameCraftTransformer3DModel(BaseDiT): """ HunyuanGameCraft Transformer - ported from official implementation. """ diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/models/dits/hunyuanvideo.py index 7396b29f2..a47751b17 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/models/dits/hunyuanvideo.py @@ -2,26 +2,20 @@ from typing import Any -import numpy as np import torch import torch.nn as nn from fastvideo.attention import DistributedAttention, LocalAttention from fastvideo.configs.models.dits import HunyuanVideoConfig -from fastvideo.configs.sample.teacache import TeaCacheParams -from fastvideo.distributed.parallel_state import get_sp_world_size -from fastvideo.forward_context import get_forward_context from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear # TODO(will-PY-refactor): RMSNorm .... from fastvideo.layers.mlp import MLP -from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, - get_rotary_pos_embed) +from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify) -from fastvideo.models.dits.base import CachableDiT -from fastvideo.models.utils import modulate +from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum from fastvideo.distributed.communication_op import sequence_model_parallel_shard, sequence_model_parallel_all_gather @@ -409,7 +403,7 @@ class MMSingleStreamBlock(nn.Module): return self.output_residual(x, output, mod_gate) -class HunyuanVideoTransformer3DModel(CachableDiT): +class HunyuanVideoTransformer3DModel(BaseDiT): """ HunyuanVideo Transformer backbone adapted for distributed training. @@ -550,10 +544,6 @@ class HunyuanVideoTransformer3DModel(CachableDiT): Returns: Tuple of (output) """ - forward_context = get_forward_context() - forward_batch = forward_context.forward_batch - enable_teacache = forward_batch is not None and forward_batch.enable_teacache - if guidance is None: guidance = torch.tensor([6016.0], device=hidden_states.device, @@ -602,38 +592,26 @@ class HunyuanVideoTransformer3DModel(CachableDiT): freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None - should_skip_forward = self.should_skip_forward_for_cached_states( - img=img, vec=vec) + # Process through double stream blocks + for index, block in enumerate(self.double_blocks): + double_block_args = [img, txt, vec, freqs_cis] + img, txt = block(*double_block_args) + # Merge txt and img to pass through single stream blocks + x = torch.cat((img, txt), 1) - if should_skip_forward: - img = self.retrieve_cached_states(img) - else: - if enable_teacache: - original_img = img.clone() + # Process through single stream blocks + if len(self.single_blocks) > 0: + for index, block in enumerate(self.single_blocks): + single_block_args = [ + x, + vec, + txt_seq_len, + freqs_cis, + ] + x = block(*single_block_args) - # Process through double stream blocks - for index, block in enumerate(self.double_blocks): - double_block_args = [img, txt, vec, freqs_cis] - img, txt = block(*double_block_args) - # Merge txt and img to pass through single stream blocks - x = torch.cat((img, txt), 1) - - # Process through single stream blocks - if len(self.single_blocks) > 0: - for index, block in enumerate(self.single_blocks): - single_block_args = [ - x, - vec, - txt_seq_len, - freqs_cis, - ] - x = block(*single_block_args) - - # Extract image features - img = x[:, :img_seq_len, ...] - - if enable_teacache: - self.maybe_cache_states(img, original_img) + # Extract image features + img = x[:, :img_seq_len, ...] # Final layer processing img = sequence_model_parallel_all_gather(img, dim=1) @@ -643,98 +621,6 @@ class HunyuanVideoTransformer3DModel(CachableDiT): return img - def maybe_cache_states(self, hidden_states: torch.Tensor, - original_hidden_states: torch.Tensor) -> None: - self.previous_residual = hidden_states - original_hidden_states - - def should_skip_forward_for_cached_states(self, **kwargs) -> bool: - - forward_context = get_forward_context() - forward_batch = forward_context.forward_batch - if forward_batch is None: - return False - current_timestep = forward_context.current_timestep - enable_teacache = forward_batch.enable_teacache - - if not enable_teacache: - return False - raise NotImplementedError( - "teacache is not supported yet for HunyuanVideo") - - teacache_params = forward_batch.teacache_params - assert teacache_params is not None, "teacache_params is not initialized" - assert isinstance( - teacache_params, - TeaCacheParams), "teacache_params is not a TeaCacheParams" - num_inference_steps = forward_batch.num_inference_steps - teache_thresh = teacache_params.teacache_thresh - - coefficients = teacache_params.coefficients - - if current_timestep == 0: - self.cnt = 0 - - inp = kwargs["img"].clone() - vec_ = kwargs["vec"].clone() - # convert to DTensor - vec_ = torch.distributed.tensor.DTensor.from_local( - vec_, - torch.distributed.DeviceMesh("cuda", - list(range(get_sp_world_size())), - mesh_dim_names=("dp", )), - [torch.distributed.tensor.Replicate()]) - - inp = torch.distributed.tensor.DTensor.from_local( - inp, - torch.distributed.DeviceMesh("cuda", - list(range(get_sp_world_size())), - mesh_dim_names=("dp", )), - [torch.distributed.tensor.Replicate()]) - - # txt_ = kwargs["txt"].clone() - - # inp = img.clone() - # vec_ = vec.clone() - # txt_ = txt.clone() - ( - img_mod1_shift, - img_mod1_scale, - img_mod1_gate, - img_mod2_shift, - img_mod2_scale, - img_mod2_gate, - ) = self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1) - normed_inp = self.double_blocks[0].img_attn_norm.norm(inp) - modulated_inp = modulate(normed_inp, - shift=img_mod1_shift, - scale=img_mod1_scale) - if self.cnt == 0 or self.cnt == num_inference_steps - 1: - should_calc = True - self.accumulated_rel_l1_distance = 0 - else: - coefficients = [ - 7.33226126e+02, -4.01131952e+02, 6.75869174e+01, - -3.14987800e+00, 9.61237896e-02 - ] - rescale_func = np.poly1d(coefficients) - assert self.previous_modulated_input is not None, "previous_modulated_input is not initialized" - self.accumulated_rel_l1_distance += rescale_func( - ((modulated_inp - self.previous_modulated_input).abs().mean() / - self.previous_modulated_input.abs().mean()).cpu().item()) - if self.accumulated_rel_l1_distance < teache_thresh: - should_calc = False - else: - should_calc = True - self.accumulated_rel_l1_distance = 0 - self.previous_modulated_input = modulated_inp - self.cnt += 1 - - return not should_calc - - def retrieve_cached_states(self, - hidden_states: torch.Tensor) -> torch.Tensor: - return hidden_states + self.previous_residual - class SingleTokenRefiner(nn.Module): """ diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py index ac5ffb11d..0da8815e2 100644 --- a/fastvideo/models/dits/hunyuanvideo15.py +++ b/fastvideo/models/dits/hunyuanvideo15.py @@ -31,7 +31,7 @@ from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify) -from fastvideo.models.dits.base import CachableDiT +from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum from fastvideo.logger import init_logger from fastvideo.forward_context import set_forward_context @@ -389,7 +389,7 @@ class MMDoubleStreamBlock(nn.Module): return img, txt -class HunyuanVideo15Transformer3DModel(CachableDiT): +class HunyuanVideo15Transformer3DModel(BaseDiT): r""" A Transformer model for video-like data used in [HunyuanVideo1.5](https://huggingface.co/tencent/HunyuanVideo1.5). """ diff --git a/fastvideo/models/dits/lingbotworld/model.py b/fastvideo/models/dits/lingbotworld/model.py index 9b32a1e8a..4cc84cb78 100644 --- a/fastvideo/models/dits/lingbotworld/model.py +++ b/fastvideo/models/dits/lingbotworld/model.py @@ -1,20 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 import math -from contextlib import nullcontext from typing import Any -import numpy as np import torch import torch.nn as nn from fastvideo.attention import DistributedAttention from fastvideo.configs.models.dits.lingbotworld import LingBotWorldVideoConfig -from fastvideo.configs.sample.wan import WanTeaCacheParams from fastvideo.distributed.communication_op import ( sequence_model_parallel_all_gather_with_unpad, sequence_model_parallel_shard) -from fastvideo.forward_context import get_forward_context from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) @@ -25,7 +21,7 @@ from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (PatchEmbed, WanCamControlPatchEmbedding) from fastvideo.logger import init_logger -from fastvideo.models.dits.base import CachableDiT +from fastvideo.models.dits.base import BaseDiT from fastvideo.models.dits.wanvideo import ( WanI2VCrossAttention, WanT2VCrossAttention, @@ -226,7 +222,7 @@ class LingBotWorldTransformerBlock(nn.Module): return hidden_states -class LingBotWorldTransformer3DModel(CachableDiT): +class LingBotWorldTransformer3DModel(BaseDiT): _fsdp_shard_conditions = LingBotWorldVideoConfig()._fsdp_shard_conditions _compile_conditions = LingBotWorldVideoConfig()._compile_conditions _supported_attention_backends = LingBotWorldVideoConfig( @@ -297,18 +293,6 @@ class LingBotWorldTransformer3DModel(CachableDiT): self.gradient_checkpointing = False self._logged_attention_mask = False - - # For type checking - self.previous_e0_even = None - self.previous_e0_odd = None - self.previous_residual_even = None - self.previous_residual_odd = None - self.is_even = True - self.should_calc_even = True - self.should_calc_odd = True - self.accumulated_rel_l1_distance_even = 0 - self.accumulated_rel_l1_distance_odd = 0 - self.cnt = 0 self.__post_init__() def forward(self, @@ -320,9 +304,6 @@ class LingBotWorldTransformer3DModel(CachableDiT): guidance=None, c2ws_plucker_emb: torch.Tensor | None = None, **kwargs) -> torch.Tensor: - forward_batch = get_forward_context().forward_batch - enable_teacache = forward_batch is not None and forward_batch.enable_teacache - orig_dtype = hidden_states.dtype if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] @@ -421,31 +402,16 @@ class LingBotWorldTransformer3DModel(CachableDiT): assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks - # if caching is enabled, we might be able to skip the forward pass - should_skip_forward = self.should_skip_forward_for_cached_states( - timestep_proj=timestep_proj, temb=temb) - - if should_skip_forward: - print("skipping forward, cached") - hidden_states = self.retrieve_cached_states(hidden_states) + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, attention_mask, c2ws_plucker_emb) else: - # if teacache is enabled, we need to cache the original hidden states - if enable_teacache: - original_hidden_states = hidden_states.clone() - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, attention_mask, c2ws_plucker_emb) - else: - for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, attention_mask, c2ws_plucker_emb) - # if teacache is enabled, we need to cache the original hidden states - - if enable_teacache: - self.maybe_cache_states(hidden_states, original_hidden_states) + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, attention_mask, + c2ws_plucker_emb) # 5. Output norm, projection & unpatchify if temb.dim() == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) @@ -473,97 +439,5 @@ class LingBotWorldTransformer3DModel(CachableDiT): return output - def maybe_cache_states(self, hidden_states: torch.Tensor, - original_hidden_states: torch.Tensor) -> None: - if self.is_even: - self.previous_residual_even = hidden_states.squeeze( - 0) - original_hidden_states - else: - self.previous_residual_odd = hidden_states.squeeze( - 0) - original_hidden_states - - def should_skip_forward_for_cached_states(self, **kwargs) -> bool: - - forward_context = get_forward_context() - forward_batch = forward_context.forward_batch - if forward_batch is None or not forward_batch.enable_teacache: - return False - teacache_params = forward_batch.teacache_params - assert teacache_params is not None, "teacache_params is not initialized" - assert isinstance( - teacache_params, - WanTeaCacheParams), "teacache_params is not a WanTeaCacheParams" - current_timestep = forward_context.current_timestep - num_inference_steps = forward_batch.num_inference_steps - - # initialize the coefficients, cutoff_steps, and ret_steps - coefficients = teacache_params.coefficients - use_ret_steps = teacache_params.use_ret_steps - cutoff_steps = teacache_params.get_cutoff_steps(num_inference_steps) - ret_steps = teacache_params.ret_steps - teacache_thresh = teacache_params.teacache_thresh - - if current_timestep == 0: - self.cnt = 0 - - timestep_proj = kwargs["timestep_proj"] - temb = kwargs["temb"] - modulated_inp = timestep_proj if use_ret_steps else temb - - if self.cnt % 2 == 0: # even -> condition - self.is_even = True - if self.cnt < ret_steps or self.cnt >= cutoff_steps: - self.should_calc_even = True - self.accumulated_rel_l1_distance_even = 0 - else: - assert self.previous_e0_even is not None, "previous_e0_even is not initialized" - assert self.accumulated_rel_l1_distance_even is not None, "accumulated_rel_l1_distance_even is not initialized" - rescale_func = np.poly1d(coefficients) - self.accumulated_rel_l1_distance_even += rescale_func( - ((modulated_inp - self.previous_e0_even).abs().mean() / - self.previous_e0_even.abs().mean()).cpu().item()) - if self.accumulated_rel_l1_distance_even < teacache_thresh: - self.should_calc_even = False - else: - self.should_calc_even = True - self.accumulated_rel_l1_distance_even = 0 - self.previous_e0_even = modulated_inp.clone() - - else: # odd -> unconditon - self.is_even = False - if self.cnt < ret_steps or self.cnt >= cutoff_steps: - self.should_calc_odd = True - self.accumulated_rel_l1_distance_odd = 0 - else: - assert self.previous_e0_odd is not None, "previous_e0_odd is not initialized" - assert self.accumulated_rel_l1_distance_odd is not None, "accumulated_rel_l1_distance_odd is not initialized" - rescale_func = np.poly1d(coefficients) - self.accumulated_rel_l1_distance_odd += rescale_func( - ((modulated_inp - self.previous_e0_odd).abs().mean() / - self.previous_e0_odd.abs().mean()).cpu().item()) - if self.accumulated_rel_l1_distance_odd < teacache_thresh: - self.should_calc_odd = False - else: - self.should_calc_odd = True - self.accumulated_rel_l1_distance_odd = 0 - self.previous_e0_odd = modulated_inp.clone() - self.cnt += 1 - should_skip_forward = False - if self.is_even: - if not self.should_calc_even: - should_skip_forward = True - else: - if not self.should_calc_odd: - should_skip_forward = True - - return should_skip_forward - - def retrieve_cached_states(self, - hidden_states: torch.Tensor) -> torch.Tensor: - if self.is_even: - return hidden_states + self.previous_residual_even - else: - return hidden_states + self.previous_residual_odd - # Entry point for model registry EntryClass = LingBotWorldTransformer3DModel diff --git a/fastvideo/models/dits/longcat.py b/fastvideo/models/dits/longcat.py index 385469c6c..03e89e63c 100644 --- a/fastvideo/models/dits/longcat.py +++ b/fastvideo/models/dits/longcat.py @@ -17,7 +17,7 @@ from fastvideo.layers.layernorm import RMSNorm, FP32LayerNorm from fastvideo.layers.activation import get_act_fn from fastvideo.layers.rotary_embedding_3d import RotaryPositionalEmbedding3D from fastvideo.attention.layer import DistributedAttention, LocalAttention -from fastvideo.models.dits.base import CachableDiT +from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum from fastvideo.third_party.longcat_video.block_sparse_attention.bsa_interface import flash_attn_bsa_3d @@ -929,7 +929,7 @@ class FinalLayer(nn.Module): # Main Model # ============================================================================ -class LongCatTransformer3DModel(CachableDiT): +class LongCatTransformer3DModel(BaseDiT): """ Native LongCat Video Transformer using FastVideo layers. """ diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index cb041c11e..bb582f3d3 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -30,7 +30,7 @@ from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_wo from fastvideo.distributed.utils import create_attention_mask_for_padding from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context from fastvideo.logger import init_logger -from fastvideo.models.dits.base import CachableDiT +from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum logger = init_logger(__name__) @@ -2204,7 +2204,7 @@ class LTXModel(torch.nn.Module): return vx, ax -class LTX2Transformer3DModel(CachableDiT): +class LTX2Transformer3DModel(BaseDiT): """ LTX-2 transformer using native FastVideo LTX-2 modules. """ diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index cb8252005..86aed7868 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -1,10 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 import math -from contextlib import nullcontext from typing import Any -import numpy as np import torch import torch.nn as nn @@ -12,12 +10,9 @@ import fastvideo.envs as envs from fastvideo.attention import (DistributedAttention, DistributedAttention_VSA, LocalAttention) from fastvideo.configs.models.dits import WanVideoConfig -from fastvideo.configs.sample.wan import WanTeaCacheParams from fastvideo.distributed.communication_op import ( - sequence_model_parallel_all_gather, sequence_model_parallel_all_gather_with_unpad, sequence_model_parallel_shard) -from fastvideo.forward_context import get_forward_context from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) @@ -29,7 +24,7 @@ from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder) from fastvideo.logger import init_logger -from fastvideo.models.dits.base import CachableDiT +from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum, current_platform from fastvideo.distributed.parallel_state import get_sp_world_size @@ -548,7 +543,7 @@ class WanTransformerBlock_VSA(nn.Module): return hidden_states -class WanTransformer3DModel(CachableDiT): +class WanTransformer3DModel(BaseDiT): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions _supported_attention_backends = WanVideoConfig( @@ -616,18 +611,6 @@ class WanTransformer3DModel(CachableDiT): self.gradient_checkpointing = False self._logged_attention_mask = False - - # For type checking - self.previous_e0_even = None - self.previous_e0_odd = None - self.previous_residual_even = None - self.previous_residual_odd = None - self.is_even = True - self.should_calc_even = True - self.should_calc_odd = True - self.accumulated_rel_l1_distance_even = 0 - self.accumulated_rel_l1_distance_odd = 0 - self.cnt = 0 self.__post_init__() def forward(self, @@ -638,9 +621,6 @@ class WanTransformer3DModel(CachableDiT): | None = None, guidance=None, **kwargs) -> torch.Tensor: - forward_batch = get_forward_context().forward_batch - enable_teacache = forward_batch is not None and forward_batch.enable_teacache - orig_dtype = hidden_states.dtype if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] @@ -728,31 +708,15 @@ class WanTransformer3DModel(CachableDiT): assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks - # if caching is enabled, we might be able to skip the forward pass - should_skip_forward = self.should_skip_forward_for_cached_states( - timestep_proj=timestep_proj, temb=temb) - - if should_skip_forward: - print("skipping forward, cached") - hidden_states = self.retrieve_cached_states(hidden_states) + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, attention_mask) else: - # if teacache is enabled, we need to cache the original hidden states - if enable_teacache: - original_hidden_states = hidden_states.clone() - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, attention_mask) - else: - for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, attention_mask) - # if teacache is enabled, we need to cache the original hidden states - - if enable_teacache: - self.maybe_cache_states(hidden_states, original_hidden_states) + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, attention_mask) # 5. Output norm, projection & unpatchify if temb.dim() == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) @@ -780,97 +744,5 @@ class WanTransformer3DModel(CachableDiT): return output - def maybe_cache_states(self, hidden_states: torch.Tensor, - original_hidden_states: torch.Tensor) -> None: - if self.is_even: - self.previous_residual_even = hidden_states.squeeze( - 0) - original_hidden_states - else: - self.previous_residual_odd = hidden_states.squeeze( - 0) - original_hidden_states - - def should_skip_forward_for_cached_states(self, **kwargs) -> bool: - - forward_context = get_forward_context() - forward_batch = forward_context.forward_batch - if forward_batch is None or not forward_batch.enable_teacache: - return False - teacache_params = forward_batch.teacache_params - assert teacache_params is not None, "teacache_params is not initialized" - assert isinstance( - teacache_params, - WanTeaCacheParams), "teacache_params is not a WanTeaCacheParams" - current_timestep = forward_context.current_timestep - num_inference_steps = forward_batch.num_inference_steps - - # initialize the coefficients, cutoff_steps, and ret_steps - coefficients = teacache_params.coefficients - use_ret_steps = teacache_params.use_ret_steps - cutoff_steps = teacache_params.get_cutoff_steps(num_inference_steps) - ret_steps = teacache_params.ret_steps - teacache_thresh = teacache_params.teacache_thresh - - if current_timestep == 0: - self.cnt = 0 - - timestep_proj = kwargs["timestep_proj"] - temb = kwargs["temb"] - modulated_inp = timestep_proj if use_ret_steps else temb - - if self.cnt % 2 == 0: # even -> condition - self.is_even = True - if self.cnt < ret_steps or self.cnt >= cutoff_steps: - self.should_calc_even = True - self.accumulated_rel_l1_distance_even = 0 - else: - assert self.previous_e0_even is not None, "previous_e0_even is not initialized" - assert self.accumulated_rel_l1_distance_even is not None, "accumulated_rel_l1_distance_even is not initialized" - rescale_func = np.poly1d(coefficients) - self.accumulated_rel_l1_distance_even += rescale_func( - ((modulated_inp - self.previous_e0_even).abs().mean() / - self.previous_e0_even.abs().mean()).cpu().item()) - if self.accumulated_rel_l1_distance_even < teacache_thresh: - self.should_calc_even = False - else: - self.should_calc_even = True - self.accumulated_rel_l1_distance_even = 0 - self.previous_e0_even = modulated_inp.clone() - - else: # odd -> unconditon - self.is_even = False - if self.cnt < ret_steps or self.cnt >= cutoff_steps: - self.should_calc_odd = True - self.accumulated_rel_l1_distance_odd = 0 - else: - assert self.previous_e0_odd is not None, "previous_e0_odd is not initialized" - assert self.accumulated_rel_l1_distance_odd is not None, "accumulated_rel_l1_distance_odd is not initialized" - rescale_func = np.poly1d(coefficients) - self.accumulated_rel_l1_distance_odd += rescale_func( - ((modulated_inp - self.previous_e0_odd).abs().mean() / - self.previous_e0_odd.abs().mean()).cpu().item()) - if self.accumulated_rel_l1_distance_odd < teacache_thresh: - self.should_calc_odd = False - else: - self.should_calc_odd = True - self.accumulated_rel_l1_distance_odd = 0 - self.previous_e0_odd = modulated_inp.clone() - self.cnt += 1 - should_skip_forward = False - if self.is_even: - if not self.should_calc_even: - should_skip_forward = True - else: - if not self.should_calc_odd: - should_skip_forward = True - - return should_skip_forward - - def retrieve_cached_states(self, - hidden_states: torch.Tensor) -> torch.Tensor: - if self.is_even: - return hidden_states + self.previous_residual_even - else: - return hidden_states + self.previous_residual_odd - # Entry point for model registry EntryClass = WanTransformer3DModel diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 3ec7340ed..ea97f2e8b 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -21,7 +21,6 @@ import time from collections import OrderedDict from fastvideo.attention import AttentionMetadata -from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams class PipelineLoggingInfo: @@ -205,10 +204,6 @@ class ForwardBatch: save_video: bool = True return_frames: bool = False - # TeaCache parameters - enable_teacache: bool = False - teacache_params: TeaCacheParams | WanTeaCacheParams | None = None - # STA parameters STA_param: list | None = None is_cfg_negative: bool = False diff --git a/fastvideo/tests/transformers/test_hunyuangamecraft.py b/fastvideo/tests/transformers/test_hunyuangamecraft.py index a64a6133a..db1e5881f 100644 --- a/fastvideo/tests/transformers/test_hunyuangamecraft.py +++ b/fastvideo/tests/transformers/test_hunyuangamecraft.py @@ -112,7 +112,7 @@ def test_hunyuangamecraft_transformer(): encoder_hidden_states = [text_states, text_states_2] - forward_batch = ForwardBatch(data_type="video", enable_teacache=False) + forward_batch = ForwardBatch(data_type="video") with torch.no_grad(): with torch.amp.autocast('cuda', dtype=precision): diff --git a/tests/local_tests/transformers/test_gamecraft_parity.py b/tests/local_tests/transformers/test_gamecraft_parity.py index 2dcc96988..fd5a35b7a 100644 --- a/tests/local_tests/transformers/test_gamecraft_parity.py +++ b/tests/local_tests/transformers/test_gamecraft_parity.py @@ -402,7 +402,7 @@ def test_gamecraft_transformer_parity(): from fastvideo.forward_context import set_forward_context from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - forward_batch = ForwardBatch(data_type="video", enable_teacache=False) + forward_batch = ForwardBatch(data_type="video") # Run FastVideo model with torch.no_grad():