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

Add max_skip_steps option; Merge TeaCache nodes; Update TeaCache for HunyuanVideo and LTX-Video

Browse Source
This commit is contained in:
YunjieYu
2025-03-10 19:23:38 +08:00
parent b3429ef3de
commit 1a11412dea
10 changed files with 592 additions and 543 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+19 -13
View File
@@ -6,6 +6,10 @@ 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 10 2025: ComfyUI-TeaCache adds max_skip_steps option and has made some changes for ease of use:
- Add max_skip_steps option to enjoy a good trade-off between quality and speed for Wan2.1 models. The best settings are shown in the usage section.
- Merge TeaCache For Img Gen and TeaCache For Vid Gen nodes into a single TeaCache node.
- Update TeaCache for HunyuanVideo and LTX-Video to match the offical ComfyUI update.
- 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!
@@ -41,25 +45,27 @@ Installation via ComfyUI-Manager is preferred. Simply search for ComfyUI-TeaCach
## Usage
### 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:
To use TeaCache node, simply add `TeaCache` 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 and max_skip_steps ​for different models:
<div align="center">
| 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 |
| Models | rel_l1_thresh | max_skip_steps | speedup |
|:-------------------:|:-----------------:|:-----------------:|:-----------------:|
| FLUX | 0.4 | 3 | ~2x |
| PuLID-FLUX | 0.4 | 3 | ~1.7x |
| HunyuanVideo | 0.15 | 3 | ~1.9x |
| LTX-Video | 0.06 | 3 | ~1.7x |
| CogVideoX | 0.3 | 3 | ~2x |
| Wan2.1-T2V-1.3B | 0.08 | 3 | ~1.6x |
| Wan2.1-T2V-14B | 0.2 | 3 | ~1.9x |
| Wan2.1-I2V-14B-480P | 0.26 | 3 | ~1.9x |
| Wan2.1-I2V-14B-720P | 0.25 | 3 | ~1.9x |
</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.
If the video after applying TeaCache has multicolored block artifacts, a lower range of motion or even the still frames, please reduce max_skip_steps to 1 or 2.
The demo workflows ([flux](./examples/flux.json), [pulid_flux](./examples/pulid_flux.json), [hunyuanvideo](./examples/hunyuanvideo.json), [ltx_video](./examples/ltx_video.json), [cogvideox](./examples/cogvideox.json), [wan2.1_t2v](./examples/wan2.1_t2v.json) and [wan2.1_i2v](./examples/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:
examples/teacache_cogvideox.json → examples/cogvideox.json
View File
examples/teacache_flux.json → examples/flux.json
+129 -103
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 40,
"last_link_id": 122,
"last_node_id": 41,
"last_link_id": 125,
"nodes": [
{
"id": 11,
@@ -11,7 +11,7 @@
],
"size": [
315,
106
122
],
"flags": {},
"order": 0,
@@ -21,11 +21,11 @@
{
"name": "CLIP",
"type": "CLIP",
"shape": 3,
"links": [
10
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -34,7 +34,8 @@
"widgets_values": [
"t5xxl_fp16.safetensors",
"clip_l.safetensors",
"flux"
"flux",
"default"
]
},
{
@@ -62,11 +63,11 @@
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"shape": 3,
"links": [
42
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -122,11 +123,11 @@
{
"name": "VAE",
"type": "VAE",
"shape": 3,
"links": [
12
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -180,14 +181,14 @@
{
"name": "INT",
"type": "INT",
"widget": {
"name": "width"
},
"links": [
112,
115
],
"slot_index": 0,
"widget": {
"name": "width"
}
"slot_index": 0
}
],
"title": "width",
@@ -220,14 +221,14 @@
{
"name": "INT",
"type": "INT",
"widget": {
"name": "height"
},
"links": [
113,
114
],
"slot_index": 0,
"widget": {
"name": "height"
}
"slot_index": 0
}
],
"title": "height",
@@ -250,7 +251,7 @@
],
"size": [
315,
130
170
],
"flags": {},
"order": 13,
@@ -259,38 +260,46 @@
{
"name": "model",
"type": "MODEL",
"link": 121,
"link": 125,
"slot_index": 0
},
{
"name": "width",
"type": "INT",
"link": 115,
"slot_index": 1,
"pos": [
10,
84
],
"widget": {
"name": "width"
}
},
"link": 115,
"slot_index": 1
},
{
"name": "height",
"type": "INT",
"link": 114,
"slot_index": 2,
"pos": [
10,
108
],
"widget": {
"name": "height"
}
},
"link": 114,
"slot_index": 2
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"shape": 3,
"links": [
54,
55
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -335,11 +344,11 @@
{
"name": "GUIDER",
"type": "GUIDER",
"shape": 3,
"links": [
30
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -397,17 +406,17 @@
{
"name": "output",
"type": "LATENT",
"shape": 3,
"links": [
24
],
"slot_index": 0,
"shape": 3
"slot_index": 0
},
{
"name": "denoised_output",
"type": "LATENT",
"links": null,
"shape": 3
"shape": 3,
"links": null
}
],
"properties": {
@@ -434,10 +443,10 @@
{
"name": "SAMPLER",
"type": "SAMPLER",
"shape": 3,
"links": [
19
],
"shape": 3
]
}
],
"properties": {
@@ -456,7 +465,7 @@
],
"size": [
315,
106
126
],
"flags": {},
"order": 10,
@@ -465,29 +474,37 @@
{
"name": "width",
"type": "INT",
"link": 112,
"pos": [
10,
36
],
"widget": {
"name": "width"
}
},
"link": 112
},
{
"name": "height",
"type": "INT",
"link": 113,
"pos": [
10,
60
],
"widget": {
"name": "height"
}
},
"link": 113
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"shape": 3,
"links": [
116
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -608,10 +625,10 @@
{
"name": "SIGMAS",
"type": "SIGMAS",
"shape": 3,
"links": [
20
],
"shape": 3
]
}
],
"properties": {
@@ -652,41 +669,6 @@
"ComfyUI"
]
},
{
"id": 25,
"type": "RandomNoise",
"pos": [
480,
768
],
"size": [
315,
82
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "NOISE",
"type": "NOISE",
"links": [
37
],
"shape": 3
}
],
"properties": {
"Node name for S&R": "RandomNoise"
},
"widgets_values": [
725991276578445,
"randomize"
],
"color": "#2a363b",
"bgcolor": "#3f5159"
},
{
"id": 12,
"type": "UNETLoader",
@@ -699,18 +681,18 @@
82
],
"flags": {},
"order": 8,
"order": 7,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"shape": 3,
"links": [
122
124
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -724,11 +706,11 @@
"bgcolor": "#335"
},
{
"id": 40,
"type": "TeaCacheForImgGen",
"id": 41,
"type": "TeaCache",
"pos": [
77.93648529052734,
1151.0115966796875
64.00762939453125,
1143.05322265625
],
"size": [
315,
@@ -741,27 +723,62 @@
{
"name": "model",
"type": "MODEL",
"link": 122
"link": 124
}
],
"outputs": [
{
"name": "MODEL",
"name": "model",
"type": "MODEL",
"links": [
121
125
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForImgGen"
"Node name for S&R": "TeaCache"
},
"widgets_values": [
"flux",
0.4,
30
3
]
},
{
"id": 25,
"type": "RandomNoise",
"pos": [
480,
768
],
"size": [
315,
82
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "NOISE",
"type": "NOISE",
"shape": 3,
"links": [
37
]
}
],
"properties": {
"Node name for S&R": "RandomNoise"
},
"widgets_values": [
725991276578445,
"fixed"
],
"color": "#2a363b",
"bgcolor": "#3f5159"
}
],
"links": [
@@ -902,18 +919,18 @@
"LATENT"
],
[
121,
40,
124,
12,
0,
30,
41,
0,
"MODEL"
],
[
122,
12,
125,
41,
0,
40,
30,
0,
"MODEL"
]
@@ -922,13 +939,22 @@
"config": {},
"extra": {
"ds": {
"scale": 0.5644739300537778,
"scale": 0.6830134553650713,
"offset": [
789.503519235993,
-441.9800029782946
585.0708074418511,
-310.9613966916353
]
},
"groupNodes": {}
"groupNodes": {},
"node_versions": {
"comfy-core": "0.3.26",
"ComfyUI-TeaCache": "b3429ef3dea426d2f167e348b44cd2f5a3674e7d"
},
"ue_links": [],
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
examples/teacache_hunyuanvideo.json → examples/hunyuanvideo.json
+108 -94
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 81,
"last_link_id": 224,
"last_node_id": 82,
"last_link_id": 229,
"nodes": [
{
"id": 16,
@@ -21,10 +21,10 @@
{
"name": "SAMPLER",
"type": "SAMPLER",
"shape": 3,
"links": [
19
],
"shape": 3
]
}
],
"properties": {
@@ -59,11 +59,11 @@
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"shape": 3,
"links": [
129
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -107,11 +107,11 @@
{
"name": "GUIDER",
"type": "GUIDER",
"shape": 3,
"links": [
30
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -119,40 +119,6 @@
},
"widgets_values": []
},
{
"id": 10,
"type": "VAELoader",
"pos": [
0,
420
],
"size": [
350,
60
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "VAE",
"type": "VAE",
"links": [
206,
211
],
"slot_index": 0,
"shape": 3
}
],
"properties": {
"Node name for S&R": "VAELoader"
},
"widgets_values": [
"hunyuan_video_vae_bf16.safetensors"
]
},
{
"id": 11,
"type": "DualCLIPLoader",
@@ -162,21 +128,21 @@
],
"size": [
350,
106
122
],
"flags": {},
"order": 2,
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "CLIP",
"type": "CLIP",
"shape": 3,
"links": [
205
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -185,7 +151,8 @@
"widgets_values": [
"clip_l.safetensors",
"llava_llama3_fp8_scaled.safetensors",
"hunyuan_video"
"hunyuan_video",
"default"
]
},
{
@@ -239,7 +206,7 @@
170
],
"flags": {},
"order": 3,
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [],
@@ -300,18 +267,18 @@
{
"name": "output",
"type": "LATENT",
"shape": 3,
"links": [
181,
210
],
"slot_index": 0,
"shape": 3
"slot_index": 0
},
{
"name": "denoised_output",
"type": "LATENT",
"links": null,
"shape": 3
"shape": 3,
"links": null
}
],
"properties": {
@@ -373,7 +340,7 @@
110
],
"flags": {},
"order": 4,
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [],
@@ -402,7 +369,7 @@
{
"name": "model",
"type": "MODEL",
"link": 224,
"link": 229,
"slot_index": 0
}
],
@@ -410,10 +377,10 @@
{
"name": "SIGMAS",
"type": "SIGMAS",
"shape": 3,
"links": [
20
],
"shape": 3
]
}
],
"properties": {
@@ -437,7 +404,7 @@
130
],
"flags": {},
"order": 5,
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
@@ -469,7 +436,7 @@
],
"size": [
292.3946838378906,
596.3946533203125
620.3946533203125
],
"flags": {},
"order": 17,
@@ -483,20 +450,20 @@
{
"name": "audio",
"type": "AUDIO",
"link": null,
"shape": 7
"shape": 7,
"link": null
},
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"link": null,
"shape": 7
"shape": 7,
"link": null
},
{
"name": "vae",
"type": "VAE",
"link": null,
"shape": 7
"shape": 7,
"link": null
}
],
"outputs": [
@@ -517,17 +484,20 @@
"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_00011.mp4",
"filename": "teacache_00082.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 24
"frame_rate": 24,
"workflow": "teacache_00082.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/teacache_00082.mp4"
},
"muted": false
}
@@ -551,7 +521,7 @@
{
"name": "model",
"type": "MODEL",
"link": 223
"link": 228
}
],
"outputs": [
@@ -629,17 +599,17 @@
82
],
"flags": {},
"order": 6,
"order": 5,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "NOISE",
"type": "NOISE",
"shape": 3,
"links": [
37
],
"shape": 3
]
}
],
"properties": {
@@ -664,18 +634,18 @@
82
],
"flags": {},
"order": 7,
"order": 6,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"shape": 3,
"links": [
222
227
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -689,11 +659,11 @@
"bgcolor": "#335"
},
{
"id": 81,
"type": "TeaCacheForVidGen",
"id": 82,
"type": "TeaCache",
"pos": [
20.529560089111328,
102.00926208496094
14.720276832580566,
89.97886657714844
],
"size": [
315,
@@ -706,27 +676,61 @@
{
"name": "model",
"type": "MODEL",
"link": 222
"link": 227
}
],
"outputs": [
{
"name": "MODEL",
"name": "model",
"type": "MODEL",
"links": [
223,
224
228,
229
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForVidGen"
"Node name for S&R": "TeaCache"
},
"widgets_values": [
"hunyuan_video",
0.15,
25
3
]
},
{
"id": 10,
"type": "VAELoader",
"pos": [
-2.535106658935547,
466.2682800292969
],
"size": [
350,
60
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "VAE",
"type": "VAE",
"shape": 3,
"links": [
206,
211
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAELoader"
},
"widgets_values": [
"hunyuan_video_vae_bf16.safetensors"
]
}
],
@@ -844,24 +848,24 @@
"IMAGE"
],
[
222,
227,
12,
0,
81,
82,
0,
"MODEL"
],
[
223,
81,
228,
82,
0,
67,
0,
"MODEL"
],
[
224,
81,
229,
82,
0,
17,
0,
@@ -872,16 +876,26 @@
"config": {},
"extra": {
"ds": {
"scale": 0.6830134553650712,
"scale": 0.9090909090909091,
"offset": [
553.6828346433346,
75.84624922798709
424.45821237860434,
139.84374290591396
]
},
"groupNodes": {},
"workspace_info": {
"id": "hSBfPb8KVVUoyCyWT3oHM"
}
},
"node_versions": {
"comfy-core": "0.3.26",
"ComfyUI-VideoHelperSuite": "124c913ccdd8a585734ea758c35fa1bab8499c99",
"ComfyUI-TeaCache": "b3429ef3dea426d2f167e348b44cd2f5a3674e7d"
},
"ue_links": [],
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
examples/teacache_ltx_video.json → examples/ltx_video.json
+48 -34
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 89,
"last_link_id": 188,
"last_node_id": 90,
"last_link_id": 191,
"nodes": [
{
"id": 71,
@@ -20,8 +20,8 @@
{
"name": "latent",
"type": "LATENT",
"link": 168,
"shape": 7
"shape": 7,
"link": 168
}
],
"outputs": [
@@ -279,7 +279,7 @@
],
"size": [
315,
82
98
],
"flags": {},
"order": 3,
@@ -301,7 +301,8 @@
},
"widgets_values": [
"t5xxl_fp16.safetensors",
"ltxv"
"ltxv",
"default"
]
},
{
@@ -354,7 +355,7 @@
],
"size": [
312.7515869140625,
616.7515869140625
640.7515869140625
],
"flags": {},
"order": 12,
@@ -368,20 +369,20 @@
{
"name": "audio",
"type": "AUDIO",
"link": null,
"shape": 7
"shape": 7,
"link": null
},
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"link": null,
"shape": 7
"shape": 7,
"link": null
},
{
"name": "vae",
"type": "VAE",
"link": null,
"shape": 7
"shape": 7,
"link": null
}
],
"outputs": [
@@ -402,17 +403,20 @@
"pix_fmt": "yuv420p",
"crf": 19,
"save_metadata": true,
"trim_to_audio": false,
"pingpong": false,
"save_output": true,
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"filename": "ltxv_00032.mp4",
"filename": "ltxv_00076.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 24
"frame_rate": 24,
"workflow": "ltxv_00076.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/ltxv_00076.mp4"
},
"muted": false
}
@@ -436,7 +440,7 @@
{
"name": "model",
"type": "MODEL",
"link": 188
"link": 191
},
{
"name": "positive",
@@ -509,7 +513,7 @@
"name": "MODEL",
"type": "MODEL",
"links": [
187
190
],
"slot_index": 0
},
@@ -535,11 +539,11 @@
]
},
{
"id": 89,
"type": "TeaCacheForVidGen",
"id": 90,
"type": "TeaCache",
"pos": [
862.1427612304688,
-115.31109619140625
865.708251953125,
-110.11214447021484
],
"size": [
315,
@@ -552,26 +556,26 @@
{
"name": "model",
"type": "MODEL",
"link": 187
"link": 190
}
],
"outputs": [
{
"name": "MODEL",
"name": "model",
"type": "MODEL",
"links": [
188
191
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForVidGen"
"Node name for S&R": "TeaCache"
},
"widgets_values": [
"ltxv",
0.06,
30
3
]
}
],
@@ -681,16 +685,16 @@
"IMAGE"
],
[
187,
190,
44,
0,
89,
90,
0,
"MODEL"
],
[
188,
89,
191,
90,
0,
72,
0,
@@ -701,12 +705,22 @@
"config": {},
"extra": {
"ds": {
"scale": 0.6830134553650712,
"scale": 0.7513148009015777,
"offset": [
67.19056247475572,
283.72050363030013
91.73110553274871,
304.2520181935422
]
}
},
"node_versions": {
"comfy-core": "0.3.26",
"ComfyUI-VideoHelperSuite": "124c913ccdd8a585734ea758c35fa1bab8499c99",
"ComfyUI-TeaCache": "b3429ef3dea426d2f167e348b44cd2f5a3674e7d"
},
"ue_links": [],
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
examples/teacache_pulid_flux.json → examples/pulid_flux.json
+76 -64
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 72,
"last_link_id": 152,
"last_node_id": 73,
"last_link_id": 155,
"nodes": [
{
"id": 26,
@@ -29,11 +29,11 @@
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"shape": 3,
"links": [
107
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -106,11 +106,11 @@
{
"name": "LATENT",
"type": "LATENT",
"shape": 3,
"links": [
86
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -143,11 +143,11 @@
{
"name": "SAMPLER",
"type": "SAMPLER",
"shape": 3,
"links": [
85
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -176,17 +176,17 @@
{
"name": "IMAGE",
"type": "IMAGE",
"shape": 3,
"links": [
126
],
"slot_index": 0,
"shape": 3
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
"shape": 3,
"links": null
}
],
"properties": {
@@ -216,11 +216,11 @@
{
"name": "FACEANALYSIS",
"type": "FACEANALYSIS",
"shape": 3,
"links": [
124
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -249,11 +249,11 @@
{
"name": "EVA_CLIP",
"type": "EVA_CLIP",
"shape": 3,
"links": [
123
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -280,11 +280,11 @@
{
"name": "PULIDFLUX",
"type": "PULIDFLUX",
"shape": 3,
"links": [
125
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -339,17 +339,17 @@
{
"name": "output",
"type": "LATENT",
"shape": 3,
"links": [
87
],
"slot_index": 0,
"shape": 3
"slot_index": 0
},
{
"name": "denoised_output",
"type": "LATENT",
"links": null,
"shape": 3
"shape": 3,
"links": null
}
],
"properties": {
@@ -387,11 +387,11 @@
{
"name": "IMAGE",
"type": "IMAGE",
"shape": 3,
"links": [
89
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -435,7 +435,7 @@
],
"size": [
315,
106
122
],
"flags": {},
"order": 6,
@@ -445,11 +445,11 @@
{
"name": "CLIP",
"type": "CLIP",
"shape": 3,
"links": [
132
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -458,7 +458,8 @@
"widgets_values": [
"t5xxl_fp16.safetensors",
"clip_l.safetensors",
"flux"
"flux",
"default"
]
},
{
@@ -480,11 +481,11 @@
{
"name": "VAE",
"type": "VAE",
"shape": 3,
"links": [
88
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -513,11 +514,11 @@
{
"name": "NOISE",
"type": "NOISE",
"shape": 3,
"links": [
84
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -558,11 +559,11 @@
{
"name": "SIGMAS",
"type": "SIGMAS",
"shape": 3,
"links": [
93
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -593,12 +594,12 @@
{
"name": "MODEL",
"type": "MODEL",
"shape": 3,
"links": [
141,
147
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -627,7 +628,7 @@
{
"name": "model",
"type": "MODEL",
"link": 152
"link": 155
},
{
"name": "conditioning",
@@ -639,11 +640,11 @@
{
"name": "GUIDER",
"type": "GUIDER",
"shape": 3,
"links": [
83
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -694,25 +695,25 @@
{
"name": "attn_mask",
"type": "MASK",
"link": null,
"shape": 7
"shape": 7,
"link": null
},
{
"name": "prior_image",
"type": "IMAGE",
"link": null,
"shape": 7
"shape": 7,
"link": null
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"shape": 3,
"links": [
148
154
],
"slot_index": 0,
"shape": 3
"slot_index": 0
}
],
"properties": {
@@ -730,15 +731,15 @@
]
},
{
"id": 71,
"type": "TeaCacheForImgGen",
"id": 73,
"type": "TeaCache",
"pos": [
1228.98779296875,
504.7470397949219
1214.045166015625,
477.5345153808594
],
"size": [
315,
82
106
],
"flags": {},
"order": 14,
@@ -747,25 +748,26 @@
{
"name": "model",
"type": "MODEL",
"link": 148
"link": 154
}
],
"outputs": [
{
"name": "MODEL",
"name": "model",
"type": "MODEL",
"links": [
152
155
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForImgGen"
"Node name for S&R": "TeaCache"
},
"widgets_values": [
"flux",
0.4
0.4,
3
]
}
],
@@ -907,16 +909,16 @@
"MODEL"
],
[
148,
154,
62,
0,
71,
73,
0,
"MODEL"
],
[
152,
71,
155,
73,
0,
47,
0,
@@ -927,12 +929,22 @@
"config": {},
"extra": {
"ds": {
"scale": 0.620921323059155,
"scale": 0.6209213230591561,
"offset": [
52.21858017949011,
279.11370195488706
71.23485125658658,
174.91854810034815
]
}
},
"node_versions": {
"comfy-core": "0.3.26",
"ComfyUI-PuLID-Flux-Enhanced": "04e1b52320f1f14383afe18959349703623c5b88",
"ComfyUI-TeaCache": "b3429ef3dea426d2f167e348b44cd2f5a3674e7d"
},
"ue_links": [],
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
examples/teacache_wan2.1_i2v.json → examples/wan2.1_i2v.json
+26 -25
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 56,
"last_link_id": 114,
"last_node_id": 57,
"last_link_id": 117,
"nodes": [
{
"id": 39,
@@ -262,7 +262,7 @@
{
"name": "model",
"type": "MODEL",
"link": 114
"link": 117
}
],
"outputs": [
@@ -513,13 +513,13 @@
"hidden": false,
"paused": false,
"params": {
"filename": "teacache_00044.mp4",
"filename": "teacache_00092.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 16,
"workflow": "teacache_00044.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/teacache_00044.mp4"
"workflow": "teacache_00092.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/teacache_00092.mp4"
},
"muted": false
}
@@ -588,7 +588,7 @@
"name": "MODEL",
"type": "MODEL",
"links": [
113
116
],
"slot_index": 0
}
@@ -602,15 +602,15 @@
]
},
{
"id": 56,
"type": "TeaCacheForVidGen",
"id": 57,
"type": "TeaCache",
"pos": [
336.0735168457031,
-95.73665618896484
354.6571044921875,
-123.38716125488281
],
"size": [
315,
82
106
],
"flags": {},
"order": 8,
@@ -619,7 +619,7 @@
{
"name": "model",
"type": "MODEL",
"link": 113
"link": 116
}
],
"outputs": [
@@ -627,17 +627,18 @@
"name": "model",
"type": "MODEL",
"links": [
114
117
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForVidGen"
"Node name for S&R": "TeaCache"
},
"widgets_values": [
"wan2.1_i2v_480p_14B",
0.26
0.26,
3
]
}
],
@@ -771,16 +772,16 @@
"IMAGE"
],
[
113,
116,
37,
0,
56,
57,
0,
"MODEL"
],
[
114,
56,
117,
57,
0,
54,
0,
@@ -791,16 +792,16 @@
"config": {},
"extra": {
"ds": {
"scale": 0.7627768444385795,
"scale": 0.8264462809917354,
"offset": [
336.0382964059074,
205.80199013908017
191.78943826948876,
255.9140408750332
]
},
"node_versions": {
"comfy-core": "0.3.22",
"comfy-core": "0.3.26",
"ComfyUI-VideoHelperSuite": "124c913ccdd8a585734ea758c35fa1bab8499c99",
"ComfyUI-TeaCache": "062efbab2ec1296ba721d9985c5cd80c08da0e42"
"ComfyUI-TeaCache": "b3429ef3dea426d2f167e348b44cd2f5a3674e7d"
},
"ue_links": [],
"VHS_latentpreview": false,
examples/teacache_wan2.1_t2v.json → examples/wan2.1_t2v.json
+58 -57
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 51,
"last_link_id": 98,
"last_node_id": 52,
"last_link_id": 101,
"nodes": [
{
"id": 39,
@@ -299,51 +299,18 @@
"hidden": false,
"paused": false,
"params": {
"filename": "teacache_00029.mp4",
"filename": "teacache_00084.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 16,
"workflow": "teacache_00029.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/teacache_00029.mp4"
"workflow": "teacache_00084.png",
"fullpath": "/home/yuyunjie/code/ComfyUI/output/teacache_00084.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",
@@ -421,7 +388,7 @@
{
"name": "model",
"type": "MODEL",
"link": 98
"link": 101
}
],
"outputs": [
@@ -442,15 +409,48 @@
]
},
{
"id": 51,
"type": "TeaCacheForVidGen",
"id": 37,
"type": "UNETLoader",
"pos": [
295.4561767578125,
-147.876708984375
20,
40
],
"size": [
346.7470703125,
82
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
100
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "UNETLoader"
},
"widgets_values": [
"wan2.1_t2v_1.3B_fp16.safetensors",
"default"
]
},
{
"id": 52,
"type": "TeaCache",
"pos": [
386.0345458984375,
-121.91252136230469
],
"size": [
315,
82
106
],
"flags": {},
"order": 6,
@@ -459,7 +459,7 @@
{
"name": "model",
"type": "MODEL",
"link": 97
"link": 100
}
],
"outputs": [
@@ -467,17 +467,18 @@
"name": "model",
"type": "MODEL",
"links": [
98
101
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "TeaCacheForVidGen"
"Node name for S&R": "TeaCache"
},
"widgets_values": [
"wan2.1_t2v_1.3B",
0.08
0.08,
3
]
}
],
@@ -555,16 +556,16 @@
"IMAGE"
],
[
97,
100,
37,
0,
51,
52,
0,
"MODEL"
],
[
98,
51,
101,
52,
0,
48,
0,
@@ -575,16 +576,16 @@
"config": {},
"extra": {
"ds": {
"scale": 0.6934334949441631,
"scale": 0.7513148009015777,
"offset": [
509.8357423302538,
381.2577148997867
302.7703147246433,
297.73529230242883
]
},
"node_versions": {
"comfy-core": "0.3.22",
"comfy-core": "0.3.26",
"ComfyUI-VideoHelperSuite": "124c913ccdd8a585734ea758c35fa1bab8499c99",
"ComfyUI-TeaCache": "062efbab2ec1296ba721d9985c5cd80c08da0e42"
"ComfyUI-TeaCache": "b3429ef3dea426d2f167e348b44cd2f5a3674e7d"
},
"ue_links": [],
"VHS_latentpreview": false,
nodes.py
+126 -151
View File
@@ -8,14 +8,15 @@ 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.lightricks.symmetric_patchifier import latent_to_pixel_coords
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],
"hunyuan_video": [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02],
"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],
@@ -26,7 +27,7 @@ def poly1d(coefficients, x):
result = torch.zeros_like(x)
for i, coeff in enumerate(coefficients):
result += coeff * (x ** (len(coefficients) - 1 - i))
return result.abs()
return result
def teacache_flux_forward(
self,
@@ -42,8 +43,9 @@ def teacache_flux_forward(
attn_mask: Tensor = None,
) -> Tensor:
patches_replace = transformer_options.get("patches_replace", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh")
coefficients = transformer_options.get("coefficients")
max_skip_steps = transformer_options.get("max_skip_steps")
if img.ndim != 3 or txt.ndim != 3:
raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -75,17 +77,20 @@ def teacache_flux_forward(
if not hasattr(self, 'accumulated_rel_l1_distance'):
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
elif self.skip_steps == max_skip_steps:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 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:
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
self.skip_steps += 1
else:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
self.previous_modulated_input = modulated_inp
@@ -193,24 +198,35 @@ def teacache_hunyuanvideo_forward(
timesteps: Tensor,
y: Tensor,
guidance: Tensor = None,
guiding_frame_index=None,
control=None,
transformer_options={},
) -> Tensor:
patches_replace = transformer_options.get("patches_replace", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh")
coefficients = transformer_options.get("coefficients")
max_skip_steps = transformer_options.get("max_skip_steps")
initial_shape = list(img.shape)
# running on sequences img
img = self.img_in(img)
vec = self.time_in(timestep_embedding(timesteps, 256, time_factor=1.0).to(img.dtype))
vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
if guiding_frame_index is not None:
token_replace_vec = self.time_in(timestep_embedding(guiding_frame_index, 256, time_factor=1.0))
vec_ = self.vector_in(y[:, :self.params.vec_in_dim])
vec = torch.cat([(vec_ + token_replace_vec).unsqueeze(1), (vec_ + vec).unsqueeze(1)], dim=1)
frame_tokens = (initial_shape[-1] // self.patch_size[-1]) * (initial_shape[-2] // self.patch_size[-2])
modulation_dims = [(0, frame_tokens, 0), (frame_tokens, None, 1)]
modulation_dims_txt = [(0, None, 1)]
else:
vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
modulation_dims = None
modulation_dims_txt = None
if self.params.guidance_embed:
if guidance is None:
raise ValueError("Didn't get guidance strength for guidance distilled model.")
vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
if guidance is not None:
vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
if txt_mask is not None and not torch.is_floating_point(txt_mask):
txt_mask = (txt_mask - 1).to(img.dtype) * torch.finfo(img.dtype).max
@@ -240,17 +256,20 @@ def teacache_hunyuanvideo_forward(
if not hasattr(self, 'accumulated_rel_l1_distance'):
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
elif self.skip_steps == max_skip_steps:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 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:
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
self.skip_steps += 1
else:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
self.previous_modulated_input = modulated_inp
@@ -262,14 +281,14 @@ def teacache_hunyuanvideo_forward(
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
out["img"], out["txt"] = block(img=args["img"], txt=args["txt"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"])
out["img"], out["txt"] = block(img=args["img"], txt=args["txt"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"], modulation_dims_img=args["modulation_dims_img"], modulation_dims_txt=args["modulation_dims_txt"])
return out
out = blocks_replace[("double_block", i)]({"img": img, "txt": txt, "vec": vec, "pe": pe, "attention_mask": attn_mask}, {"original_block": block_wrap})
out = blocks_replace[("double_block", i)]({"img": img, "txt": txt, "vec": vec, "pe": pe, "attention_mask": attn_mask, 'modulation_dims_img': modulation_dims, 'modulation_dims_txt': modulation_dims_txt}, {"original_block": block_wrap})
txt = out["txt"]
img = out["img"]
else:
img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask)
img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask, modulation_dims_img=modulation_dims, modulation_dims_txt=modulation_dims_txt)
if control is not None: # Controlnet
control_i = control.get("input")
@@ -284,13 +303,13 @@ def teacache_hunyuanvideo_forward(
if ("single_block", i) in blocks_replace:
def block_wrap(args):
out = {}
out["img"] = block(args["img"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"])
out["img"] = block(args["img"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"], modulation_dims=args["modulation_dims"])
return out
out = blocks_replace[("single_block", i)]({"img": img, "vec": vec, "pe": pe, "attention_mask": attn_mask}, {"original_block": block_wrap})
out = blocks_replace[("single_block", i)]({"img": img, "vec": vec, "pe": pe, "attention_mask": attn_mask, 'modulation_dims': modulation_dims}, {"original_block": block_wrap})
img = out["img"]
else:
img = block(img, vec=vec, pe=pe, attn_mask=attn_mask)
img = block(img, vec=vec, pe=pe, attn_mask=attn_mask, modulation_dims=modulation_dims)
if control is not None: # Controlnet
control_o = control.get("output")
@@ -302,7 +321,7 @@ def teacache_hunyuanvideo_forward(
img = img[:, : img_len]
self.previous_residual = img - ori_img
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
img = self.final_layer(img, vec, modulation_dims=modulation_dims) # (N, T, patch_size ** 2 * out_channels)
shape = initial_shape[-3:]
for i in range(len(shape)):
@@ -319,57 +338,37 @@ def teacache_ltxvmodel_forward(
context,
attention_mask,
frame_rate=25,
guiding_latent=None,
guiding_latent_noise_scale=0,
transformer_options={},
keyframe_idxs=None,
**kwargs
):
patches_replace = transformer_options.get("patches_replace", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
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],
orig_height=x.shape[3],
orig_width=x.shape[4],
batch_size=x.shape[0],
scale_grid=((1 / frame_rate) * 8, 32, 32),
device=x.device,
)
if guiding_latent is not None:
ts = torch.ones([x.shape[0], 1, x.shape[2], x.shape[3], x.shape[4]], device=x.device, dtype=x.dtype)
input_ts = timestep.view([timestep.shape[0]] + [1] * (x.ndim - 1))
ts *= input_ts
ts[:, :, 0] = guiding_latent_noise_scale * (input_ts[:, :, 0] ** 2)
timestep = self.patchifier.patchify(ts)
input_x = x.clone()
x[:, :, 0] = guiding_latent[:, :, 0]
if guiding_latent_noise_scale > 0:
if self.generator is None:
self.generator = torch.Generator(device=x.device).manual_seed(42)
elif self.generator.device != x.device:
self.generator = torch.Generator(device=x.device).set_state(self.generator.get_state())
noise_shape = [guiding_latent.shape[0], guiding_latent.shape[1], 1, guiding_latent.shape[3], guiding_latent.shape[4]]
scale = guiding_latent_noise_scale * (input_ts ** 2)
guiding_noise = scale * torch.randn(size=noise_shape, device=x.device, generator=self.generator)
x[:, :, 0] = guiding_noise[:, :, 0] + x[:, :, 0] * (1.0 - scale[:, :, 0])
max_skip_steps = transformer_options.get("max_skip_steps")
orig_shape = list(x.shape)
x = self.patchifier.patchify(x)
x, latent_coords = self.patchifier.patchify(x)
pixel_coords = latent_to_pixel_coords(
latent_coords=latent_coords,
scale_factors=self.vae_scale_factors,
causal_fix=self.causal_temporal_positioning,
)
if keyframe_idxs is not None:
pixel_coords[:, :, -keyframe_idxs.shape[2]:] = keyframe_idxs
fractional_coords = pixel_coords.to(torch.float32)
fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)
x = self.patchify_proj(x)
timestep = timestep * 1000.0
attention_mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1]))
attention_mask = attention_mask.masked_fill(attention_mask.to(torch.bool), float("-inf")) # not sure about this
# attention_mask = (context != 0).any(dim=2).to(dtype=x.dtype)
if attention_mask is not None and not torch.is_floating_point(attention_mask):
attention_mask = (attention_mask - 1).to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(x.dtype).max
pe = precompute_freqs_cis(indices_grid, dim=self.inner_dim, out_dtype=x.dtype)
pe = precompute_freqs_cis(fractional_coords, dim=self.inner_dim, out_dtype=x.dtype)
batch_size = x.shape[0]
timestep, embedded_timestep = self.adaln_single(
@@ -406,21 +405,23 @@ def teacache_ltxvmodel_forward(
if not hasattr(self, 'accumulated_rel_l1_distance'):
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
elif self.skip_steps == max_skip_steps:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 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:
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
self.skip_steps += 1
else:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
self.previous_modulated_input = modulated_inp
if not should_calc:
x += self.previous_residual
else:
@@ -463,10 +464,6 @@ def teacache_ltxvmodel_forward(
out_channels=orig_shape[1] // math.prod(self.patchifier.patch_size),
)
if guiding_latent is not None:
x[:, :, 0] = (input_x[:, :, 0] - guiding_latent[:, :, 0]) / input_ts[:, :, 0]
# print("res", x)
return x
def teacache_wanmodel_forward(self, x, timestep, context, clip_fea=None, transformer_options={}, **kwargs):
@@ -494,8 +491,9 @@ def teacache_wanmodel_forward_orig(
freqs=None,
transformer_options={},
):
rel_l1_thresh = transformer_options.get("rel_l1_thresh", {})
rel_l1_thresh = transformer_options.get("rel_l1_thresh")
coefficients = transformer_options.get("coefficients")
max_skip_steps = transformer_options.get("max_skip_steps")
# embeddings
x = self.patch_embedding(x.float()).to(x.dtype)
@@ -526,19 +524,22 @@ def teacache_wanmodel_forward_orig(
if not hasattr(self, 'accumulated_rel_l1_distance'):
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
elif self.skip_steps == max_skip_steps:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 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:
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
self.skip_steps += 1
else:
should_calc = True
self.accumulated_rel_l1_distance = 0
self.skip_steps = 0
self.previous_modulated_input = modulated_inp
self.previous_modulated_input = modulated_inp
if not should_calc:
x += self.previous_residual
@@ -555,14 +556,15 @@ def teacache_wanmodel_forward_orig(
x = self.unpatchify(x, grid_sizes)
return x
class TeaCacheForImgGen:
class TeaCache:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL", {"tooltip": "The image diffusion model the TeaCache will be applied to."}),
"model_type": (["flux"],),
"rel_l1_thresh": ("FLOAT", {"default": 0.4, "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."})
"model": ("MODEL", {"tooltip": "The diffusion model the TeaCache will be applied to."}),
"model_type": (["flux", "ltxv", "hunyuan_video", "wan2.1_t2v_1.3B", "wan2.1_t2v_14B", "wan2.1_i2v_480p_14B", "wan2.1_i2v_720p_14B"], {"default": "flux", "tooltip": "Supported diffusion model."}),
"rel_l1_thresh": ("FLOAT", {"default": 0.4, "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."}),
"max_skip_steps": ([1, 2, 3], {"default": 3, "tooltip": "Max continuous skip steps."})
}
}
@@ -570,9 +572,9 @@ class TeaCacheForImgGen:
RETURN_NAMES = ("model",)
FUNCTION = "apply_teacache"
CATEGORY = "TeaCache"
TITLE = "TeaCache For Img Gen"
TITLE = "TeaCache"
def apply_teacache(self, model, model_type: str, rel_l1_thresh: float):
def apply_teacache(self, model, model_type: str, rel_l1_thresh: float, max_skip_steps: int):
if rel_l1_thresh == 0:
return (model,)
@@ -580,67 +582,25 @@ class TeaCacheForImgGen:
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"]["max_skip_steps"] = max_skip_steps
new_model.model_options["transformer_options"]["coefficients"] = SUPPORTED_MODELS_COEFFICIENTS[model_type]
diffusion_model = new_model.get_model_object("diffusion_model")
if model_type == "flux":
forward_name = "forward_orig"
replaced_forward_fn = teacache_flux_forward.__get__(
diffusion_model,
diffusion_model.__class__
)
else:
raise ValueError(f"Unknown type {model_type}")
def unet_wrapper_function(model_function, kwargs):
input = kwargs["input"]
timestep = kwargs["timestep"]
c = kwargs["c"]
with patch.object(diffusion_model, forward_name, replaced_forward_fn):
return model_function(input, timestep, **c)
new_model.set_model_unet_function_wrapper(unet_wrapper_function)
return (new_model,)
class TeaCacheForVidGen:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL", {"tooltip": "The video diffusion model the TeaCache will be applied to."}),
"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."})
}
}
RETURN_TYPES = ("MODEL",)
RETURN_NAMES = ("model",)
FUNCTION = "apply_teacache"
CATEGORY = "TeaCache"
TITLE = "TeaCache For Vid Gen"
def apply_teacache(self, model, model_type: str, rel_l1_thresh: float):
if rel_l1_thresh == 0:
return (model,)
new_model = model.clone()
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 "hunyuan_video" in model_type:
if "flux" in model_type:
context = patch.multiple(
diffusion_model,
forward_orig=teacache_hunyuanvideo_forward.__get__(diffusion_model, diffusion_model.__class__)
forward_orig=teacache_flux_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 "hunyuan_video" in model_type:
context = patch.multiple(
diffusion_model,
forward_orig=teacache_hunyuanvideo_forward.__get__(diffusion_model, diffusion_model.__class__)
)
elif "wan2.1" in model_type:
context = patch.multiple(
diffusion_model,
@@ -654,6 +614,22 @@ class TeaCacheForVidGen:
input = kwargs["input"]
timestep = kwargs["timestep"]
c = kwargs["c"]
# referenced from https://github.com/kijai/ComfyUI-KJNodes/blob/d126b62cebee81ea14ec06ea7cd7526999cb0554/nodes/model_optimization_nodes.py#L868
sigmas = c["transformer_options"]["sample_sigmas"]
matched_step_index = (sigmas == timestep[0]).nonzero()
if len(matched_step_index) > 0:
current_step_index = matched_step_index.item()
else:
current_step_index = 0
for i in range(len(sigmas) - 1):
# walk from beginning of steps until crossing the timestep
if (sigmas[i] - timestep) * (sigmas[i + 1] - timestep) <= 0:
current_step_index = i
break
if current_step_index == 0 and hasattr(diffusion_model, 'accumulated_rel_l1_distance'):
delattr(diffusion_model, 'accumulated_rel_l1_distance')
with context:
return model_function(input, timestep, **c)
@@ -766,8 +742,7 @@ class CompileModel:
NODE_CLASS_MAPPINGS = {
"TeaCacheForImgGen": TeaCacheForImgGen,
"TeaCacheForVidGen": TeaCacheForVidGen,
"TeaCache": TeaCache,
"CompileModel": CompileModel
}
pyproject.toml
+2 -2
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-TeaCache"
name = "teacache"
description = "Unofficial implementation of [ali-vilab/TeaCache](https://github.com/ali-vilab/TeaCache) for ComfyUI"
version = "1.3.2"
version = "1.4.0"
license = {file = "LICENSE"}
[project.urls]