nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Misc] Remove Teacache (#1121)

Browse Source
This commit is contained in:
Zhang Peiyuan
2026-02-22 16:53:07 -08:00
committed by GitHub
parent 8092f02e6d
commit 128d0b7fc5
27 changed files with 94 additions and 771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui/examples/FastHunyuan-diffusers.json
+1 -7
View File
@@ -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
}
}
comfyui/examples/Wan2.1-I2V-14B-480P-Diffusers.json
+2 -8
View File
@@ -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
}
}
comfyui/video_generator/inference_args.py
-5
View File
@@ -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
comfyui/web/js/FastVideo.core.js
+1 -1
View File
@@ -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"]
docs/index.md
-1
View File
@@ -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
docs/inference/optimizations.md
+14 -56
View File
@@ -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`.
docs/inference/support_matrix.md
+20 -20
View File
@@ -46,26 +46,26 @@ pipeline initialization and sampling.
}
</style>
| 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<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>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<br>720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| LongCat T2V 13.6B | See note** | 480P<br>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<br>720P | ❌ | ✅ | ⭕ | ⭕ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ✅ | ⭕ | ⭕ |
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ✅ | ✅ | ⭕ | ⭕ |
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>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<br>720P | ⭕ | ⭕ | ⭕ | ⭕ |
| LongCat T2V 13.6B | See note** | 480P<br>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.
examples/inference/optimizations/README.md
-4
View File
@@ -3,7 +3,3 @@
```bash
python examples/inference/optimizations/attention_example.py
```
```bash
python examples/inference/optimizations/teacache_example.py
```
examples/inference/optimizations/teacache_example.py
-47
View File
@@ -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()
fastvideo/configs/models/base.py
+1 -1
View File
@@ -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
fastvideo/configs/sample/base.py
-3
View File
@@ -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
fastvideo/configs/sample/hunyuan.py
+1 -10
View File
@@ -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):
fastvideo/configs/sample/hunyuangamecraft.py
+1 -12
View File
@@ -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
fastvideo/configs/sample/teacache.py
-41
View File
@@ -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
fastvideo/configs/sample/wan.py
+1 -53
View File
@@ -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, ))
# =============================================
fastvideo/models/dits/base.py
-57
View File
@@ -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")
fastvideo/models/dits/cosmos.py
-20
View File
@@ -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]
fastvideo/models/dits/hunyuangamecraft.py
+2 -2
View File
@@ -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.
"""
fastvideo/models/dits/hunyuanvideo.py
+21 -135
View File
@@ -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):
"""
fastvideo/models/dits/hunyuanvideo15.py
+2 -2
View File
@@ -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).
"""
fastvideo/models/dits/lingbotworld/model.py
+11 -137
View File
@@ -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
fastvideo/models/dits/longcat.py
+2 -2
View File
@@ -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.
"""
fastvideo/models/dits/ltx2.py
+2 -2
View File
@@ -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.
"""
fastvideo/models/dits/wanvideo.py
+10 -138
View File
@@ -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
fastvideo/pipelines/pipeline_batch_info.py
-5
View File
@@ -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
fastvideo/tests/transformers/test_hunyuangamecraft.py
+1 -1
View File
@@ -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):
tests/local_tests/transformers/test_gamecraft_parity.py
+1 -1
View File
@@ -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():