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