From f74a6cb427244f095029cef67e4961ee73c93c00 Mon Sep 17 00:00:00 2001 From: hkz Date: Thu, 23 Jan 2025 15:52:48 +0800 Subject: [PATCH] Support TeaCache (#180) * Support TeaCache * Update app * Update teacache --------- Co-authored-by: bubbliiiing <3323290568@qq.com> --- app.py | 9 +- comfyui/comfyui_nodes.py | 23 +- ...syanimatev5.1_workflow_control_camera.json | 189 ++++----- ...imatev5.1_workflow_control_trajectory.json | 364 +++++++++--------- .../v5.1/easyanimatev5.1_workflow_i2v.json | 4 +- .../v5.1/easyanimatev5.1_workflow_t2v.json | 4 +- .../v5.1/easyanimatev5.1_workflow_v2v.json | 2 + .../easyanimatev5.1_workflow_v2v_control.json | 55 ++- comfyui/v5/easyanimatev5_workflow_v2v.json | 20 +- .../easyanimatev5_workflow_v2v_control.json | 20 +- easyanimate/models/transformer3d.py | 200 ++++++++-- easyanimate/ui/ui.py | 23 +- predict_i2v.py | 9 + predict_t2v.py | 9 + predict_v2v.py | 9 + predict_v2v_control.py | 9 + 16 files changed, 605 insertions(+), 344 deletions(-) diff --git a/app.py b/app.py index eafead6..eec6879 100644 --- a/app.py +++ b/app.py @@ -24,6 +24,11 @@ if __name__ == "__main__": # EasyAnimateV4, V5 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload" # EasyAnimateV5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" GPU_memory_mode = "model_cpu_offload_and_qfloat8" + # EasyAnimateV5.1 support TeaCache. + enable_teacache = True + # Recommended to be set between 0.05 and 0.1. A larger threshold can cache more steps, speeding up the inference process, + # but it may cause slight differences between the generated content and the original content. + teacache_threshold = 0.1 # Use torch.float16 if GPU does not support torch.bfloat16 # ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16 weight_dtype = torch.bfloat16 @@ -44,11 +49,11 @@ if __name__ == "__main__": savedir_sample = "samples" if ui_mode == "modelscope": - demo, controller = ui_modelscope(model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, weight_dtype) + demo, controller = ui_modelscope(model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype) elif ui_mode == "eas": demo, controller = ui_eas(edition, config_path, model_name, savedir_sample) else: - demo, controller = ui(GPU_memory_mode, weight_dtype) + demo, controller = ui(GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype) # launch gradio app, _, _ = demo.queue(status_update_rate=1).launch( diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py index 51b0680..c9cf642 100644 --- a/comfyui/comfyui_nodes.py +++ b/comfyui/comfyui_nodes.py @@ -442,7 +442,7 @@ class EasyAnimateT2VSampler: FUNCTION = "process" CATEGORY = "EasyAnimateWrapper" - def process(self, easyanimate_model, prompt, negative_prompt, video_length, width, height, is_image, seed, steps, cfg, scheduler): + def process(self, easyanimate_model, prompt, negative_prompt, video_length, width, height, is_image, seed, steps, cfg, scheduler, teacache_threshold=0.10, enable_teacache=False): global transformer_cpu_cache global lora_path_before device = mm.get_torch_device() @@ -459,6 +459,9 @@ class EasyAnimateT2VSampler: # Load Sampler pipeline.scheduler = all_cheduler_dict[scheduler].from_pretrained(model_name, subfolder='scheduler') + if enable_teacache: + pipeline.transformer.enable_teacache(steps, teacache_threshold) + generator= torch.Generator(device).manual_seed(seed) video_length = 1 if is_image else video_length @@ -588,6 +591,8 @@ class EasyAnimateV5_T2VSampler(EasyAnimateT2VSampler): "default": 'Flow' } ), + "teacache_threshold": ("FLOAT", {"default": 0.10, "min": 0.00, "max": 1.00, "step": 0.005}), + "enable_teacache":([False, True], {"default": True,}), }, } @@ -649,7 +654,7 @@ class EasyAnimateI2VSampler: FUNCTION = "process" CATEGORY = "EasyAnimateWrapper" - def process(self, easyanimate_model, prompt, negative_prompt, video_length, base_resolution, seed, steps, cfg, scheduler, start_img=None, end_img=None): + def process(self, easyanimate_model, prompt, negative_prompt, video_length, base_resolution, seed, steps, cfg, scheduler, start_img=None, end_img=None, teacache_threshold=0.10, enable_teacache=False): global transformer_cpu_cache global lora_path_before device = mm.get_torch_device() @@ -674,6 +679,9 @@ class EasyAnimateI2VSampler: # Load Sampler pipeline.scheduler = all_cheduler_dict[scheduler].from_pretrained(model_name, subfolder='scheduler') + if enable_teacache: + pipeline.transformer.enable_teacache(steps, teacache_threshold) + generator= torch.Generator(device).manual_seed(seed) with torch.no_grad(): @@ -780,7 +788,9 @@ class EasyAnimateV5_I2VSampler(EasyAnimateI2VSampler): { "default": 'Flow' } - ) + ), + "teacache_threshold": ("FLOAT", {"default": 0.10, "min": 0.00, "max": 1.00, "step": 0.005}), + "enable_teacache":([False, True], {"default": True,}), }, "optional":{ "start_img": ("IMAGE",), @@ -849,7 +859,7 @@ class EasyAnimateV2VSampler: FUNCTION = "process" CATEGORY = "EasyAnimateWrapper" - def process(self, easyanimate_model, prompt, negative_prompt, video_length, base_resolution, seed, steps, cfg, denoise_strength, scheduler, validation_video=None, control_video=None, ref_image=None, camera_conditions=None): + def process(self, easyanimate_model, prompt, negative_prompt, video_length, base_resolution, seed, steps, cfg, denoise_strength, scheduler, validation_video=None, control_video=None, ref_image=None, camera_conditions=None, teacache_threshold=0.10, enable_teacache=False): global transformer_cpu_cache global lora_path_before @@ -892,6 +902,9 @@ class EasyAnimateV2VSampler: # Load Sampler pipeline.scheduler = all_cheduler_dict[scheduler].from_pretrained(model_name, subfolder='scheduler') + if enable_teacache: + pipeline.transformer.enable_teacache(steps, teacache_threshold) + generator= torch.Generator(device).manual_seed(seed) with torch.no_grad(): @@ -1033,6 +1046,8 @@ class EasyAnimateV5_V2VSampler(EasyAnimateV2VSampler): "default": 'Flow' } ), + "teacache_threshold": ("FLOAT", {"default": 0.10, "min": 0.00, "max": 1.00, "step": 0.005}), + "enable_teacache":([False, True], {"default": True,}), }, "optional":{ "validation_video": ("IMAGE",), diff --git a/comfyui/v5.1/easyanimatev5.1_workflow_control_camera.json b/comfyui/v5.1/easyanimatev5.1_workflow_control_camera.json index c9e43f3..55a7555 100644 --- a/comfyui/v5.1/easyanimatev5.1_workflow_control_camera.json +++ b/comfyui/v5.1/easyanimatev5.1_workflow_control_camera.json @@ -214,89 +214,6 @@ "bf16" ] }, - { - "id": 131, - "type": "EasyAnimateV5_V2VSampler", - "pos": { - "0": 821, - "1": 242 - }, - "size": { - "0": 504, - "1": 350 - }, - "flags": {}, - "order": 12, - "mode": 0, - "inputs": [ - { - "name": "easyanimate_model", - "type": "EASYANIMATESMODEL", - "link": 271 - }, - { - "name": "prompt", - "type": "STRING_PROMPT", - "link": 272 - }, - { - "name": "negative_prompt", - "type": "STRING_PROMPT", - "link": 273 - }, - { - "name": "validation_video", - "type": "IMAGE", - "link": null, - "shape": 7 - }, - { - "name": "control_video", - "type": "IMAGE", - "link": null, - "shape": 7 - }, - { - "name": "ref_image", - "type": "IMAGE", - "link": 274, - "shape": 7 - }, - { - "name": "camera_conditions", - "type": "STRING", - "link": 275, - "widget": { - "name": "camera_conditions" - }, - "shape": 7 - } - ], - "outputs": [ - { - "name": "images", - "type": "IMAGE", - "links": [ - 276 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "EasyAnimateV5_V2VSampler" - }, - "widgets_values": [ - 49, - 512, - 43, - "fixed", - 43, - 6, - 1, - "Flow", - "" - ] - }, { "id": 129, "type": "CameraTrajectoryFromChaoJie", @@ -454,7 +371,7 @@ }, "size": [ 390, - 535.4285714285714 + 546 ], "flags": {}, "order": 13, @@ -532,10 +449,10 @@ "0": 819, "1": 658 }, - "size": [ - 517.6458089787227, - 93.61251593411134 - ], + "size": { + "0": 517.6458129882812, + "1": 93.61251831054688 + }, "flags": {}, "order": 10, "mode": 0, @@ -549,6 +466,91 @@ ], "color": "#432", "bgcolor": "#653" + }, + { + "id": 131, + "type": "EasyAnimateV5_V2VSampler", + "pos": { + "0": 822, + "1": 211 + }, + "size": { + "0": 504, + "1": 394 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "easyanimate_model", + "type": "EASYANIMATESMODEL", + "link": 271 + }, + { + "name": "prompt", + "type": "STRING_PROMPT", + "link": 272 + }, + { + "name": "negative_prompt", + "type": "STRING_PROMPT", + "link": 273 + }, + { + "name": "validation_video", + "type": "IMAGE", + "link": null, + "shape": 7 + }, + { + "name": "control_video", + "type": "IMAGE", + "link": null, + "shape": 7 + }, + { + "name": "ref_image", + "type": "IMAGE", + "link": 274, + "shape": 7 + }, + { + "name": "camera_conditions", + "type": "STRING", + "link": 275, + "widget": { + "name": "camera_conditions" + }, + "shape": 7 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 276 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EasyAnimateV5_V2VSampler" + }, + "widgets_values": [ + 49, + 512, + 43, + "fixed", + 43, + 6, + 1, + "Flow", + 0.10, + true, + "" + ] } ], "links": [ @@ -662,11 +664,16 @@ "config": {}, "extra": { "ds": { - "scale": 0.6209213230591558, + "scale": 1.1, "offset": [ - 417.7460035994012, - -70.36580413723722 + -465.8996857769304, + 51.92597569190605 ] + }, + "node_versions": { + "EasyAnimate": "de24d49f07f6d9b12b4e98de98ec959a8b44b989", + "comfy-core": "v0.2.7-3-g8afb97c", + "ComfyUI-VideoHelperSuite": "70faa9bcef65932ab72e7404d6373fb300013a2e" } }, "version": 0.4 diff --git a/comfyui/v5.1/easyanimatev5.1_workflow_control_trajectory.json b/comfyui/v5.1/easyanimatev5.1_workflow_control_trajectory.json index 1dd2ddc..7d88aa5 100644 --- a/comfyui/v5.1/easyanimatev5.1_workflow_control_trajectory.json +++ b/comfyui/v5.1/easyanimatev5.1_workflow_control_trajectory.json @@ -211,86 +211,6 @@ "color": "#432", "bgcolor": "#653" }, - { - "id": 106, - "type": "VHS_VideoCombine", - "pos": { - "0": 1497, - "1": 204 - }, - "size": [ - 390, - 546 - ], - "flags": {}, - "order": 14, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 261, - "slot_index": 0, - "label": "图像", - "shape": 7 - }, - { - "name": "audio", - "type": "AUDIO", - "link": null, - "label": "音频", - "shape": 7 - }, - { - "name": "meta_batch", - "type": "VHS_BatchManager", - "link": null, - "label": "批次管理", - "shape": 7 - }, - { - "name": "vae", - "type": "VAE", - "link": null, - "shape": 7 - } - ], - "outputs": [ - { - "name": "Filenames", - "type": "VHS_FILENAMES", - "links": null, - "slot_index": 0, - "shape": 3, - "label": "文件名" - } - ], - "properties": { - "Node name for S&R": "VHS_VideoCombine" - }, - "widgets_values": { - "frame_rate": 8, - "loop_count": 0, - "filename_prefix": "EasyAnimate", - "format": "video/h264-mp4", - "pix_fmt": "yuv420p", - "crf": 22, - "save_metadata": true, - "pingpong": false, - "save_output": true, - "videopreview": { - "hidden": false, - "paused": false, - "params": { - "filename": "EasyAnimate_00105.mp4", - "subfolder": "", - "type": "output", - "format": "video/h264-mp4", - "frame_rate": 8 - } - } - } - }, { "id": 99, "type": "LoadEasyAnimateModel", @@ -321,96 +241,12 @@ }, "widgets_values": [ "EasyAnimateV5.1-12b-zh-Control", - "model_cpu_offload", + "model_cpu_offload_and_qfloat8", "Control", "easyanimate_video_v5.1_magvit_qwen.yaml", "bf16" ] }, - { - "id": 111, - "type": "EasyAnimateV5_V2VSampler", - "pos": { - "0": 905, - "1": 201 - }, - "size": { - "0": 504, - "1": 350 - }, - "flags": {}, - "order": 13, - "mode": 0, - "inputs": [ - { - "name": "easyanimate_model", - "type": "EASYANIMATESMODEL", - "link": 256 - }, - { - "name": "prompt", - "type": "STRING_PROMPT", - "link": 257 - }, - { - "name": "negative_prompt", - "type": "STRING_PROMPT", - "link": 258 - }, - { - "name": "validation_video", - "type": "IMAGE", - "link": null, - "shape": 7 - }, - { - "name": "control_video", - "type": "IMAGE", - "link": 259, - "shape": 7 - }, - { - "name": "ref_image", - "type": "IMAGE", - "link": 260, - "shape": 7 - }, - { - "name": "camera_conditions", - "type": "STRING", - "link": null, - "widget": { - "name": "camera_conditions" - }, - "shape": 7 - } - ], - "outputs": [ - { - "name": "images", - "type": "IMAGE", - "links": [ - 261, - 263 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "EasyAnimateV5_V2VSampler" - }, - "widgets_values": [ - 49, - 512, - 43, - "fixed", - 50, - 6, - 1, - "Flow", - "" - ] - }, { "id": 97, "type": "SplineEditor", @@ -498,7 +334,7 @@ }, "size": [ 530, - 650.4 + 310 ], "flags": {}, "order": 12, @@ -621,10 +457,10 @@ "0": 1140.1396484375, "1": 909.9193115234375 }, - "size": [ - 315, - 82 - ], + "size": { + "0": 315, + "1": 82 + }, "flags": { "collapsed": false }, @@ -715,7 +551,7 @@ }, "size": [ 530, - 630 + 310 ], "flags": {}, "order": 16, @@ -790,10 +626,10 @@ "0": 1544.139404296875, "1": 1047.919189453125 }, - "size": [ - 645, - 812 - ], + "size": { + "0": 645, + "1": 812 + }, "flags": {}, "order": 9, "mode": 0, @@ -858,6 +694,172 @@ null, null ] + }, + { + "id": 111, + "type": "EasyAnimateV5_V2VSampler", + "pos": { + "0": 876, + "1": 147 + }, + "size": { + "0": 504, + "1": 394 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "easyanimate_model", + "type": "EASYANIMATESMODEL", + "link": 256 + }, + { + "name": "prompt", + "type": "STRING_PROMPT", + "link": 257 + }, + { + "name": "negative_prompt", + "type": "STRING_PROMPT", + "link": 258 + }, + { + "name": "validation_video", + "type": "IMAGE", + "link": null, + "shape": 7 + }, + { + "name": "control_video", + "type": "IMAGE", + "link": 259, + "shape": 7 + }, + { + "name": "ref_image", + "type": "IMAGE", + "link": 260, + "shape": 7 + }, + { + "name": "camera_conditions", + "type": "STRING", + "link": null, + "widget": { + "name": "camera_conditions" + }, + "shape": 7 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 261, + 263 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EasyAnimateV5_V2VSampler" + }, + "widgets_values": [ + 49, + 512, + 43, + "fixed", + 50, + 6, + 1, + "Flow", + 0.10, + true, + "" + ] + }, + { + "id": 106, + "type": "VHS_VideoCombine", + "pos": { + "0": 1496, + "1": 153 + }, + "size": [ + 390, + 310 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 261, + "slot_index": 0, + "label": "图像", + "shape": 7 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "label": "音频", + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "label": "批次管理", + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "slot_index": 0, + "shape": 3, + "label": "文件名" + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 8, + "loop_count": 0, + "filename_prefix": "EasyAnimate", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 22, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "EasyAnimate_00105.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 8 + } + } + } } ], "links": [ @@ -1027,11 +1029,17 @@ "config": {}, "extra": { "ds": { - "scale": 0.5054470284993016, + "scale": 1.1, "offset": [ - -239.62344289721364, - 114.98345539297088 + -465.8996857769304, + 51.92597569190605 ] + }, + "node_versions": { + "EasyAnimate": "de24d49f07f6d9b12b4e98de98ec959a8b44b989", + "comfy-core": "v0.2.7-3-g8afb97c", + "ComfyUI-KJNodes": "4c5c26a2c91de356212419ac8bc7fcf9869527e9", + "ComfyUI-VideoHelperSuite": "70faa9bcef65932ab72e7404d6373fb300013a2e" } }, "version": 0.4 diff --git a/comfyui/v5.1/easyanimatev5.1_workflow_i2v.json b/comfyui/v5.1/easyanimatev5.1_workflow_i2v.json index 90daeb3..6d96d3c 100644 --- a/comfyui/v5.1/easyanimatev5.1_workflow_i2v.json +++ b/comfyui/v5.1/easyanimatev5.1_workflow_i2v.json @@ -289,7 +289,9 @@ "fixed", 50, 6, - "Flow" + "Flow", + 0.10, + true ] }, { diff --git a/comfyui/v5.1/easyanimatev5.1_workflow_t2v.json b/comfyui/v5.1/easyanimatev5.1_workflow_t2v.json index 74b10e7..b8496c4 100644 --- a/comfyui/v5.1/easyanimatev5.1_workflow_t2v.json +++ b/comfyui/v5.1/easyanimatev5.1_workflow_t2v.json @@ -84,7 +84,9 @@ "fixed", 50, 6, - "Flow" + "Flow", + 0.10, + true ] }, { diff --git a/comfyui/v5.1/easyanimatev5.1_workflow_v2v.json b/comfyui/v5.1/easyanimatev5.1_workflow_v2v.json index 3b5b504..419526d 100644 --- a/comfyui/v5.1/easyanimatev5.1_workflow_v2v.json +++ b/comfyui/v5.1/easyanimatev5.1_workflow_v2v.json @@ -305,6 +305,8 @@ 6, 0.7000000000000001, "Flow", + 0.10, + true, "" ] }, diff --git a/comfyui/v5.1/easyanimatev5.1_workflow_v2v_control.json b/comfyui/v5.1/easyanimatev5.1_workflow_v2v_control.json index 2665324..3ef4133 100644 --- a/comfyui/v5.1/easyanimatev5.1_workflow_v2v_control.json +++ b/comfyui/v5.1/easyanimatev5.1_workflow_v2v_control.json @@ -154,7 +154,7 @@ }, "size": [ 390.9534912109375, - 973.1686096191406 + 546.5720947265625 ], "flags": {}, "order": 9, @@ -234,7 +234,7 @@ }, "size": { "0": 336, - "1": 350 + "1": 394 }, "flags": {}, "order": 8, @@ -275,6 +275,15 @@ "type": "IMAGE", "link": null, "shape": 7 + }, + { + "name": "camera_conditions", + "type": "STRING", + "link": null, + "widget": { + "name": "camera_conditions" + }, + "shape": 7 } ], "outputs": [ @@ -300,6 +309,8 @@ 6, 1, "Flow", + 0.10, + true, "" ] }, @@ -349,7 +360,7 @@ }, "size": [ 252.056640625, - 685.7 + 262 ], "flags": {}, "order": 6, @@ -432,10 +443,10 @@ "0": -192, "1": -293 }, - "size": [ - 427.074951171875, - 143.9142608642578 - ], + "size": { + "0": 427.074951171875, + "1": 143.9142608642578 + }, "flags": {}, "order": 7, "mode": 0, @@ -495,14 +506,14 @@ ], "groups": [ { - "title": "Prompts", + "title": "Upload Your Video", "bounding": [ 218, - -127, - 450, - 483 + 385, + 487, + 789 ], - "color": "#3f789e", + "color": "#a1309b", "font_size": 24, "flags": {} }, @@ -519,14 +530,14 @@ "flags": {} }, { - "title": "Upload Your Video", + "title": "Prompts", "bounding": [ 218, - 385, - 487, - 789 + -127, + 450, + 483 ], - "color": "#a1309b", + "color": "#3f789e", "font_size": 24, "flags": {} } @@ -534,14 +545,18 @@ "config": {}, "extra": { "ds": { - "scale": 0.5644739300537782, + "scale": 0.8264462809917354, "offset": [ - 634.4708817322136, - 478.05663679245043 + -156.13347668602108, + 275.2525393282698 ] }, "workspace_info": { "id": "776b62b4-bd17-4ed3-9923-b7aad000b1ea" + }, + "node_versions": { + "EasyAnimate": "de24d49f07f6d9b12b4e98de98ec959a8b44b989", + "ComfyUI-VideoHelperSuite": "70faa9bcef65932ab72e7404d6373fb300013a2e" } }, "version": 0.4 diff --git a/comfyui/v5/easyanimatev5_workflow_v2v.json b/comfyui/v5/easyanimatev5_workflow_v2v.json index 6903ecf..e3d9d76 100644 --- a/comfyui/v5/easyanimatev5_workflow_v2v.json +++ b/comfyui/v5/easyanimatev5_workflow_v2v.json @@ -226,6 +226,21 @@ "type": "IMAGE", "link": null, "shape": 7 + }, + { + "name": "ref_image", + "type": "IMAGE", + "link": null, + "shape": 7 + }, + { + "name": "camera_conditions", + "type": "STRING", + "link": null, + "widget": { + "name": "camera_conditions" + }, + "shape": 7 } ], "outputs": [ @@ -250,7 +265,10 @@ 35, 7, 0.7, - "DDIM" + "DDIM", + 0.10, + true, + "" ] }, { diff --git a/comfyui/v5/easyanimatev5_workflow_v2v_control.json b/comfyui/v5/easyanimatev5_workflow_v2v_control.json index 9af90be..4eab714 100644 --- a/comfyui/v5/easyanimatev5_workflow_v2v_control.json +++ b/comfyui/v5/easyanimatev5_workflow_v2v_control.json @@ -391,6 +391,21 @@ "type": "IMAGE", "link": 53, "shape": 7 + }, + { + "name": "ref_image", + "type": "IMAGE", + "link": null, + "shape": 7 + }, + { + "name": "camera_conditions", + "type": "STRING", + "link": null, + "widget": { + "name": "camera_conditions" + }, + "shape": 7 } ], "outputs": [ @@ -415,7 +430,10 @@ 35, 6, 1, - "DDIM" + "DDIM", + 0.10, + true, + "" ] }, { diff --git a/easyanimate/models/transformer3d.py b/easyanimate/models/transformer3d.py index 7793ac2..53b7ead 100644 --- a/easyanimate/models/transformer3d.py +++ b/easyanimate/models/transformer3d.py @@ -87,6 +87,40 @@ class Transformer3DModelOutput(BaseOutput): sample: torch.FloatTensor +class TeaCache(): + """ + Timestep Embedding Aware Cache, a training-free caching approach that estimates and leverages + the fluctuating differences among model outputs across timesteps, thereby accelerating the inference. + Please refer to: + 1. https://github.com/ali-vilab/TeaCache. + 2. Liu, Feng, et al. "Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model." arXiv preprint arXiv:2411.19108 (2024). + """ + def __init__(self, coefficients: list[float], num_steps: int, rel_l1_thresh: float = 0.0): + if num_steps < 1: + raise ValueError("`num_steps` must be greater than 0 but is {num_steps}.") + if rel_l1_thresh < 0: + raise ValueError("`rel_l1_thresh` must be greater than or equal to 0 but is {rel_l1_thresh}.") + self.coefficients = coefficients + self.cnt = 0 + self.num_steps = num_steps + self.rel_l1_thresh = rel_l1_thresh + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = None + self.previous_residual = None + self.rescale_func = np.poly1d(self.coefficients) + + @staticmethod + def compute_rel_l1_distance(prev, cur): + rel_l1_distance = (torch.abs(cur - prev).mean()) / torch.abs(prev).mean() + + return rel_l1_distance.cpu().item() + + def reset(self): + self.cnt = 0 + self.previous_modulated_input = None + self.previous_residual = None + + class Transformer3DModel(ModelMixin, ConfigMixin): """ A 3D Transformer model for image-like data. @@ -1428,7 +1462,21 @@ class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin): ) self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * out_channels) + self.teacache = None + self.gradient_checkpointing = False + + def enable_teacache( + self, + num_steps: int, + rel_l1_thresh: float, + coefficients: list[float] = [-10.47857366, 8.33844143, -0.78477557, 0.68798618, 0.0136149] + ): + # The coefficient was obtained by sampling videos from T2V CompBench using EasyAnimateV5.1-12b-zh-InP. + # This coefficient can be applied to both the EasyAnimateV5.1-12b-zh and EasyAnimateV5.1-12b-Control. + # The coefficients for EasyAnimateV5.1-7b-zh-InP should be: + # [-3.64204720e+03, 1.43764725e+03, -1.93045263e+02, 1.09596499e+01, -1.70663507e-01] + self.teacache = TeaCache(coefficients, num_steps, rel_l1_thresh=rel_l1_thresh) def _set_gradient_checkpointing(self, module, value=False): self.gradient_checkpointing = value @@ -1499,48 +1547,122 @@ class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin): clip_encoder_hidden_states = self.clip_proj(clip_encoder_hidden_states) encoder_hidden_states = torch.concat([clip_encoder_hidden_states, ref_latents], dim=1) - - # 4. Transformer blocks - for i, block in enumerate(self.transformer_blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - def create_custom_forward(module, return_dict=None): - def custom_forward(*inputs): - if return_dict is not None: - return module(*inputs, return_dict=return_dict) - else: - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - hidden_states, - encoder_hidden_states, - temb, - image_rotary_emb, - video_length, - height // self.patch_size, - width // self.patch_size, - **ckpt_kwargs, - ) + + # TeaCache + if self.teacache is not None: + inp = hidden_states.clone() + temb_ = temb.clone() + encoder_hidden_states_ = encoder_hidden_states.clone() + modulated_inp, _, _, _ = self.transformer_blocks[0].norm1(inp, encoder_hidden_states_, temb_) + if self.teacache.cnt == 0 or self.teacache.cnt == self.teacache.num_steps - 1: + should_calc = True + self.teacache.accumulated_rel_l1_distance = 0 else: - hidden_states, encoder_hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - image_rotary_emb=image_rotary_emb, - num_frames=video_length, - height=height // self.patch_size, - width=width // self.patch_size - ) + rel_l1_distance = self.teacache.compute_rel_l1_distance(self.teacache.previous_modulated_input, modulated_inp) + self.teacache.accumulated_rel_l1_distance += self.teacache.rescale_func(rel_l1_distance) + if self.teacache.accumulated_rel_l1_distance < self.teacache.rel_l1_thresh: + should_calc = False + else: + should_calc = True + self.teacache.accumulated_rel_l1_distance = 0 + self.teacache.previous_modulated_input = modulated_inp + self.teacache.cnt += 1 + if self.teacache.cnt == self.teacache.num_steps: + # self.cnt = 0 + self.teacache.reset() - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - hidden_states = self.norm_final(hidden_states) - hidden_states = hidden_states[:, encoder_hidden_states.size()[1]:] + # TeaCache + if self.teacache is not None: + if not should_calc: + hidden_states += self.teacache.previous_residual + else: + ori_hidden_states = hidden_states.clone() + + # 4. Transformer blocks + for i, block in enumerate(self.transformer_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + video_length, + height // self.patch_size, + width // self.patch_size, + **ckpt_kwargs, + ) + else: + hidden_states, encoder_hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + num_frames=video_length, + height=height // self.patch_size, + width=width // self.patch_size + ) + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + hidden_states = self.norm_final(hidden_states) + hidden_states = hidden_states[:, encoder_hidden_states.size()[1]:] + + # 5. Final block + hidden_states = self.norm_out(hidden_states, temb=temb) + self.teacache.previous_residual = hidden_states - ori_hidden_states + else: + # 4. Transformer blocks + for i, block in enumerate(self.transformer_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + video_length, + height // self.patch_size, + width // self.patch_size, + **ckpt_kwargs, + ) + else: + hidden_states, encoder_hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + num_frames=video_length, + height=height // self.patch_size, + width=width // self.patch_size + ) + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + hidden_states = self.norm_final(hidden_states) + hidden_states = hidden_states[:, encoder_hidden_states.size()[1]:] + + # 5. Final block + hidden_states = self.norm_out(hidden_states, temb=temb) - # 5. Final block - hidden_states = self.norm_out(hidden_states, temb=temb) hidden_states = self.proj_out(hidden_states) # 6. Unpatchify diff --git a/easyanimate/ui/ui.py b/easyanimate/ui/ui.py index 85fe9b3..d1e2f0c 100755 --- a/easyanimate/ui/ui.py +++ b/easyanimate/ui/ui.py @@ -67,7 +67,7 @@ css = """ """ class EasyAnimateController: - def __init__(self, GPU_memory_mode, weight_dtype): + def __init__(self, GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype): # config dirs self.basedir = os.getcwd() self.config_dir = os.path.join(self.basedir, "config") @@ -97,6 +97,8 @@ class EasyAnimateController: self.base_model_path = "none" self.lora_model_path = "none" self.GPU_memory_mode = GPU_memory_mode + self.enable_teacache = enable_teacache + self.teacache_threshold = teacache_threshold self.weight_dtype = weight_dtype self.edition = "v5.1" @@ -462,6 +464,9 @@ class EasyAnimateController: # lora part self.pipeline = merge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider) + if self.edition == "v5.1" and self.enable_teacache: + self.pipeline.transformer.enable_teacache(sample_step_slider, self.teacache_threshold) + try: if self.model_type == "Inpaint": if self.transformer.config.in_channels != self.vae.config.latent_channels: @@ -661,8 +666,8 @@ class EasyAnimateController: return gr.Image.update(visible=False, value=None), gr.Video.update(value=save_sample_path, visible=True), "Success" -def ui(GPU_memory_mode, weight_dtype): - controller = EasyAnimateController(GPU_memory_mode, weight_dtype) +def ui(GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype): + controller = EasyAnimateController(GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype) with gr.Blocks(css=css) as demo: gr.Markdown( @@ -1000,7 +1005,7 @@ def ui(GPU_memory_mode, weight_dtype): class EasyAnimateController_Modelscope: - def __init__(self, model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, weight_dtype): + def __init__(self, model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype): # Basic dir self.basedir = os.getcwd() self.personalized_model_dir = os.path.join(self.basedir, "models", "Personalized_Model") @@ -1012,6 +1017,8 @@ class EasyAnimateController_Modelscope: # Config and model path self.model_type = model_type self.edition = edition + self.enable_teacache = enable_teacache + self.teacache_threshold = teacache_threshold self.weight_dtype = weight_dtype self.inference_config = OmegaConf.load(config_path) Choosen_AutoencoderKL = name_to_autoencoder_magvit[ @@ -1259,6 +1266,10 @@ class EasyAnimateController_Modelscope: # lora part self.pipeline = merge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider) + if self.edition == "v5.1" and self.enable_teacache: + print(f"Enable TeaCache with threshold: {self.teacache_threshold}.") + self.pipeline.transformer.enable_teacache(sample_step_slider, self.teacache_threshold) + try: if self.model_type == "Inpaint": if self.vae.cache_mag_vae: @@ -1373,8 +1384,8 @@ class EasyAnimateController_Modelscope: return gr.Image.update(visible=False, value=None), gr.Video.update(value=save_sample_path, visible=True), "Success" -def ui_modelscope(model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, weight_dtype): - controller = EasyAnimateController_Modelscope(model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, weight_dtype) +def ui_modelscope(model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype): + controller = EasyAnimateController_Modelscope(model_type, edition, config_path, model_name, savedir_sample, GPU_memory_mode, enable_teacache, teacache_threshold, weight_dtype) with gr.Blocks(css=css) as demo: gr.Markdown( diff --git a/predict_i2v.py b/predict_i2v.py index 61ec857..654467f 100644 --- a/predict_i2v.py +++ b/predict_i2v.py @@ -33,6 +33,11 @@ from easyanimate.utils.utils import get_image_to_video_latent, save_videos_grid # EasyAnimateV4, V5 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload" # EasyAnimateV5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" GPU_memory_mode = "model_cpu_offload_and_qfloat8" +# EasyAnimateV5.1 support TeaCache. +enable_teacache = True +# Recommended to be set between 0.05 and 0.1. A larger threshold can cache more steps, speeding up the inference process, +# but it may cause slight differences between the generated content and the original content. +teacache_threshold = 0.1 # Config and model path config_path = "config/easyanimate_video_v5.1_magvit_qwen.yaml" @@ -248,6 +253,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() +if "v5.1" in config_path and enable_teacache: + print(f"Enable TeaCache with threshold: {teacache_threshold}.") + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold) + generator = torch.Generator(device="cuda").manual_seed(seed) if lora_path is not None: diff --git a/predict_t2v.py b/predict_t2v.py index 505345b..42f5b58 100644 --- a/predict_t2v.py +++ b/predict_t2v.py @@ -35,6 +35,11 @@ from easyanimate.utils.utils import get_image_to_video_latent, save_videos_grid # EasyAnimateV4, V5 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload" # EasyAnimateV5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" GPU_memory_mode = "model_cpu_offload_and_qfloat8" +# EasyAnimateV5.1 support TeaCache. +enable_teacache = True +# Recommended to be set between 0.05 and 0.1. A larger threshold can cache more steps, speeding up the inference process, +# but it may cause slight differences between the generated content and the original content. +teacache_threshold = 0.1 # Config and model path config_path = "config/easyanimate_video_v5.1_magvit_qwen.yaml" @@ -256,6 +261,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() +if "v5.1" in config_path and enable_teacache: + print(f"Enable TeaCache with threshold: {teacache_threshold}.") + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold) + generator = torch.Generator(device="cuda").manual_seed(seed) if lora_path is not None: diff --git a/predict_v2v.py b/predict_v2v.py index 5e54876..44aa569 100644 --- a/predict_v2v.py +++ b/predict_v2v.py @@ -34,6 +34,11 @@ from easyanimate.utils.utils import (get_video_to_video_latent, # EasyAnimateV4, V5 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload" # EasyAnimateV5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" GPU_memory_mode = "model_cpu_offload_and_qfloat8" +# EasyAnimateV5.1 support TeaCache. +enable_teacache = True +# Recommended to be set between 0.05 and 0.1. A larger threshold can cache more steps, speeding up the inference process, +# but it may cause slight differences between the generated content and the original content. +teacache_threshold = 0.1 # Config and model path config_path = "config/easyanimate_video_v5.1_magvit_qwen.yaml" @@ -243,6 +248,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() +if "v5.1" in config_path and enable_teacache: + print(f"Enable TeaCache with threshold: {teacache_threshold}.") + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold) + generator = torch.Generator(device="cuda").manual_seed(seed) if lora_path is not None: diff --git a/predict_v2v_control.py b/predict_v2v_control.py index e1d2678..d9640e4 100644 --- a/predict_v2v_control.py +++ b/predict_v2v_control.py @@ -34,6 +34,11 @@ from diffusers import FlowMatchEulerDiscreteScheduler # EasyAnimateV5 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload" # EasyAnimateV5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" GPU_memory_mode = "model_cpu_offload_and_qfloat8" +# EasyAnimateV5.1 support TeaCache. +enable_teacache = True +# Recommended to be set between 0.05 and 0.1. A larger threshold can cache more steps, speeding up the inference process, +# but it may cause slight differences between the generated content and the original content. +teacache_threshold = 0.1 # Config and model path config_path = "config/easyanimate_video_v5.1_magvit_qwen.yaml" @@ -228,6 +233,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() +if "v5.1" in config_path and enable_teacache: + print(f"Enable TeaCache with threshold: {teacache_threshold}.") + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold) + generator = torch.Generator(device="cuda").manual_seed(seed) if lora_path is not None: