nodes/welltop-cn-ComfyUI-TeaCache
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

support Wan2.1 series model

Browse Source
This commit is contained in:
YunjieYu
2025-03-06 21:17:28 +08:00
parent 062efbab2e
commit 71c535a082
6 changed files with 1557 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+25 -5
View File
@@ -6,6 +6,9 @@ Timestep Embedding Aware Cache ([TeaCache](https://github.com/ali-vilab/TeaCache
TeaCache has now been integrated into ComfyUI and is compatible with the ComfyUI native nodes. ComfyUI-TeaCache is easy to use, simply connect the TeaCache node with the ComfyUI native nodes for seamless usage.
## Updates
- Mar 6 2025: ComfyUI-TeaCache supports Wan2.1:
- It can achieve a 1.5x lossless speedup and a 2x speedup without much visual quality degradation.
- Support Text to Video and Image to Video!
- Jan 17 2025: ComfyUI-TeaCache supports CogVideoX:
- It can achieve a 1.5x lossless speedup and a 2x speedup without much visual quality degradation.
- Support Text to Video and Image to Video!
@@ -40,12 +43,23 @@ Installation via ComfyUI-Manager is preferred. Simply search for ComfyUI-TeaCach
### TeaCache
To use TeaCache node, simply add `TeaCache For Img Gen` or `TeaCache For Vid Gen` node to your workflow after `Load Diffusion Model` node or `Load LoRA` node (if you need LoRA). Generally, TeaCache can achieve a speedup of 1.5x to 3x with acceptable visual quality loss. The following table gives the recommended rel_l1_thresh ​for different models:
| | FLUX | PuLID-FLUX | HunyuanVideo | LTX-Video | CogVideoX |
|:---------------:|:-----------------:|:------------------:|:------------------:|:---------------:|:---------------:|
| rel_l1_thresh | 0.4 | 0.4 | 0.15 | 0.06 | 0.3 |
| speedup | ~2x | ~1.7x | ~2x | ~1.7x | ~2x |
<div align="center">
The demo workflows ([teacache_flux](./examples/teacache_flux.json), [teacache_pulid_flux](./examples/teacache_pulid_flux.json), [teacache_hunyuanvideo](./examples/teacache_hunyuanvideo.json), [teacache_ltx_video](./examples/teacache_ltx_video.json), [teacache_cogvideox](./examples/teacache_cogvideox.json)) are placed in examples folder.
| Models | rel_l1_thresh | speedup |
|:-------------------:|:-----------------:|:------------------:|
| FLUX | 0.4 | ~2x |
| PuLID-FLUX | 0.4 | ~1.7x |
| HunyuanVideo | 0.15 | ~2x |
| LTX-Video | 0.06 | ~1.7x |
| CogVideoX | 0.3 | ~2x |
| Wan2.1-T2V-1.3B | 0.08 | ~1.7x |
| Wan2.1-T2V-14B | 0.1 | ~2.3x |
| Wan2.1-I2V-14B-480P | 0.26 | ~2x |
| Wan2.1-I2V-14B-720P | 0.25 | ~1.7x |
</div>
The demo workflows ([teacache_flux](./examples/teacache_flux.json), [teacache_pulid_flux](./examples/teacache_pulid_flux.json), [teacache_hunyuanvideo](./examples/teacache_hunyuanvideo.json), [teacache_ltx_video](./examples/teacache_ltx_video.json), [teacache_cogvideox](./examples/teacache_cogvideox.json), [teacache_wan2.1_t2v](./examples/teacache_wan2.1_t2v.json) and [teacache_wan2.1_i2v](./examples/teacache_wan2.1_i2v.json)) are placed in examples folder.
### Compile Model
To use Compile Model node, simply add `Compile Model` node to your workflow after `Load Diffusion Model` node or `TeaCache` node. Compile Model uses `torch.compile` to enhance the model performance by compiling model into more efficient intermediate representations (IRs). This compilation process leverages backend compilers to generate optimized code, which can significantly speed up inference. The compilation may take long time when you run the workflow at first, but once it is compiled, inference is extremely fast. The usage is shown below:
@@ -67,5 +81,11 @@ https://github.com/user-attachments/assets/8fce9b48-2243-46f1-b411-80e4a53f6f7d
- <p><strong>CogVideoX</strong></p>
https://github.com/user-attachments/assets/de8f5400-e885-446d-936f-c026a78ba5c2
- <p><strong>Wan2.1-T2V</strong></p>
https://github.com/user-attachments/assets/2c7428b6-f038-43ab-aae1-efc003a55342
- <p><strong>Wan2.1-I2V</strong></p>
https://github.com/user-attachments/assets/ae6991e2-d995-4812-a3ed-dd581bb7701b
## Acknowledgments
Thanks to TeaCache repo owner [ali-vilab/TeaCache: Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model](https://github.com/ali-vilab/TeaCache)
assets/flux_dev_example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

examples/teacache_wan2.1_i2v.json
+812
View File
@@ -0,0 +1,812 @@
{
"last_node_id": 56,
"last_link_id": 114,
"nodes": [
{
"id": 39,
"type": "VAELoader",
"pos": [
866.3932495117188,
499.18597412109375
],
"size": [
306.36004638671875,
58
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "VAE",
"type": "VAE",
"links": [
76,
99
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAELoader"
},
"widgets_values": [
"wan_2.1_vae.safetensors"
]
},
{
"id": 7,
"type": "CLIPTextEncode",
"pos": [
413,
389
],
"size": [
425.27801513671875,
180.6060791015625
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 75
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
98
],
"slot_index": 0
}
],
"title": "CLIP Text Encode (Negative Prompt)",
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
[
false,
true
]
],
"color": "#322",
"bgcolor": "#533"
},
{
"id": 50,
"type": "WanImageToVideo",
"pos": [
673.0507202148438,
627.272705078125
],
"size": [
342.5999755859375,
210
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "positive",
"type": "CONDITIONING",
"link": 97
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 98
},
{
"name": "vae",
"type": "VAE",
"link": 99
},
{
"name": "clip_vision_output",
"type": "CLIP_VISION_OUTPUT",
"shape": 7,
"link": 107
},
{
"name": "start_image",
"type": "IMAGE",
"shape": 7,
"link": 106
}
],
"outputs": [
{
"name": "positive",
"type": "CONDITIONING",
"links": [
101
],
"slot_index": 0
},
{
"name": "negative",
"type": "CONDITIONING",
"links": [
102
],
"slot_index": 1
},
{
"name": "latent",
"type": "LATENT",
"links": [
103
],
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "WanImageToVideo"
},
"widgets_values": [
512,
512,
33,
1
]
},
{
"id": 6,
"type": "CLIPTextEncode",
"pos": [
415,
186
],
"size": [
422.84503173828125,
164.31304931640625
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 74
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
97
],
"slot_index": 0
}
],
"title": "CLIP Text Encode (Positive Prompt)",
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"a cute anime girl with massive fennec ears and a big fluffy tail wearing a maid outfit turning around",
[
false,
true
]
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 52,
"type": "LoadImage",
"pos": [
20,
760
],
"size": [
315,
314
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
106,
109
],
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"flux_dev_example.png",
"image"
]
},
{
"id": 54,
"type": "ModelSamplingSD3",
"pos": [
510,
70
],
"size": [
315,
58
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 114
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
111
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ModelSamplingSD3"
},
"widgets_values": [
8
]
},
{
"id": 38,
"type": "CLIPLoader",
"pos": [
20,
190
],
"size": [
390,
98
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "CLIP",
"type": "CLIP",
"links": [
74,
75
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPLoader"
},
"widgets_values": [
"umt5_xxl_fp8_e4m3fn_scaled.safetensors",
"wan",
"default"
]
},
{
"id": 49,
"type": "CLIPVisionLoader",
"pos": [
20,
640
],
"size": [
315,
58
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "CLIP_VISION",
"type": "CLIP_VISION",
"links": [
94
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPVisionLoader"
},
"widgets_values": [
"clip_vision_h.safetensors"
]
},
{
"id": 3,
"type": "KSampler",
"pos": [
863,
187
],
"size": [
315,
262
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 111
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 101
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 102
},
{
"name": "latent_image",
"type": "LATENT",
"link": 103
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
35
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
987948718394761,
"fixed",
20,
6,
"uni_pc",
"simple",
1
]
},
{
"id": 8,
"type": "VAEDecode",
"pos": [
1210,
190
],
"size": [
210,
46
],
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 35
},
{
"name": "vae",
"type": "VAE",
"link": 76
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
112
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
},
"widgets_values": []
},
{
"id": 55,
"type": "VHS_VideoCombine",
"pos": [
1465.71044921875,
191.89413452148438
],
"size": [
292.3946838378906,
334
],
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 112
},
{
"name": "audio",
"type": "AUDIO",
"shape": 7,
"link": null
},
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"shape": 7,
"link": null
},
{
"name": "vae",
"type": "VAE",
"shape": 7,
"link": null
}
],
"outputs": [
{
"name": "Filenames",
"type": "VHS_FILENAMES",
"links": null
}
],
"properties": {
"Node name for S&R": "VHS_VideoCombine"
},
"widgets_values": {
"frame_rate": 16,
"loop_count": 0,
"filename_prefix": "teacache",
"format": "video/h264-mp4",
"pix_fmt": "yuv420p",
"crf": 19,
"save_metadata": true,
"trim_to_audio": false,
"pingpong": false,
"save_output": true,
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"filename": "teacache_00044.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 16,
"workflow": "teacache_00044.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/teacache_00044.mp4"
},
"muted": false
}
}
},
{
"id": 51,
"type": "CLIPVisionEncode",
"pos": [
366.8671569824219,
640.624267578125
],
"size": [
253.60000610351562,
78
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "clip_vision",
"type": "CLIP_VISION",
"link": 94
},
{
"name": "image",
"type": "IMAGE",
"link": 109
}
],
"outputs": [
{
"name": "CLIP_VISION_OUTPUT",
"type": "CLIP_VISION_OUTPUT",
"links": [
107
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPVisionEncode"
},
"widgets_values": [
"none"
]
},
{
"id": 37,
"type": "UNETLoader",
"pos": [
20,
70
],
"size": [
346.7470703125,
82
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
113
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "UNETLoader"
},
"widgets_values": [
"wan2.1_i2v_480p_14B_bf16.safetensors",
"default"
]
},
{
"id": 56,
"type": "TeaCacheForVidGen",
"pos": [
336.0735168457031,
-95.73665618896484
],
"size": [
315,
82
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 113
}
],
"outputs": [
{
"name": "model",
"type": "MODEL",
"links": [
114
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForVidGen"
},
"widgets_values": [
"wan2.1_i2v_480p_14B",
0.26
]
}
],
"links": [
[
35,
3,
0,
8,
0,
"LATENT"
],
[
74,
38,
0,
6,
0,
"CLIP"
],
[
75,
38,
0,
7,
0,
"CLIP"
],
[
76,
39,
0,
8,
1,
"VAE"
],
[
94,
49,
0,
51,
0,
"CLIP_VISION"
],
[
97,
6,
0,
50,
0,
"CONDITIONING"
],
[
98,
7,
0,
50,
1,
"CONDITIONING"
],
[
99,
39,
0,
50,
2,
"VAE"
],
[
101,
50,
0,
3,
1,
"CONDITIONING"
],
[
102,
50,
1,
3,
2,
"CONDITIONING"
],
[
103,
50,
2,
3,
3,
"LATENT"
],
[
106,
52,
0,
50,
4,
"IMAGE"
],
[
107,
51,
0,
50,
3,
"CLIP_VISION_OUTPUT"
],
[
109,
52,
0,
51,
1,
"IMAGE"
],
[
111,
54,
0,
3,
0,
"MODEL"
],
[
112,
8,
0,
55,
0,
"IMAGE"
],
[
113,
37,
0,
56,
0,
"MODEL"
],
[
114,
56,
0,
54,
0,
"MODEL"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.7627768444385795,
"offset": [
336.0382964059074,
205.80199013908017
]
},
"node_versions": {
"comfy-core": "0.3.22",
"ComfyUI-VideoHelperSuite": "124c913ccdd8a585734ea758c35fa1bab8499c99",
"ComfyUI-TeaCache": "062efbab2ec1296ba721d9985c5cd80c08da0e42"
},
"ue_links": [],
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
examples/teacache_wan2.1_t2v.json
+596
View File
@@ -0,0 +1,596 @@
{
"last_node_id": 51,
"last_link_id": 98,
"nodes": [
{
"id": 39,
"type": "VAELoader",
"pos": [
866.3932495117188,
499.18597412109375
],
"size": [
306.36004638671875,
58
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "VAE",
"type": "VAE",
"links": [
76
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAELoader"
},
"widgets_values": [
"wan_2.1_vae.safetensors"
]
},
{
"id": 7,
"type": "CLIPTextEncode",
"pos": [
413,
389
],
"size": [
425.27801513671875,
180.6060791015625
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 75
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
52
],
"slot_index": 0
}
],
"title": "CLIP Text Encode (Negative Prompt)",
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
[
false,
true
]
],
"color": "#322",
"bgcolor": "#533"
},
{
"id": 38,
"type": "CLIPLoader",
"pos": [
12.94982624053955,
184.6981658935547
],
"size": [
390,
98
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "CLIP",
"type": "CLIP",
"links": [
74,
75
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPLoader"
},
"widgets_values": [
"umt5_xxl_fp8_e4m3fn_scaled.safetensors",
"wan",
"default"
]
},
{
"id": 40,
"type": "EmptyHunyuanLatentVideo",
"pos": [
520,
620
],
"size": [
315,
130
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
91
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "EmptyHunyuanLatentVideo"
},
"widgets_values": [
832,
480,
33,
1
]
},
{
"id": 6,
"type": "CLIPTextEncode",
"pos": [
415,
186
],
"size": [
422.84503173828125,
164.31304931640625
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 74
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
46
],
"slot_index": 0
}
],
"title": "CLIP Text Encode (Positive Prompt)",
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"a fox moving quickly in a beautiful winter scenery nature trees mountains daytime tracking camera",
[
false,
true
]
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 8,
"type": "VAEDecode",
"pos": [
1210,
190
],
"size": [
210,
46
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 35
},
{
"name": "vae",
"type": "VAE",
"link": 76
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
96
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
},
"widgets_values": []
},
{
"id": 49,
"type": "VHS_VideoCombine",
"pos": [
1476.3587646484375,
192.16001892089844
],
"size": [
292.3946838378906,
334
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 96
},
{
"name": "audio",
"type": "AUDIO",
"shape": 7,
"link": null
},
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"shape": 7,
"link": null
},
{
"name": "vae",
"type": "VAE",
"shape": 7,
"link": null
}
],
"outputs": [
{
"name": "Filenames",
"type": "VHS_FILENAMES",
"links": null
}
],
"properties": {
"Node name for S&R": "VHS_VideoCombine"
},
"widgets_values": {
"frame_rate": 16,
"loop_count": 0,
"filename_prefix": "teacache",
"format": "video/h264-mp4",
"pix_fmt": "yuv420p",
"crf": 19,
"save_metadata": true,
"trim_to_audio": false,
"pingpong": false,
"save_output": true,
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"filename": "teacache_00029.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 16,
"workflow": "teacache_00029.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/teacache_00029.mp4"
},
"muted": false
}
}
},
{
"id": 37,
"type": "UNETLoader",
"pos": [
20,
40
],
"size": [
346.7470703125,
82
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
97
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "UNETLoader"
},
"widgets_values": [
"wan2.1_t2v_1.3B_fp16.safetensors",
"default"
]
},
{
"id": 3,
"type": "KSampler",
"pos": [
863,
187
],
"size": [
315,
262
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 95
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 46
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 52
},
{
"name": "latent_image",
"type": "LATENT",
"link": 91
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
35
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
82628696717253,
"fixed",
30,
6,
"uni_pc",
"simple",
1
]
},
{
"id": 48,
"type": "ModelSamplingSD3",
"pos": [
440,
50
],
"size": [
210,
58
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 98
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
95
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ModelSamplingSD3"
},
"widgets_values": [
8
]
},
{
"id": 51,
"type": "TeaCacheForVidGen",
"pos": [
295.4561767578125,
-147.876708984375
],
"size": [
315,
82
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 97
}
],
"outputs": [
{
"name": "model",
"type": "MODEL",
"links": [
98
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForVidGen"
},
"widgets_values": [
"wan2.1_t2v_1.3B",
0.08
]
}
],
"links": [
[
35,
3,
0,
8,
0,
"LATENT"
],
[
46,
6,
0,
3,
1,
"CONDITIONING"
],
[
52,
7,
0,
3,
2,
"CONDITIONING"
],
[
74,
38,
0,
6,
0,
"CLIP"
],
[
75,
38,
0,
7,
0,
"CLIP"
],
[
76,
39,
0,
8,
1,
"VAE"
],
[
91,
40,
0,
3,
3,
"LATENT"
],
[
95,
48,
0,
3,
0,
"MODEL"
],
[
96,
8,
0,
49,
0,
"IMAGE"
],
[
97,
37,
0,
51,
0,
"MODEL"
],
[
98,
51,
0,
48,
0,
"MODEL"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.6934334949441631,
"offset": [
509.8357423302538,
381.2577148997867
]
},
"node_versions": {
"comfy-core": "0.3.22",
"ComfyUI-VideoHelperSuite": "124c913ccdd8a585734ea758c35fa1bab8499c99",
"ComfyUI-TeaCache": "062efbab2ec1296ba721d9985c5cd80c08da0e42"
},
"ue_links": [],
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
nodes.py
+122 -18
View File
@@ -1,19 +1,32 @@
import math
import torch
import comfy
from torch import Tensor
from einops import repeat
from unittest.mock import patch
from comfy.ldm.flux.layers import timestep_embedding
from comfy.ldm.lightricks.model import precompute_freqs_cis
from comfy.ldm.common_dit import rms_norm
from comfy.ldm.wan.model import sinusoidal_embedding_1d
SUPPORTED_MODELS_COEFFICIENTS = {
"flux": [4.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 2.64230861e-01],
"hunyuan_video": [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02],
"ltxv": [2.14700694e+01, -1.28016453e+01, 2.31279151e+00, 7.92487521e-01, 9.69274326e-03],
"wan2.1_t2v_1.3B": [2.39676752e+03, -1.31110545e+03, 2.01331979e+02, -8.29855975e+00, 1.37887774e-01],
"wan2.1_t2v_14B": [-5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404],
"wan2.1_i2v_480p_14B": [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01],
"wan2.1_i2v_720p_14B": [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683]
}
def poly1d(coefficients, x):
result = torch.zeros_like(x)
for i, coeff in enumerate(coefficients):
result += coeff * (x ** (len(coefficients) - 1 - i))
return result
return result.abs()
def teacache_flux_forward(
self,
@@ -30,6 +43,7 @@ def teacache_flux_forward(
) -> Tensor:
patches_replace = transformer_options.get("patches_replace", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
coefficients = transformer_options.get("coefficients")
if img.ndim != 3 or txt.ndim != 3:
raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -63,7 +77,6 @@ def teacache_flux_forward(
self.accumulated_rel_l1_distance = 0
else:
try:
coefficients = [4.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 2.64230861e-01]
self.accumulated_rel_l1_distance += poly1d(coefficients, ((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()))
if self.accumulated_rel_l1_distance < rel_l1_thresh:
should_calc = False
@@ -185,6 +198,7 @@ def teacache_hunyuanvideo_forward(
) -> Tensor:
patches_replace = transformer_options.get("patches_replace", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
coefficients = transformer_options.get("coefficients")
initial_shape = list(img.shape)
# running on sequences img
@@ -228,7 +242,6 @@ def teacache_hunyuanvideo_forward(
self.accumulated_rel_l1_distance = 0
else:
try:
coefficients = [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02]
self.accumulated_rel_l1_distance += poly1d(coefficients, ((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()))
if self.accumulated_rel_l1_distance < rel_l1_thresh:
should_calc = False
@@ -313,6 +326,7 @@ def teacache_ltxvmodel_forward(
):
patches_replace = transformer_options.get("patches_replace", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
coefficients = transformer_options.get("coefficients")
indices_grid = self.patchifier.get_grid(
orig_num_frames=x.shape[2],
@@ -394,7 +408,6 @@ def teacache_ltxvmodel_forward(
self.accumulated_rel_l1_distance = 0
else:
try:
coefficients = [2.14700694e+01, -1.28016453e+01, 2.31279151e+00, 7.92487521e-01, 9.69274326e-03]
self.accumulated_rel_l1_distance += poly1d(coefficients, ((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()))
if self.accumulated_rel_l1_distance < rel_l1_thresh:
should_calc = False
@@ -456,6 +469,92 @@ def teacache_ltxvmodel_forward(
# print("res", x)
return x
def teacache_wanmodel_forward(self, x, timestep, context, clip_fea=None, transformer_options={}, **kwargs):
bs, c, t, h, w = x.shape
x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
patch_size = self.patch_size
t_len = ((t + (patch_size[0] // 2)) // patch_size[0])
h_len = ((h + (patch_size[1] // 2)) // patch_size[1])
w_len = ((w + (patch_size[2] // 2)) // patch_size[2])
img_ids = torch.zeros((t_len, h_len, w_len, 3), device=x.device, dtype=x.dtype)
img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(0, t_len - 1, steps=t_len, device=x.device, dtype=x.dtype).reshape(-1, 1, 1)
img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).reshape(1, -1, 1)
img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).reshape(1, 1, -1)
img_ids = repeat(img_ids, "t h w c -> b (t h w) c", b=bs)
freqs = self.rope_embedder(img_ids).movedim(1, 2)
return self.forward_orig(x, timestep, context, clip_fea, freqs, transformer_options)[:, :, :t, :h, :w]
def teacache_wanmodel_forward_orig(
self,
x,
t,
context,
clip_fea=None,
freqs=None,
transformer_options={},
):
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
coefficients = transformer_options.get("coefficients")
# embeddings
x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:]
x = x.flatten(2).transpose(1, 2)
# time embeddings
e = self.time_embedding(
sinusoidal_embedding_1d(self.freq_dim, t).to(dtype=x[0].dtype))
e0 = self.time_projection(e).unflatten(1, (6, self.dim))
# context
context = self.text_embedding(context)
if clip_fea is not None and self.img_emb is not None:
context_clip = self.img_emb(clip_fea) # bs x 257 x dim
context = torch.concat([context_clip, context], dim=1)
# arguments
kwargs = dict(
e=e0,
freqs=freqs,
context=context)
# enable teacache
modulated_inp = e
if not hasattr(self, 'accumulated_rel_l1_distance'):
should_calc = True
self.accumulated_rel_l1_distance = 0
else:
try:
self.accumulated_rel_l1_distance += poly1d(coefficients, ((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()))
if self.accumulated_rel_l1_distance < rel_l1_thresh:
should_calc = False
else:
should_calc = True
self.accumulated_rel_l1_distance = 0
except:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.previous_modulated_input = modulated_inp
if not should_calc:
x += self.previous_residual
else:
ori_x = x.clone()
for block in self.blocks:
x = block(x, **kwargs)
self.previous_residual = x - ori_x
# head
x = self.head(x, e)
# unpatchify
x = self.unpatchify(x, grid_sizes)
return x
class TeaCacheForImgGen:
@classmethod
def INPUT_TYPES(s):
@@ -509,7 +608,7 @@ class TeaCacheForVidGen:
return {
"required": {
"model": ("MODEL", {"tooltip": "The video diffusion model the TeaCache will be applied to."}),
"model_type": (["hunyuan_video", "ltxv"],),
"model_type": (["hunyuan_video", "ltxv", "wan2.1_t2v_1.3B", "wan2.1_t2v_14B", "wan2.1_i2v_480p_14B", "wan2.1_i2v_720p_14B"],),
"rel_l1_thresh": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "How strongly to cache the output of diffusion model. This value must be non-negative."})
}
}
@@ -528,20 +627,25 @@ class TeaCacheForVidGen:
if 'transformer_options' not in new_model.model_options:
new_model.model_options['transformer_options'] = {}
new_model.model_options["transformer_options"]["rel_l1_thresh"] = rel_l1_thresh
new_model.model_options["transformer_options"]["coefficients"] = SUPPORTED_MODELS_COEFFICIENTS[model_type]
diffusion_model = new_model.get_model_object("diffusion_model")
if model_type == "hunyuan_video":
forward_name = "forward_orig"
replaced_forward_fn = teacache_hunyuanvideo_forward.__get__(
diffusion_model,
diffusion_model.__class__
)
elif model_type == "ltxv":
forward_name = "forward"
replaced_forward_fn = teacache_ltxvmodel_forward.__get__(
diffusion_model,
diffusion_model.__class__
)
if "hunyuan_video" in model_type:
context = patch.multiple(
diffusion_model,
forward_orig=teacache_hunyuanvideo_forward.__get__(diffusion_model, diffusion_model.__class__)
)
elif "ltxv" in model_type:
context = patch.multiple(
diffusion_model,
forward=teacache_ltxvmodel_forward.__get__(diffusion_model, diffusion_model.__class__)
)
elif "wan2.1" in model_type:
context = patch.multiple(
diffusion_model,
forward=teacache_wanmodel_forward.__get__(diffusion_model, diffusion_model.__class__),
forward_orig=teacache_wanmodel_forward_orig.__get__(diffusion_model, diffusion_model.__class__)
)
else:
raise ValueError(f"Unknown type {model_type}")
@@ -549,7 +653,7 @@ class TeaCacheForVidGen:
input = kwargs["input"]
timestep = kwargs["timestep"]
c = kwargs["c"]
with patch.object(diffusion_model, forward_name, replaced_forward_fn):
with context:
return model_function(input, timestep, **c)
new_model.set_model_unet_function_wrapper(unet_wrapper_function)
pyproject.toml
+2 -2
View File
@@ -1,7 +1,7 @@
[project]
name = "teacache"
name = "ComfyUI-TeaCache"
description = "Unofficial implementation of [ali-vilab/TeaCache](https://github.com/ali-vilab/TeaCache) for ComfyUI"
version = "1.2.4"
version = "1.3.0"
license = {file = "LICENSE"}
[project.urls]