diff --git a/README.md b/README.md index 8e4631c..115609e 100644 --- a/README.md +++ b/README.md @@ -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:
-| 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 |
-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: diff --git a/examples/teacache_cogvideox.json b/examples/cogvideox.json similarity index 100% rename from examples/teacache_cogvideox.json rename to examples/cogvideox.json diff --git a/examples/teacache_flux.json b/examples/flux.json similarity index 89% rename from examples/teacache_flux.json rename to examples/flux.json index 6aaa302..af5ab17 100644 --- a/examples/teacache_flux.json +++ b/examples/flux.json @@ -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 } \ No newline at end of file diff --git a/examples/teacache_hunyuanvideo.json b/examples/hunyuanvideo.json similarity index 88% rename from examples/teacache_hunyuanvideo.json rename to examples/hunyuanvideo.json index 92505cc..26a6f26 100644 --- a/examples/teacache_hunyuanvideo.json +++ b/examples/hunyuanvideo.json @@ -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 } \ No newline at end of file diff --git a/examples/teacache_ltx_video.json b/examples/ltx_video.json similarity index 90% rename from examples/teacache_ltx_video.json rename to examples/ltx_video.json index 720f491..65e0c0f 100644 --- a/examples/teacache_ltx_video.json +++ b/examples/ltx_video.json @@ -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 } \ No newline at end of file diff --git a/examples/teacache_pulid_flux.json b/examples/pulid_flux.json similarity index 88% rename from examples/teacache_pulid_flux.json rename to examples/pulid_flux.json index b35d5d8..d7d75bd 100644 --- a/examples/teacache_pulid_flux.json +++ b/examples/pulid_flux.json @@ -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 } \ No newline at end of file diff --git a/examples/teacache_wan2.1_i2v.json b/examples/wan2.1_i2v.json similarity index 95% rename from examples/teacache_wan2.1_i2v.json rename to examples/wan2.1_i2v.json index bc32f1a..b1a7527 100644 --- a/examples/teacache_wan2.1_i2v.json +++ b/examples/wan2.1_i2v.json @@ -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, diff --git a/examples/teacache_wan2.1_t2v.json b/examples/wan2.1_t2v.json similarity index 94% rename from examples/teacache_wan2.1_t2v.json rename to examples/wan2.1_t2v.json index 05c25b7..999f98a 100644 --- a/examples/teacache_wan2.1_t2v.json +++ b/examples/wan2.1_t2v.json @@ -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, diff --git a/nodes.py b/nodes.py index c9458a5..78b5200 100644 --- a/nodes.py +++ b/nodes.py @@ -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 } diff --git a/pyproject.toml b/pyproject.toml index acddfa8..28bfcd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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]