fix get_meta_file.py and autocast

---------

Co-authored-by: bubbliiiing <3323290568@qq.com>

README.md

@@ -117,22 +117,19 @@ We need about 60GB available on disk (for saving weights), please check!
 The video size for EasyAnimateV5.1-12B can be generated by different GPU Memory, including:
 | GPU memory | 384x672x25 | 384x672x49 | 576x1008x25 | 576x1008x49 | 768x1344x25 | 768x1344x49 |
 |------------|------------|------------|------------|------------|------------|------------|
-| 16GB | 🧡 | 🧡 | ❌ | ❌ | ❌ | ❌ |
+| 16GB | 🧡 | ⭕️ | ⭕️ | ⭕️ | ❌ | ❌ |
 | 24GB | 🧡 | 🧡 | 🧡 | 🧡 | 🧡 | ❌ | 
 | 40GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 | 80GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 
-Due to the float16 weights of qwen2-vl-7b, it cannot run on a 16GB GPU. If your GPU memory is 16GB, please visit [Huggingface](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct-GPTQ-Int8) or [Modelscope](https://modelscope.cn/models/Qwen/Qwen2-VL-7B-Instruct-GPTQ-Int8) to download the quantized version of qwen2-vl-7b to replace the original text encoder, and install the corresponding dependency libraries (auto-gptq, optimum).
-
-The video size for EasyAnimateV5-7B can be generated by different GPU Memory, including:
+The video size for EasyAnimateV5.1-7B can be generated by different GPU Memory, including:
 | GPU memory |384x672x25|384x672x49|576x1008x25|576x1008x49|768x1344x25|768x1344x49|
 |----------|----------|----------|----------|----------|----------|----------|
-| 16GB | 🧡 | 🧡 | ❌ | ❌ | ❌ | ❌ | 
+| 16GB | 🧡 | 🧡 | ⭕️ | ⭕️ | ❌ | ❌ | 
 | 24GB | ✅ | ✅ | ✅ | 🧡 | 🧡 | ❌ | 
 | 40GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 | 80GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 
-
 ✅ indicates it can run under "model_cpu_offload", 🧡 represents it can run under "model_cpu_offload_and_qfloat8", ⭕️ indicates it can run under "sequential_cpu_offload", ❌ means it can't run. Please note that running with sequential_cpu_offload will be slower.
 
 Some GPUs that do not support torch.bfloat16, such as 2080ti and V100, require changing the weight_dtype in app.py and predict files to torch.float16 in order to run.
@@ -501,6 +498,14 @@ For details on setting some parameters, please refer to [Readme Train](scripts/R
 
 EasyAnimateV5.1:
 
+7B:
+| Name | Type | Storage Space | Hugging Face | Model Scope | Description |
+|--|--|--|--|--|--|
+| EasyAnimateV5.1-7b-zh-InP | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-InP) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-InP) | Official image-to-video weights. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+| EasyAnimateV5.1-7b-zh-Control | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control) | Official video control weights, supporting various control conditions such as Canny, Depth, Pose, MLSD, and trajectory control. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+| EasyAnimateV5.1-7b-zh-Control-Camera | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control-Camera) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control-Camera) | Official video camera control weights, supporting direction generation control by inputting camera motion trajectories. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+| EasyAnimateV5.1-7b-zh | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh) | Official text-to-video weights. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+
 12B:
 | Name | Type | Storage Space | Hugging Face | Model Scope | Description |
 |--|--|--|--|--|--|

README_ja-JP.md

@@ -118,17 +118,15 @@ Linuxの詳細：
 EasyAnimateV5.1-12Bのビデオサイズは異なるGPUメモリにより生成できます。以下の表をご覧ください：
 | GPUメモリ |384x672x25|384x672x49|576x1008x25|576x1008x49|768x1344x25|768x1344x49|
 |----------|----------|----------|----------|----------|----------|----------|
-| 16GB | 🧡 | 🧡 | ❌ | ❌ | ❌ | ❌ | 
+| 16GB | 🧡 | ⭕️ | ⭕️ | ⭕️ | ❌ | ❌ | 
 | 24GB | 🧡 | 🧡 | 🧡 | 🧡 | 🧡 | ❌ | 
 | 40GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 | 80GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 
-qwen2-vl-7bのfloat16の重みのため、16GBのVRAMでは実行できません。もしお使いのVRAMが16GBである場合は、[Huggingface](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct-GPTQ-
-
-EasyAnimateV5-7Bのビデオサイズは異なるGPUメモリにより生成できます。以下の表をご覧ください：
+EasyAnimateV5.1-7Bのビデオサイズは異なるGPUメモリにより生成できます。以下の表をご覧ください：
 | GPU memory |384x672x25|384x672x49|576x1008x25|576x1008x49|768x1344x25|768x1344x49|
 |----------|----------|----------|----------|----------|----------|----------|
-| 16GB | 🧡 | 🧡 | ❌ | ❌ | ❌ | ❌ | 
+| 16GB | 🧡 | 🧡 | ⭕️ | ⭕️ | ❌ | ❌ | 
 | 24GB | ✅ | ✅ | ✅ | 🧡 | 🧡 | ❌ | 
 | 40GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 | 80GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 

README_zh-CN.md

@@ -115,17 +115,15 @@ Linux 的详细信息：
 EasyAnimateV5.1-12B的视频大小可以由不同的GPU Memory生成，包括：
 | GPU memory |384x672x25|384x672x49|576x1008x25|576x1008x49|768x1344x25|768x1344x49|
 |----------|----------|----------|----------|----------|----------|----------|
-| 16GB | 🧡 | 🧡 | ❌ | ❌ | ❌ | ❌ | 
+| 16GB | 🧡 | ⭕️ | ⭕️ | ⭕️ | ❌ | ❌ | 
 | 24GB | 🧡 | 🧡 | 🧡 | 🧡 | 🧡 | ❌ | 
 | 40GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 | 80GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 
-由于qwen2-vl-7b的float16的权重，无法在16GB显存下运行，如果您的显存是16GB，请前往[Huggingface](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct-GPTQ-Int8)或者[Modelscope](https://modelscope.cn/models/Qwen/Qwen2-VL-7B-Instruct-GPTQ-Int8)下载量化后的qwen2-vl-7b对原有的text encoder进行替换，并安装对应的依赖库（auto-gptq, optimum）。
-
-EasyAnimateV5-7B的视频大小可以由不同的GPU Memory生成，包括：
+EasyAnimateV5.1-7B的视频大小可以由不同的GPU Memory生成，包括：
 | GPU memory |384x672x25|384x672x49|576x1008x25|576x1008x49|768x1344x25|768x1344x49|
 |----------|----------|----------|----------|----------|----------|----------|
-| 16GB | 🧡 | 🧡 | ❌ | ❌ | ❌ | ❌ | 
+| 16GB | 🧡 | 🧡 | ⭕️ | ⭕️ | ❌ | ❌ | 
 | 24GB | ✅ | ✅ | ✅ | 🧡 | 🧡 | ❌ | 
 | 40GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
 | 80GB | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 
@@ -495,6 +493,14 @@ sh scripts/train.sh
 # 模型地址
 EasyAnimateV5.1:
 
+7B:
+| 名称 | 种类 | 存储空间 | Hugging Face | Model Scope | 描述 |
+|--|--|--|--|--|--|
+| EasyAnimateV5.1-7b-zh-InP | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-InP) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-InP)| 官方的图生视频权重。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+| EasyAnimateV5.1-7b-zh-Control | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control)| 官方的视频控制权重，支持不同的控制条件，如Canny、Depth、Pose、MLSD等，同时支持使用轨迹控制。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+| EasyAnimateV5.1-7b-zh-Control-Camera | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control-Camera) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control-Camera)| 官方的视频相机控制权重，支持通过输入相机运动轨迹控制生成方向。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+| EasyAnimateV5.1-7b-zh | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh)| 官方的文生视频权重。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+
 12B:
 | 名称 | 种类 | 存储空间 | Hugging Face | Model Scope | 描述 |
 |--|--|--|--|--|--|

app.py

@@ -21,9 +21,13 @@ if __name__ == "__main__":
     # resulting in slower speeds but saving a large amount of GPU memory.
     # 
     # EasyAnimateV1, V2 and V3 support "model_cpu_offload" "sequential_cpu_offload"
-    # 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" 
+    # EasyAnimateV4, V5 and V5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload"
     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.08
     # 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 +48,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(

comfyui/README.md

@@ -38,6 +38,14 @@ pip install -r comfyui/requirements.txt
 
 EasyAnimateV5.1:
 
+7B:
+| Name | Type | Storage Space | Hugging Face | Model Scope | Description |
+|--|--|--|--|--|--|
+| EasyAnimateV5.1-7b-zh-InP | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-InP) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-InP) | Official image-to-video weights. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+| EasyAnimateV5.1-7b-zh-Control | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control) | Official video control weights, supporting various control conditions such as Canny, Depth, Pose, MLSD, and trajectory control. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+| EasyAnimateV5.1-7b-zh-Control-Camera | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control-Camera) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control-Camera) | Official video camera control weights, supporting direction generation control by inputting camera motion trajectories. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+| EasyAnimateV5.1-7b-zh | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh) | Official text-to-video weights. Supports video prediction at multiple resolutions (512, 768, 1024), trained with 49 frames at 8 frames per second, and supports for multilingual prediction. |
+
 12B:
 | Name | Type | Storage Space | Hugging Face | Model Scope | Description |
 |--|--|--|--|--|--|

comfyui/README_zh-CN.md

@@ -36,6 +36,15 @@ pip install -r comfyui/requirements.txt
 ## 将模型下载到`ComfyUI/models/EasyAnimate/`
 
 EasyAnimateV5.1：
+
+7B:
+| 名称 | 种类 | 存储空间 | Hugging Face | Model Scope | 描述 |
+|--|--|--|--|--|--|
+| EasyAnimateV5.1-7b-zh-InP | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-InP) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-InP)| 官方的图生视频权重。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+| EasyAnimateV5.1-7b-zh-Control | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control)| 官方的视频控制权重，支持不同的控制条件，如Canny、Depth、Pose、MLSD等，同时支持使用轨迹控制。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+| EasyAnimateV5.1-7b-zh-Control-Camera | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh-Control-Camera) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh-Control-Camera)| 官方的视频相机控制权重，支持通过输入相机运动轨迹控制生成方向。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+| EasyAnimateV5.1-7b-zh | EasyAnimateV5.1 | 30 GB | [🤗Link](https://huggingface.co/alibaba-pai/EasyAnimateV5.1-7b-zh) | [😄Link](https://modelscope.cn/models/PAI/EasyAnimateV5.1-7b-zh)| 官方的文生视频权重。支持多分辨率（512，768，1024）的视频预测，支持多分辨率（512，768，1024）的视频预测，以49帧、每秒8帧进行训练，支持多语言预测 |
+
 12B：
 |名称|类型|存储空间|拥抱面|型号范围|描述|
 |--|--|--|--|--|--|

comfyui/comfyui_nodes.py

@@ -39,7 +39,8 @@ from ..easyanimate.pipeline.pipeline_easyanimate_control import \
 from ..easyanimate.utils.lora_utils import merge_lora, unmerge_lora
 from ..easyanimate.utils.utils import (get_image_to_video_latent, get_image_latent,
                                        get_video_to_video_latent)
-from ..easyanimate.utils.fp8_optimization import convert_weight_dtype_wrapper
+from ..easyanimate.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                                convert_weight_dtype_wrapper)
 from ..easyanimate.ui.ui import ddpm_scheduler_dict, flow_scheduler_dict, all_cheduler_dict
 
 # Compatible with Alibaba EAS for quick launch
@@ -98,6 +99,11 @@ class LoadEasyAnimateModel:
                         'EasyAnimateV5-12b-zh-InP',
                         'EasyAnimateV5-12b-zh-Control',
                         'EasyAnimateV5-12b-zh',
+                        'EasyAnimateV5.1-7b-zh',
+                        'EasyAnimateV5.1-7b-zh-InP',
+                        'EasyAnimateV5.1-7b-zh-Control',
+                        'EasyAnimateV5.1-7b-zh-Control-Camera',
+                        'EasyAnimateV5.1-12b-zh',
                         'EasyAnimateV5.1-12b-zh-InP',
                         'EasyAnimateV5.1-12b-zh-Control',
                         'EasyAnimateV5.1-12b-zh-Control-Camera',
@@ -174,7 +180,7 @@ class LoadEasyAnimateModel:
             model_name, 
             subfolder="vae"
         ).to(weight_dtype)
-        if config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and weight_dtype == torch.float16:
+        if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
             vae.upcast_vae = True
         # Update pbar
         pbar.update(1)
@@ -185,7 +191,7 @@ class LoadEasyAnimateModel:
         ]
 
         transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
-        if weight_dtype == torch.float16:
+        if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
             transformer_additional_kwargs["upcast_attention"] = True
 
         transformer = Choosen_Transformer3DModel.from_pretrained_2d(
@@ -299,11 +305,23 @@ class LoadEasyAnimateModel:
                 transformer=transformer,
                 scheduler=scheduler,
             )
+
         if GPU_memory_mode == "sequential_cpu_offload":
+            pipeline._manual_cpu_offload_in_sequential_cpu_offload = []
+            for name, _text_encoder in zip(["text_encoder", "text_encoder_2"], [pipeline.text_encoder, pipeline.text_encoder_2]):
+                if isinstance(_text_encoder, Qwen2VLForConditionalGeneration):
+                    if hasattr(_text_encoder, "visual"):
+                        del _text_encoder.visual
+                    convert_model_weight_to_float8(_text_encoder)
+                    convert_weight_dtype_wrapper(_text_encoder, weight_dtype)
+                    pipeline._manual_cpu_offload_in_sequential_cpu_offload = [name]
             pipeline.enable_sequential_cpu_offload()
         elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
-            pipeline.enable_model_cpu_offload()
+            for _text_encoder in [pipeline.text_encoder, pipeline.text_encoder_2]:
+                if hasattr(_text_encoder, "visual"):
+                    del _text_encoder.visual
             convert_weight_dtype_wrapper(transformer, weight_dtype)
+            pipeline.enable_model_cpu_offload()
         else:
             pipeline.enable_model_cpu_offload()
         easyanimate_model = {
@@ -442,7 +460,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 +477,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 +609,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 +672,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 +697,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 +806,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 +877,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 +920,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 +1064,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",),

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.08,
+        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

comfyui/v5.1/easyanimatev5.1_workflow_control_trajectory.json

@@ -140,7 +140,7 @@
         "Node name for S&R": "EasyAnimate_TextBox"
       },
       "widgets_values": [
-        "一只棕褐色的狗在摇晃脑袋，坐在一个舒适的房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里的灯光柔和温暖，营造出舒适的氛围。"
+        "一只棕褐色的狗正摇晃着脑袋，坐在一个舒适的房间里的浅色沙发上。沙发看起来柔软而宽敞，为这只活泼的狗狗提供了一个完美的休息地点。在狗的后面，靠墙摆放着一个架子，架子上挂着一幅精美的镶框画，画中描绘着一些美丽的风景或场景。画框周围装饰着粉红色的花朵，这些花朵不仅增添了房间的色彩，还带来了一丝自然和生机。房间里的灯光柔和而温暖，从天花板上的吊灯和角落里的台灯散发出来，营造出一种温馨舒适的氛围。整个空间给人一种宁静和谐的感觉，仿佛时间在这里变得缓慢而美好。"
       ]
     },
     {
@@ -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.08,
+        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

comfyui/v5.1/easyanimatev5.1_workflow_i2v.json

@@ -289,7 +289,9 @@
         "fixed",
         50,
         6,
-        "Flow"
+        "Flow",
+        0.08,
+        true
       ]
     },
     {

comfyui/v5.1/easyanimatev5.1_workflow_t2v.json

@@ -84,7 +84,9 @@
         "fixed",
         50,
         6,
-        "Flow"
+        "Flow",
+        0.08,
+        true
       ]
     },
     {
@@ -319,7 +321,7 @@
         "Node name for S&R": "EasyAnimate_TextBox"
       },
       "widgets_values": [
-        "一只棕褐色的狗在摇晃脑袋，坐在一个舒适的房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里的灯光柔和温暖，营造出舒适的氛围。"
+        "一只棕褐色的狗正摇晃着脑袋，坐在一个舒适的房间里的浅色沙发上。沙发看起来柔软而宽敞，为这只活泼的狗狗提供了一个完美的休息地点。在狗的后面，靠墙摆放着一个架子，架子上挂着一幅精美的镶框画，画中描绘着一些美丽的风景或场景。画框周围装饰着粉红色的花朵，这些花朵不仅增添了房间的色彩，还带来了一丝自然和生机。房间里的灯光柔和而温暖，从天花板上的吊灯和角落里的台灯散发出来，营造出一种温馨舒适的氛围。整个空间给人一种宁静和谐的感觉，仿佛时间在这里变得缓慢而美好。"
       ]
     }
   ],

comfyui/v5.1/easyanimatev5.1_workflow_v2v.json

@@ -117,7 +117,7 @@
         "Node name for S&R": "EasyAnimate_TextBox"
       },
       "widgets_values": [
-        "一只穿着小外套的猫咪正在花园秋千上安静地弹吉他。晚霞的余光洒在它柔软的毛皮上，和煦的微风轻轻拂过，周围斑驳的光影随着音乐的旋律轻轻摇曳。"
+        "一只穿着小外套的猫咪正安静地坐在花园的秋千上弹吉他。它的小外套精致而合身，增添了几分俏皮与可爱。晚霞的余光洒在它柔软的毛皮上，给它的毛发镀上了一层温暖的金色光辉。和煦的微风轻轻拂过，带来阵阵花香和草木的气息，令人心旷神怡。周围斑驳的光影随着音乐的旋律轻轻摇曳，仿佛整个花园都在为这只小猫咪的演奏伴舞。阳光透过树叶间的缝隙，投下一片片光影交错的图案，与悠扬的吉他声交织在一起，营造出一种梦幻而宁静的氛围。猫咪专注而投入地弹奏着，每一个音符都似乎充满了魔力，让这个傍晚变得更加美好。"
       ]
     },
     {
@@ -305,6 +305,8 @@
         6,
         0.7000000000000001,
         "Flow",
+        0.08,
+        true,
         ""
       ]
     },

comfyui/v5.1/easyanimatev5.1_workflow_v2v_control.json

@@ -117,7 +117,7 @@
         "Node name for S&R": "EasyAnimate_TextBox"
       },
       "widgets_values": [
-        "一个穿着及膝白色无袖连衣裙和白色高跟凉鞋的美女在一个光线充足、木地板的房间里跳舞。房间的背景是一扇紧闭的门、一个展示透明玻璃瓶酒精饮料的架子和一个部分可见的深色沙发。"
+        "在这个阳光明媚的户外花园里，美女身穿一袭及膝的白色无袖连衣裙，裙摆在她轻盈的舞姿中轻柔地摆动，宛如一只翩翩起舞的蝴蝶。阳光透过树叶间洒下斑驳的光影，映衬出她柔和的脸庞和清澈的眼眸，显得格外优雅。仿佛每一个动作都在诉说着青春与活力，她在草地上旋转，裙摆随之飞扬，仿佛整个花园都因她的舞动而欢愉。周围五彩缤纷的花朵在微风中摇曳，玫瑰、菊花、百合，各自释放出阵阵香气，营造出一种轻松而愉快的氛围。"
       ]
     },
     {
@@ -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.08,
+        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

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,
+        ""
       ]
     },
     {

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,
+        ""
       ]
     },
     {

easyanimate/models/processor.py

@@ -318,8 +318,8 @@ except:
     print("Flash Attention is not installed. Please install with `pip install flash-attn`, if you want to use SWA.")
 
 class EasyAnimateSWAttnProcessor2_0:
-    def __init__(self, window_size=1024):
-        self.window_size = window_size
+    def __init__(self, cross_attention_size=1024):
+        self.cross_attention_size = cross_attention_size
 
     def __call__(
         self,
@@ -334,6 +334,7 @@ class EasyAnimateSWAttnProcessor2_0:
         attn2: Attention = None,
     ) -> torch.Tensor:
         text_seq_length = encoder_hidden_states.size(1)
+        windows_size = height * width
 
         batch_size, sequence_length, _ = (
             hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
@@ -387,7 +388,7 @@ class EasyAnimateSWAttnProcessor2_0:
             
         query = query.transpose(1, 2).to(value)
         key = key.transpose(1, 2).to(value)
-        interval = max((query.size(1) - text_seq_length) // (self.window_size - text_seq_length), 1)
+        interval = max((query.size(1) - text_seq_length) // (self.cross_attention_size - text_seq_length), 1)
 
         cross_key = torch.cat([key[:, :text_seq_length], key[:, text_seq_length::interval]], dim=1)
         cross_val = torch.cat([value[:, :text_seq_length], value[:, text_seq_length::interval]], dim=1)
@@ -418,8 +419,8 @@ class EasyAnimateSWAttnProcessor2_0:
         value = torch.cat(new_values, dim=2)
         
         # apply attention
-        hidden_states = flash_attn_func(query, key, value, dropout_p=0.0, causal=False, window_size=(self.window_size, self.window_size))
-        
+        hidden_states = flash_attn_func(query, key, value, dropout_p=0.0, causal=False, window_size=(windows_size, windows_size))
+
         hidden_states = torch.tensor_split(hidden_states, 6, 2)
         new_hidden_states = [hidden_states[0]]
         for index, mode in enumerate(

easyanimate/models/transformer3d.py

@@ -87,6 +87,56 @@ 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
+
+
+def get_teacache_coefficients(model_name):
+    # The coefficients for EasyAnimateV5-7b-zh-InP should be:
+    # [-3.64204720e+03, 1.43764725e+03, -1.93045263e+02, 1.09596499e+01, -1.70663507e-01]
+    if "v5.1-7b" in model_name.lower():
+        # The coefficient was obtained by sampling videos from T2V CompBench using EasyAnimateV5.1-7b-zh-InP.
+        # This coefficient can be applied to both the EasyAnimateV5.1-7b-zh and EasyAnimateV5.1-7b-Control.
+        return [1.07862322, -4.19362456, 3.06725828, 0.33161686, 0.02374758]
+    elif "v5.1-12b" in model_name.lower():
+        # 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.
+        return [-10.47857366, 8.33844143, -0.78477557, 0.68798618, 0.0136149]
+    else:
+        print(f"The model {model_name} is not supported by TeaCache.")
+        return None
+
+
 class Transformer3DModel(ModelMixin, ConfigMixin):
     """
     A 3D Transformer model for image-like data.
@@ -1428,7 +1478,17 @@ 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]
+    ):
+        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 +1559,124 @@ 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.to(modulated_inp.device), 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.cpu()
+            self.teacache.cnt += 1
+            if self.teacache.cnt == self.teacache.num_steps:
+                # self.cnt = 0
+                self.teacache.reset()
+            del inp, temb_, encoder_hidden_states_
 
-        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.to(modulated_inp.device)
+            else:
+                ori_hidden_states = hidden_states.clone().cpu()
+
+                # 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.cpu() - ori_hidden_states
+                del 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

easyanimate/pipeline/pipeline_easyanimate.py

@@ -240,14 +240,69 @@ class EasyAnimatePipeline(DiffusionPipeline):
         )
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.manual_cpu_offload_flag = False
+
+    def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
+        from diffusers.pipelines.pipeline_utils import is_accelerate_available, is_accelerate_version
+
+        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
+            from accelerate import cpu_offload
+            from accelerate import cpu_offload_with_hook
+        else:
+            raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")
+        self.remove_all_hooks()
+
+        is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1
+        if is_pipeline_device_mapped:
+            raise ValueError(
+                "It seems like you have activated a device mapping strategy on the pipeline so calling `enable_sequential_cpu_offload() isn't allowed. You can call `reset_device_map()` first and then call `enable_sequential_cpu_offload()`."
+            )
+
+        torch_device = torch.device(device)
+        device_index = torch_device.index
+
+        if gpu_id is not None and device_index is not None:
+            raise ValueError(
+                f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}"
+                f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}"
+            )
+
+        # _offload_gpu_id should be set to passed gpu_id (or id in passed `device`) or default to previously set id or default to 0
+        self._offload_gpu_id = gpu_id or torch_device.index or getattr(self, "_offload_gpu_id", 0)
+
+        device_type = torch_device.type
+        device = torch.device(f"{device_type}:{self._offload_gpu_id}")
+        self._offload_device = device
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            device_mod = getattr(torch, self.device.type, None)
+            if hasattr(device_mod, "empty_cache") and device_mod.is_available():
+                device_mod.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        for name, model in self.components.items():
+            if not isinstance(model, torch.nn.Module):
+                continue
+
+            if name in self._manual_cpu_offload_in_sequential_cpu_offload:
+                pass
+            else:
+                # make sure to offload buffers if not all high level weights
+                # are of type nn.Module
+                offload_buffers = len(model._parameters) > 0
+                cpu_offload(model, device, offload_buffers=offload_buffers)
 
-    def enable_sequential_cpu_offload(self, *args, **kwargs):
-        super().enable_sequential_cpu_offload(*args, **kwargs)
         if hasattr(self.transformer, "clip_projection") and self.transformer.clip_projection is not None:
             import accelerate
             accelerate.hooks.remove_hook_from_module(self.transformer.clip_projection, recurse=True)
             self.transformer.clip_projection = self.transformer.clip_projection.to("cuda")
 
+        self.manual_cpu_offload_flag = True
+
+    def enable_model_cpu_offload(self, *args, **kwargs):
+        super().enable_model_cpu_offload(*args, **kwargs)
+        self.manual_cpu_offload_flag = True
+
     def encode_prompt(
         self,
         prompt: str,
@@ -855,6 +910,12 @@ class EasyAnimatePipeline(DiffusionPipeline):
         else:
             dtype = self.transformer.dtype
 
+        if self.manual_cpu_offload_flag:
+            if isinstance(self.text_encoder, Qwen2VLForConditionalGeneration):
+                self.text_encoder.to(device)
+            if isinstance(self.text_encoder_2, Qwen2VLForConditionalGeneration) and self.text_encoder_2 is not None:
+                self.text_encoder_2.to(device)
+
         # 3. Encode input prompt
         (
             prompt_embeds,
@@ -899,6 +960,13 @@ class EasyAnimatePipeline(DiffusionPipeline):
             prompt_attention_mask_2 = None
             negative_prompt_attention_mask_2 = None
 
+        if self.manual_cpu_offload_flag:
+            if isinstance(self.text_encoder, Qwen2VLForConditionalGeneration):
+                self.text_encoder.to("cpu")
+            if isinstance(self.text_encoder_2, Qwen2VLForConditionalGeneration) and self.text_encoder_2 is not None:
+                self.text_encoder_2.to("cpu")
+            torch.cuda.empty_cache()
+
         # 4. Prepare timesteps
         if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
             timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)

easyanimate/pipeline/pipeline_easyanimate_control.py

@@ -269,14 +269,69 @@ class EasyAnimateControlPipeline(DiffusionPipeline):
         self.mask_processor = VaeImageProcessor(
             vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True
         )
+        self.manual_cpu_offload_flag = False
+
+    def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
+        from diffusers.pipelines.pipeline_utils import is_accelerate_available, is_accelerate_version
+
+        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
+            from accelerate import cpu_offload
+            from accelerate import cpu_offload_with_hook
+        else:
+            raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")
+        self.remove_all_hooks()
+
+        is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1
+        if is_pipeline_device_mapped:
+            raise ValueError(
+                "It seems like you have activated a device mapping strategy on the pipeline so calling `enable_sequential_cpu_offload() isn't allowed. You can call `reset_device_map()` first and then call `enable_sequential_cpu_offload()`."
+            )
+
+        torch_device = torch.device(device)
+        device_index = torch_device.index
+
+        if gpu_id is not None and device_index is not None:
+            raise ValueError(
+                f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}"
+                f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}"
+            )
+
+        # _offload_gpu_id should be set to passed gpu_id (or id in passed `device`) or default to previously set id or default to 0
+        self._offload_gpu_id = gpu_id or torch_device.index or getattr(self, "_offload_gpu_id", 0)
+
+        device_type = torch_device.type
+        device = torch.device(f"{device_type}:{self._offload_gpu_id}")
+        self._offload_device = device
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            device_mod = getattr(torch, self.device.type, None)
+            if hasattr(device_mod, "empty_cache") and device_mod.is_available():
+                device_mod.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        for name, model in self.components.items():
+            if not isinstance(model, torch.nn.Module):
+                continue
+
+            if name in self._manual_cpu_offload_in_sequential_cpu_offload:
+                pass
+            else:
+                # make sure to offload buffers if not all high level weights
+                # are of type nn.Module
+                offload_buffers = len(model._parameters) > 0
+                cpu_offload(model, device, offload_buffers=offload_buffers)
 
-    def enable_sequential_cpu_offload(self, *args, **kwargs):
-        super().enable_sequential_cpu_offload(*args, **kwargs)
         if hasattr(self.transformer, "clip_projection") and self.transformer.clip_projection is not None:
             import accelerate
             accelerate.hooks.remove_hook_from_module(self.transformer.clip_projection, recurse=True)
             self.transformer.clip_projection = self.transformer.clip_projection.to("cuda")
 
+        self.manual_cpu_offload_flag = True
+
+    def enable_model_cpu_offload(self, *args, **kwargs):
+        super().enable_model_cpu_offload(*args, **kwargs)
+        self.manual_cpu_offload_flag = True
+
     def encode_prompt(
         self,
         prompt: str,
@@ -922,6 +977,12 @@ class EasyAnimateControlPipeline(DiffusionPipeline):
         else:
             dtype = self.transformer.dtype
 
+        if self.manual_cpu_offload_flag:
+            if isinstance(self.text_encoder, Qwen2VLForConditionalGeneration):
+                self.text_encoder.to(device)
+            if isinstance(self.text_encoder_2, Qwen2VLForConditionalGeneration) and self.text_encoder_2 is not None:
+                self.text_encoder_2.to(device)
+
         # 3. Encode input prompt
         (
             prompt_embeds,
@@ -966,6 +1027,13 @@ class EasyAnimateControlPipeline(DiffusionPipeline):
             prompt_attention_mask_2 = None
             negative_prompt_attention_mask_2 = None
 
+        if self.manual_cpu_offload_flag:
+            if isinstance(self.text_encoder, Qwen2VLForConditionalGeneration):
+                self.text_encoder.to("cpu")
+            if isinstance(self.text_encoder_2, Qwen2VLForConditionalGeneration) and self.text_encoder_2 is not None:
+                self.text_encoder_2.to("cpu")
+            torch.cuda.empty_cache()
+    
         # 4. Prepare timesteps
         if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
             timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)

easyanimate/pipeline/pipeline_easyanimate_inpaint.py

@@ -150,14 +150,18 @@ def resize_mask(mask, latent, process_first_frame_only=True):
 
 
 ## Add noise to reference video
-def add_noise_to_reference_video(image, ratio=None):
+def add_noise_to_reference_video(image, ratio=None, generator=None):
     if ratio is None:
         sigma = torch.normal(mean=-3.0, std=0.5, size=(image.shape[0],)).to(image.device)
         sigma = torch.exp(sigma).to(image.dtype)
     else:
         sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * ratio
     
-    image_noise = torch.randn_like(image) * sigma[:, None, None, None, None]
+    if generator is not None:
+        image_noise = torch.randn(image.size(), generator=generator, dtype=image.dtype, device=image.device) * \
+            sigma[:, None, None, None, None]
+    else:
+        image_noise = torch.randn_like(image) * sigma[:, None, None, None, None]
     image_noise = torch.where(image==-1, torch.zeros_like(image), image_noise)
     image = image + image_noise
     return image
@@ -319,14 +323,69 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
         self.mask_processor = VaeImageProcessor(
             vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True
         )
+        self.manual_cpu_offload_flag = False
+
+    def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):
+        from diffusers.pipelines.pipeline_utils import is_accelerate_available, is_accelerate_version
+
+        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
+            from accelerate import cpu_offload
+            from accelerate import cpu_offload_with_hook
+        else:
+            raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")
+        self.remove_all_hooks()
+
+        is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1
+        if is_pipeline_device_mapped:
+            raise ValueError(
+                "It seems like you have activated a device mapping strategy on the pipeline so calling `enable_sequential_cpu_offload() isn't allowed. You can call `reset_device_map()` first and then call `enable_sequential_cpu_offload()`."
+            )
+
+        torch_device = torch.device(device)
+        device_index = torch_device.index
+
+        if gpu_id is not None and device_index is not None:
+            raise ValueError(
+                f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}"
+                f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}"
+            )
+
+        # _offload_gpu_id should be set to passed gpu_id (or id in passed `device`) or default to previously set id or default to 0
+        self._offload_gpu_id = gpu_id or torch_device.index or getattr(self, "_offload_gpu_id", 0)
+
+        device_type = torch_device.type
+        device = torch.device(f"{device_type}:{self._offload_gpu_id}")
+        self._offload_device = device
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            device_mod = getattr(torch, self.device.type, None)
+            if hasattr(device_mod, "empty_cache") and device_mod.is_available():
+                device_mod.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        for name, model in self.components.items():
+            if not isinstance(model, torch.nn.Module):
+                continue
+
+            if name in self._manual_cpu_offload_in_sequential_cpu_offload:
+                pass
+            else:
+                # make sure to offload buffers if not all high level weights
+                # are of type nn.Module
+                offload_buffers = len(model._parameters) > 0
+                cpu_offload(model, device, offload_buffers=offload_buffers)
 
-    def enable_sequential_cpu_offload(self, *args, **kwargs):
-        super().enable_sequential_cpu_offload(*args, **kwargs)
         if hasattr(self.transformer, "clip_projection") and self.transformer.clip_projection is not None:
             import accelerate
             accelerate.hooks.remove_hook_from_module(self.transformer.clip_projection, recurse=True)
             self.transformer.clip_projection = self.transformer.clip_projection.to("cuda")
 
+        self.manual_cpu_offload_flag = True
+
+    def enable_model_cpu_offload(self, *args, **kwargs):
+        super().enable_model_cpu_offload(*args, **kwargs)
+        self.manual_cpu_offload_flag = True
+
     def encode_prompt(
         self,
         prompt: str,
@@ -738,7 +797,7 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
         if masked_image is not None:
             masked_image = masked_image.to(device=device, dtype=dtype)
             if self.transformer.config.add_noise_in_inpaint_model:
-                masked_image = add_noise_to_reference_video(masked_image, ratio=noise_aug_strength)
+                masked_image = add_noise_to_reference_video(masked_image, ratio=noise_aug_strength, generator=generator)
             if self.vae.quant_conv is None or self.vae.quant_conv.weight.ndim==5:
                 bs = 1
                 new_mask_pixel_values = []
@@ -809,7 +868,7 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
                 for i in range(0, video.shape[0], bs):
                     video_bs = video[i : i + bs]
                     video_bs = self.vae.encode(video_bs)[0]
-                    video_bs = video_bs.sample()
+                    video_bs = video_bs.mode()
                     new_video.append(video_bs)
                 video = torch.cat(new_video, dim = 0)
                 video = video * self.vae.config.scaling_factor
@@ -1098,7 +1157,13 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
             dtype = self.text_encoder_2.dtype
         else:
             dtype = self.transformer.dtype
-            
+
+        if self.manual_cpu_offload_flag:
+            if isinstance(self.text_encoder, Qwen2VLForConditionalGeneration):
+                self.text_encoder.to(device)
+            if isinstance(self.text_encoder_2, Qwen2VLForConditionalGeneration) and self.text_encoder_2 is not None:
+                self.text_encoder_2.to(device)
+
         # 3. Encode input prompt
         (
             prompt_embeds,
@@ -1143,6 +1208,13 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
             prompt_attention_mask_2 = None
             negative_prompt_attention_mask_2 = None
 
+        if self.manual_cpu_offload_flag:
+            if isinstance(self.text_encoder, Qwen2VLForConditionalGeneration):
+                self.text_encoder.to("cpu")
+            if isinstance(self.text_encoder_2, Qwen2VLForConditionalGeneration) and self.text_encoder_2 is not None:
+                self.text_encoder_2.to("cpu")
+            torch.cuda.empty_cache()
+
         # 4. set timesteps
         if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
             timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)

easyanimate/ui/ui.py

@@ -28,19 +28,18 @@ from transformers import (BertModel, BertTokenizer, CLIPImageProcessor,
                           T5Tokenizer)
 
 from ..data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio
-from ..models import (name_to_autoencoder_magvit,
-                                name_to_transformer3d)
-from ..pipeline.pipeline_easyanimate import \
-    EasyAnimatePipeline
-from ..pipeline.pipeline_easyanimate_control import \
-    EasyAnimateControlPipeline
-from ..pipeline.pipeline_easyanimate_inpaint import \
-    EasyAnimateInpaintPipeline
-from ..utils.fp8_optimization import convert_weight_dtype_wrapper
+from ..models import name_to_autoencoder_magvit, name_to_transformer3d
+from ..models.transformer3d import get_teacache_coefficients
+from ..pipeline.pipeline_easyanimate import EasyAnimatePipeline
+from ..pipeline.pipeline_easyanimate_control import EasyAnimateControlPipeline
+from ..pipeline.pipeline_easyanimate_inpaint import EasyAnimateInpaintPipeline
+from ..utils.fp8_optimization import (convert_model_weight_to_float8,
+                                      convert_weight_dtype_wrapper)
 from ..utils.lora_utils import merge_lora, unmerge_lora
-from ..utils.utils import (
-    get_image_to_video_latent, get_video_to_video_latent,
-    get_width_and_height_from_image_and_base_resolution, save_videos_grid)
+from ..utils.utils import (get_image_to_video_latent,
+                           get_video_to_video_latent,
+                           get_width_and_height_from_image_and_base_resolution,
+                           save_videos_grid)
 
 ddpm_scheduler_dict = {
     "Euler": EulerDiscreteScheduler,
@@ -67,7 +66,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 +96,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"
@@ -167,11 +168,11 @@ class EasyAnimateController:
             diffusion_transformer_dropdown, 
             subfolder="vae", 
         ).to(self.weight_dtype)
-        if self.inference_config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and self.weight_dtype == torch.float16:
+        if self.weight_dtype == torch.float16 and "v5.1" not in diffusion_transformer_dropdown.lower():
             self.vae.upcast_vae = True
             
         transformer_additional_kwargs = OmegaConf.to_container(self.inference_config['transformer_additional_kwargs'])
-        if self.weight_dtype == torch.float16:
+        if self.weight_dtype == torch.float16 and "v5.1" not in diffusion_transformer_dropdown.lower():
             transformer_additional_kwargs["upcast_attention"] = True
 
         # Get Transformer
@@ -292,8 +293,19 @@ class EasyAnimateController:
             )
 
         if self.GPU_memory_mode == "sequential_cpu_offload":
+            self.pipeline._manual_cpu_offload_in_sequential_cpu_offload = []
+            for name, _text_encoder in zip(["text_encoder", "text_encoder_2"], [self.pipeline.text_encoder, self.pipeline.text_encoder_2]):
+                if isinstance(_text_encoder, Qwen2VLForConditionalGeneration):
+                    if hasattr(_text_encoder, "visual"):
+                        del _text_encoder.visual
+                    convert_model_weight_to_float8(_text_encoder)
+                    convert_weight_dtype_wrapper(_text_encoder, self.weight_dtype)
+                    self.pipeline._manual_cpu_offload_in_sequential_cpu_offload = [name]
             self.pipeline.enable_sequential_cpu_offload()
         elif self.GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+            for _text_encoder in [self.pipeline.text_encoder, self.pipeline.text_encoder_2]:
+                if hasattr(_text_encoder, "visual"):
+                    del _text_encoder.visual
             self.pipeline.enable_model_cpu_offload()
             convert_weight_dtype_wrapper(self.pipeline.transformer, self.weight_dtype)
         else:
@@ -462,6 +474,11 @@ class EasyAnimateController:
             # lora part
             self.pipeline = merge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider)
         
+        coefficients = get_teacache_coefficients(self.base_model_path)
+        if coefficients is not None and self.enable_teacache:
+            print(f"Enable TeaCache with threshold: {self.teacache_threshold}.")
+            self.pipeline.transformer.enable_teacache(sample_step_slider, self.teacache_threshold, coefficients=coefficients)
+        
         try:
             if self.model_type == "Inpaint":
                 if self.transformer.config.in_channels != self.vae.config.latent_channels:
@@ -661,8 +678,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 +1017,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 +1029,9 @@ class EasyAnimateController_Modelscope:
         # Config and model path
         self.model_type = model_type
         self.edition = edition
+        self.model_name = model_name
+        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[
@@ -1021,11 +1041,11 @@ class EasyAnimateController_Modelscope:
             model_name, 
             subfolder="vae", 
         ).to(self.weight_dtype)
-        if self.inference_config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and weight_dtype == torch.float16:
+        if self.weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
             self.vae.upcast_vae = True
             
         transformer_additional_kwargs = OmegaConf.to_container(self.inference_config['transformer_additional_kwargs'])
-        if self.weight_dtype == torch.float16:
+        if self.weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
             transformer_additional_kwargs["upcast_attention"] = True
 
         # Get Transformer
@@ -1144,12 +1164,23 @@ class EasyAnimateController_Modelscope:
             )
 
         if GPU_memory_mode == "sequential_cpu_offload":
+            self.pipeline._manual_cpu_offload_in_sequential_cpu_offload = []
+            for name, _text_encoder in zip(["text_encoder", "text_encoder_2"], [self.pipeline.text_encoder, self.pipeline.text_encoder_2]):
+                if isinstance(_text_encoder, Qwen2VLForConditionalGeneration):
+                    if hasattr(_text_encoder, "visual"):
+                        del _text_encoder.visual
+                    convert_model_weight_to_float8(_text_encoder)
+                    convert_weight_dtype_wrapper(_text_encoder, weight_dtype)
+                    self.pipeline._manual_cpu_offload_in_sequential_cpu_offload = [name]
             self.pipeline.enable_sequential_cpu_offload()
         elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
-            self.pipeline.enable_model_cpu_offload()
+            for _text_encoder in [self.pipeline.text_encoder, self.pipeline.text_encoder_2]:
+                if hasattr(_text_encoder, "visual"):
+                    del _text_encoder.visual
             convert_weight_dtype_wrapper(self.pipeline.transformer, weight_dtype)
+            self.pipeline.enable_model_cpu_offload()
         else:
-            GPU_memory_mode.enable_model_cpu_offload()
+            self.pipeline.enable_model_cpu_offload()
         print("Update diffusion transformer done")
 
     def refresh_personalized_model(self):
@@ -1259,6 +1290,11 @@ class EasyAnimateController_Modelscope:
             # lora part
             self.pipeline = merge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider)
         
+        coefficients = get_teacache_coefficients(self.model_name)
+        if coefficients is not None and self.enable_teacache:
+            print(f"Enable TeaCache with threshold: {self.teacache_threshold}.")
+            self.pipeline.transformer.enable_teacache(sample_step_slider, self.teacache_threshold, coefficients=coefficients)
+        
         try:
             if self.model_type == "Inpaint":
                 if self.vae.cache_mag_vae:
@@ -1373,8 +1409,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(

easyanimate/utils/fp8_optimization.py

@@ -14,9 +14,16 @@ def autocast_model_forward(cls, origin_dtype, *inputs, **kwargs):
     cls.to(weight_dtype)
     return out
 
+def convert_model_weight_to_float8(model, exclude_module_name='embed_tokens'):
+    for name, module in model.named_modules():
+        if exclude_module_name not in name:
+            for param_name, param in module.named_parameters():
+                if exclude_module_name not in param_name:
+                    param.data = param.data.to(torch.float8_e4m3fn)
+
 def convert_weight_dtype_wrapper(module, origin_dtype):
     for name, module in module.named_modules():
-        if name == "":
+        if name == "" or "embed_tokens" in name:
             continue
         original_forward = module.forward
         if hasattr(module, "weight"):

easyanimate/vae/ldm/data/dataset_image_video.py

@@ -172,7 +172,7 @@ class ImageVideoDataset(Dataset):
                 video_reader = VideoReader(example['file_path'])
             video_length = len(video_reader)
             if self.slice_interval == "rand":
-                slice_interval = np.random.choice([1, 2, 3])
+                slice_interval = np.random.choice([1, 2, 3, 4, 5, 6, 7, 8])
             else:
                 slice_interval = int(self.slice_interval)
             clip_length = min(video_length, (self.video_len - 1) * slice_interval + 1)

easyanimate/video_caption/compute_semantic_consistency.py

@@ -162,6 +162,7 @@ def main():
         video_dataset = VideoDataset(
             dataset_inputs={args.video_path_column: splitted_video_path_list},
             video_folder=args.video_folder,
+            video_path_column=args.video_path_column,
             sample_method=args.frame_sample_method,
             num_sampled_frames=args.num_sampled_frames,
             sample_stride=args.sample_stride,

easyanimate/video_caption/compute_video_quality.py

@@ -89,10 +89,10 @@ def main():
             saved_metadata_df = pd.read_json(args.saved_path, lines=True)
 
         # Filter out the unprocessed video-caption pairs by setting the indicator=True.
-        merged_df = video_metadata_df.merge(saved_metadata_df, on="video_path", how="outer", indicator=True)
+        merged_df = video_metadata_df.merge(saved_metadata_df, on=args.video_path_column, how="outer", indicator=True)
         video_metadata_df = merged_df[merged_df["_merge"] == "left_only"]
         # Sorting to guarantee the same result for each process.
-        video_metadata_df = video_metadata_df.iloc[index_natsorted(video_metadata_df["video_path"])].reset_index(drop=True)
+        video_metadata_df = video_metadata_df.iloc[index_natsorted(video_metadata_df[args.video_path_column])].reset_index(drop=True)
         if args.caption_column is None:
             video_metadata_df = video_metadata_df[[args.video_path_column]]
         else:
@@ -160,6 +160,7 @@ def main():
         video_dataset = VideoDataset(
             dataset_inputs=splitted_video_metadata,
             video_folder=args.video_folder,
+            video_path_column=args.video_path_column,
             text_column=args.caption_column,
             sample_method=args.frame_sample_method,
             num_sampled_frames=args.num_sampled_frames

easyanimate/video_caption/filter_meta_train.py

@@ -18,6 +18,12 @@ def parse_args():
         default="video_path",
         help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
     )
+    parser.add_argument(
+        "--caption_column",
+        type=str,
+        default="caption",
+        help="The column contains the caption.",
+    )
     parser.add_argument("--video_folder", type=str, default="", help="The video folder.")
     parser.add_argument(
         "--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl)."
@@ -76,7 +82,7 @@ def main():
     )
     filtered_video_path_list = natsorted(filtered_video_path_list)
     filtered_caption_df = raw_caption_df[raw_caption_df[args.video_path_column].isin(filtered_video_path_list)]
-    train_df = filtered_caption_df.rename(columns={"video_path": "file_path", "caption": "text"})
+    train_df = filtered_caption_df.rename(columns={args.video_path_column: "file_path", args.caption_column: "text"})
     train_df["file_path"] = train_df["file_path"].map(lambda x: os.path.join(args.video_folder, x))
     train_df["type"] = "video"
     train_df.to_json(args.saved_path, orient="records", force_ascii=False, indent=2)

easyanimate/video_caption/utils/get_meta_file.py

@@ -1,30 +1,38 @@
 import argparse
 import os
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from pathlib import Path
 
 import pandas as pd
 from natsort import natsorted
-from tqdm import tqdm
 
 from .logger import logger
 
 ALL_VIDEO_EXT = set(["mp4", "webm", "mkv", "avi", "flv", "mov", "rmvb"])
 ALL_IMGAE_EXT = set(["png", "webp", "jpg", "jpeg", "bmp", "gif"])
 
-def parallel_rglob(root_path, pattern, max_workers=8):
-    root = Path(root_path)
-    futures = []
-    results = []
-    with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        for sub_path in root.iterdir():
-            if sub_path.is_dir():
-                futures.append(executor.submit(lambda p=sub_path: list(p.rglob(pattern))))
-        for future in as_completed(futures):
-            results.extend(future.result())
-    results.extend(root.glob(pattern))
 
-    return results
+def get_relative_file_paths(directory, recursive=False, ext_set=None):
+    """Get the relative paths of subfiles (recursively) in the directory that match the extension set.
+    """
+    if not recursive:
+        for entry in os.scandir(directory):
+            if entry.is_file():
+                file_name = entry.name
+                if ext_set is not None:
+                    ext = os.path.splitext(file_name)[1][1:].lower()
+                    if ext in ext_set:
+                        yield file_name
+                else:
+                    yield file_name
+    else:
+        for root, _, files in os.walk(directory):
+            for file in files:
+                relative_path = os.path.relpath(os.path.join(root, file), directory)
+                if ext_set is not None:
+                    ext = os.path.splitext(file)[1][1:].lower()
+                    if ext in ext_set:
+                        yield relative_path
+                else:
+                    yield relative_path
 
 
 def parse_args():
@@ -64,24 +72,12 @@ def main():
 
     # Use the path name instead of the file name as video_path/image_path (unique ID).
     if args.video_folder is not None:
-        video_path_list = []
-        video_folder = Path(args.video_folder)
-        for ext in tqdm(list(ALL_VIDEO_EXT)):
-            if args.recursive:
-                video_path_list += [str(file.relative_to(video_folder)) for file in parallel_rglob(video_folder, f"*.{ext}")]
-            else:
-                video_path_list += [str(file.relative_to(video_folder)) for file in video_folder.glob(f"*.{ext}")]
+        video_path_list = list(get_relative_file_paths(args.video_folder, recursive=args.recursive, ext_set=ALL_VIDEO_EXT))
         video_path_list = natsorted(video_path_list)
         meta_file_df = pd.DataFrame({args.video_path_column: video_path_list})
     
     if args.image_folder is not None:
-        image_path_list = []
-        image_folder = Path(args.image_folder)
-        for ext in tqdm(list(ALL_IMGAE_EXT)):
-            if args.recursive:
-                image_path_list += [str(file.relative_to(image_folder)) for file in parallel_rglob(video_folder, f"*.{ext}")]
-            else:
-                image_path_list += [str(file.relative_to(image_folder)) for file in image_folder.glob(f"*.{ext}")]
+        image_path_list = list(get_relative_file_paths(args.image_folder, recursive=args.recursive, ext_set=ALL_IMGAE_EXT))
         image_path_list = natsorted(image_path_list)
         meta_file_df = pd.DataFrame({args.image_path_column: image_path_list})
 

predict_i2v.py

@@ -14,9 +14,11 @@ from transformers import (BertModel, BertTokenizer, CLIPImageProcessor,
 
 from easyanimate.models import (name_to_autoencoder_magvit,
                                 name_to_transformer3d)
+from easyanimate.models.transformer3d import get_teacache_coefficients
 from easyanimate.pipeline.pipeline_easyanimate_inpaint import \
     EasyAnimateInpaintPipeline
-from easyanimate.utils.fp8_optimization import convert_weight_dtype_wrapper
+from easyanimate.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                                convert_weight_dtype_wrapper)
 from easyanimate.utils.lora_utils import merge_lora, unmerge_lora
 from easyanimate.utils.utils import get_image_to_video_latent, save_videos_grid
 
@@ -30,9 +32,13 @@ from easyanimate.utils.utils import get_image_to_video_latent, save_videos_grid
 # resulting in slower speeds but saving a large amount of GPU memory.
 # 
 # EasyAnimateV1, V2 and V3 support "model_cpu_offload" "sequential_cpu_offload"
-# 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" 
+# EasyAnimateV4, V5 and V5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload"
 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.08
 
 # Config and model path
 config_path         = "config/easyanimate_video_v5.1_magvit_qwen.yaml"
@@ -75,7 +81,7 @@ validation_image_end    = None
 # EasyAnimateV4, V5 and V5.1 support English and Chinese.
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
-prompt                  = "一只棕褐色的狗在摇晃脑袋，坐在一个舒适的房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里的灯光柔和温暖，营造出舒适的氛围。"
+prompt                  = "一只棕褐色的狗正摇晃着脑袋，坐在一个舒适的房间里的浅色沙发上。沙发看起来柔软而宽敞，为这只活泼的狗狗提供了一个完美的休息地点。在狗的后面，靠墙摆放着一个架子，架子上挂着一幅精美的镶框画，画中描绘着一些美丽的风景或场景。画框周围装饰着粉红色的花朵，这些花朵不仅增添了房间的色彩，还带来了一丝自然和生机。房间里的灯光柔和而温暖，从天花板上的吊灯和角落里的台灯散发出来，营造出一种温馨舒适的氛围。整个空间给人一种宁静和谐的感觉，仿佛时间在这里变得缓慢而美好。"
 negative_prompt         = "扭曲的身体，肢体残缺，文本字幕，漫画，静止，丑陋，错误，乱码。"
 # 
 # Using longer neg prompt such as "Blurring, mutation, deformation, distortion, dark and solid, comics, text subtitles, line art." can increase stability
@@ -96,7 +102,7 @@ Choosen_Transformer3DModel = name_to_transformer3d[
 ]
 
 transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
-if weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     transformer_additional_kwargs["upcast_attention"] = True
 
 transformer = Choosen_Transformer3DModel.from_pretrained_2d(
@@ -140,7 +146,7 @@ vae = Choosen_AutoencoderKL.from_pretrained(
     subfolder="vae",
     vae_additional_kwargs=OmegaConf.to_container(config['vae_kwargs'])
 ).to(weight_dtype)
-if config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     vae.upcast_vae = True
 
 if vae_path is not None:
@@ -241,13 +247,29 @@ pipeline = EasyAnimateInpaintPipeline(
 )
 
 if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline._manual_cpu_offload_in_sequential_cpu_offload = []
+    for name, _text_encoder in zip(["text_encoder", "text_encoder_2"], [pipeline.text_encoder, pipeline.text_encoder_2]):
+        if isinstance(_text_encoder, Qwen2VLForConditionalGeneration):
+            if hasattr(_text_encoder, "visual"):
+                del _text_encoder.visual
+            convert_model_weight_to_float8(_text_encoder)
+            convert_weight_dtype_wrapper(_text_encoder, weight_dtype)
+            pipeline._manual_cpu_offload_in_sequential_cpu_offload = [name]
     pipeline.enable_sequential_cpu_offload()
 elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
-    pipeline.enable_model_cpu_offload()
+    for _text_encoder in [pipeline.text_encoder, pipeline.text_encoder_2]:
+        if hasattr(_text_encoder, "visual"):
+            del _text_encoder.visual
     convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload()
 else:
     pipeline.enable_model_cpu_offload()
 
+coefficients = get_teacache_coefficients(model_name)
+if coefficients is not None and enable_teacache:
+    print(f"Enable TeaCache with threshold: {teacache_threshold}.")
+    pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients)
+
 generator = torch.Generator(device="cuda").manual_seed(seed)
 
 if lora_path is not None:

predict_t2v.py

@@ -7,18 +7,19 @@ from diffusers import (DDIMScheduler, DPMSolverMultistepScheduler,
                        FlowMatchEulerDiscreteScheduler, PNDMScheduler)
 from omegaconf import OmegaConf
 from PIL import Image
-from transformers import (BertModel, BertTokenizer,
-                          CLIPImageProcessor, CLIPVisionModelWithProjection,
-                          Qwen2Tokenizer, Qwen2VLForConditionalGeneration,
-                          T5EncoderModel, T5Tokenizer)
+from transformers import (BertModel, BertTokenizer, CLIPImageProcessor,
+                          CLIPVisionModelWithProjection, Qwen2Tokenizer,
+                          Qwen2VLForConditionalGeneration, T5EncoderModel,
+                          T5Tokenizer)
 
 from easyanimate.models import (name_to_autoencoder_magvit,
                                 name_to_transformer3d)
-from easyanimate.pipeline.pipeline_easyanimate import \
-    EasyAnimatePipeline
+from easyanimate.models.transformer3d import get_teacache_coefficients
+from easyanimate.pipeline.pipeline_easyanimate import EasyAnimatePipeline
 from easyanimate.pipeline.pipeline_easyanimate_inpaint import \
     EasyAnimateInpaintPipeline
-from easyanimate.utils.fp8_optimization import convert_weight_dtype_wrapper
+from easyanimate.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                                convert_weight_dtype_wrapper)
 from easyanimate.utils.lora_utils import merge_lora, unmerge_lora
 from easyanimate.utils.utils import get_image_to_video_latent, save_videos_grid
 
@@ -32,9 +33,13 @@ from easyanimate.utils.utils import get_image_to_video_latent, save_videos_grid
 # resulting in slower speeds but saving a large amount of GPU memory.
 # 
 # EasyAnimateV1, V2 and V3 support "model_cpu_offload" "sequential_cpu_offload"
-# 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" 
+# EasyAnimateV4, V5 and V5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload"
 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.08
 
 # Config and model path
 config_path         = "config/easyanimate_video_v5.1_magvit_qwen.yaml"
@@ -70,7 +75,7 @@ weight_dtype        = torch.bfloat16
 # EasyAnimateV4, V5 and V5.1 support English and Chinese.
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
-prompt              = "一只棕褐色的狗在摇晃脑袋，坐在一个舒适的房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里的灯光柔和温暖，营造出舒适的氛围。"
+prompt              = "一只棕褐色的狗正摇晃着脑袋，坐在一个舒适的房间里的浅色沙发上。沙发看起来柔软而宽敞，为这只活泼的狗狗提供了一个完美的休息地点。在狗的后面，靠墙摆放着一个架子，架子上挂着一幅精美的镶框画，画中描绘着一些美丽的风景或场景。画框周围装饰着粉红色的花朵，这些花朵不仅增添了房间的色彩，还带来了一丝自然和生机。房间里的灯光柔和而温暖，从天花板上的吊灯和角落里的台灯散发出来，营造出一种温馨舒适的氛围。整个空间给人一种宁静和谐的感觉，仿佛时间在这里变得缓慢而美好。"
 negative_prompt     = "扭曲的身体，肢体残缺，文本字幕，漫画，静止，丑陋，错误，乱码。"
 # 
 # Using longer neg prompt such as "Blurring, mutation, deformation, distortion, dark and solid, comics, text subtitles, line art." can increase stability
@@ -91,7 +96,7 @@ Choosen_Transformer3DModel = name_to_transformer3d[
 ]
 
 transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
-if weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     transformer_additional_kwargs["upcast_attention"] = True
 
 transformer = Choosen_Transformer3DModel.from_pretrained_2d(
@@ -135,7 +140,7 @@ vae = Choosen_AutoencoderKL.from_pretrained(
     subfolder="vae",
     vae_additional_kwargs=OmegaConf.to_container(config['vae_kwargs'])
 ).to(weight_dtype)
-if config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     vae.upcast_vae = True
 
 if vae_path is not None:
@@ -249,13 +254,29 @@ else:
     )
 
 if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline._manual_cpu_offload_in_sequential_cpu_offload = []
+    for name, _text_encoder in zip(["text_encoder", "text_encoder_2"], [pipeline.text_encoder, pipeline.text_encoder_2]):
+        if isinstance(_text_encoder, Qwen2VLForConditionalGeneration):
+            if hasattr(_text_encoder, "visual"):
+                del _text_encoder.visual
+            convert_model_weight_to_float8(_text_encoder)
+            convert_weight_dtype_wrapper(_text_encoder, weight_dtype)
+            pipeline._manual_cpu_offload_in_sequential_cpu_offload = [name]
     pipeline.enable_sequential_cpu_offload()
 elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    for _text_encoder in [pipeline.text_encoder, pipeline.text_encoder_2]:
+        if hasattr(_text_encoder, "visual"):
+            del _text_encoder.visual
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
     pipeline.enable_model_cpu_offload()
-    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
 else:
     pipeline.enable_model_cpu_offload()
 
+coefficients = get_teacache_coefficients(model_name)
+if coefficients is not None and enable_teacache:
+    print(f"Enable TeaCache with threshold: {teacache_threshold}.")
+    pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients)
+
 generator = torch.Generator(device="cuda").manual_seed(seed)
 
 if lora_path is not None:

predict_v2v.py

@@ -14,12 +14,13 @@ from transformers import (BertModel, BertTokenizer, CLIPImageProcessor,
 
 from easyanimate.models import (name_to_autoencoder_magvit,
                                 name_to_transformer3d)
+from easyanimate.models.transformer3d import get_teacache_coefficients
 from easyanimate.pipeline.pipeline_easyanimate_inpaint import \
     EasyAnimateInpaintPipeline
-from easyanimate.utils.fp8_optimization import convert_weight_dtype_wrapper
+from easyanimate.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                                convert_weight_dtype_wrapper)
 from easyanimate.utils.lora_utils import merge_lora, unmerge_lora
-from easyanimate.utils.utils import (get_video_to_video_latent,
-                                     save_videos_grid)
+from easyanimate.utils.utils import get_video_to_video_latent, save_videos_grid
 
 # GPU memory mode, which can be choosen in [model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
@@ -31,9 +32,13 @@ from easyanimate.utils.utils import (get_video_to_video_latent,
 # resulting in slower speeds but saving a large amount of GPU memory.
 # 
 # EasyAnimateV3 support "model_cpu_offload" "sequential_cpu_offload"
-# 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" 
+# EasyAnimateV4, V5 and V5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload"
 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.08
 
 # Config and model path
 config_path         = "config/easyanimate_video_v5.1_magvit_qwen.yaml"
@@ -63,13 +68,17 @@ fps                 = 8
 # 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
-# If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
+# If you are preparing to redraw the reference video, set validation_video and validation_video_mask. 
+# If you do not use validation_video_mask, the entire video will be redrawn; 
+# if you use validation_video_mask, as shown in asset/mask.jpg, only a portion of the video will be redrawn.
+# Please set a larger denoise_strength when using validation_video_mask, such as 1.00 instead of 0.70
 validation_video        = "asset/1.mp4"
+validation_video_mask   = None 
 denoise_strength        = 0.70
 
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
-prompt                  = "一只穿着小外套的猫咪正在花园秋千上安静地弹吉他。晚霞的余光洒在它柔软的毛皮上，和煦的微风轻轻拂过，周围斑驳的光影随着音乐的旋律轻轻摇曳。"
+prompt                  = "一只穿着小外套的猫咪正安静地坐在花园的秋千上弹吉他。它的小外套精致而合身，增添了几分俏皮与可爱。晚霞的余光洒在它柔软的毛皮上，给它的毛发镀上了一层温暖的金色光辉。和煦的微风轻轻拂过，带来阵阵花香和草木的气息，令人心旷神怡。周围斑驳的光影随着音乐的旋律轻轻摇曳，仿佛整个花园都在为这只小猫咪的演奏伴舞。阳光透过树叶间的缝隙，投下一片片光影交错的图案，与悠扬的吉他声交织在一起，营造出一种梦幻而宁静的氛围。猫咪专注而投入地弹奏着，每一个音符都似乎充满了魔力，让这个傍晚变得更加美好。"
 negative_prompt         = "扭曲的身体，肢体残缺，文本字幕，漫画，静止，丑陋，错误，乱码。"
 # 
 # Using longer neg prompt such as "Blurring, mutation, deformation, distortion, dark and solid, comics, text subtitles, line art." can increase stability
@@ -90,7 +99,7 @@ Choosen_Transformer3DModel = name_to_transformer3d[
 ]
 
 transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
-if weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     transformer_additional_kwargs["upcast_attention"] = True
 
 transformer = Choosen_Transformer3DModel.from_pretrained_2d(
@@ -134,7 +143,7 @@ vae = Choosen_AutoencoderKL.from_pretrained(
     subfolder="vae",
     vae_additional_kwargs=OmegaConf.to_container(config['vae_kwargs'])
 ).to(weight_dtype)
-if config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     vae.upcast_vae = True
 
 if vae_path is not None:
@@ -236,13 +245,29 @@ pipeline = EasyAnimateInpaintPipeline(
 )
 
 if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline._manual_cpu_offload_in_sequential_cpu_offload = []
+    for name, _text_encoder in zip(["text_encoder", "text_encoder_2"], [pipeline.text_encoder, pipeline.text_encoder_2]):
+        if isinstance(_text_encoder, Qwen2VLForConditionalGeneration):
+            if hasattr(_text_encoder, "visual"):
+                del _text_encoder.visual
+            convert_model_weight_to_float8(_text_encoder)
+            convert_weight_dtype_wrapper(_text_encoder, weight_dtype)
+            pipeline._manual_cpu_offload_in_sequential_cpu_offload = [name]
     pipeline.enable_sequential_cpu_offload()
 elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    for _text_encoder in [pipeline.text_encoder, pipeline.text_encoder_2]:
+        if hasattr(_text_encoder, "visual"):
+            del _text_encoder.visual
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
     pipeline.enable_model_cpu_offload()
-    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
 else:
     pipeline.enable_model_cpu_offload()
 
+coefficients = get_teacache_coefficients(model_name)
+if coefficients is not None and enable_teacache:
+    print(f"Enable TeaCache with threshold: {teacache_threshold}.")
+    pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients)
+
 generator = torch.Generator(device="cuda").manual_seed(seed)
 
 if lora_path is not None:
@@ -252,7 +277,7 @@ if vae.cache_mag_vae:
     video_length = int((video_length - 1) // vae.mini_batch_encoder * vae.mini_batch_encoder) + 1 if video_length != 1 else 1
 else:
     video_length = int(video_length // vae.mini_batch_encoder * vae.mini_batch_encoder) if video_length != 1 else 1
-input_video, input_video_mask, clip_image = get_video_to_video_latent(validation_video, video_length=video_length, fps=fps, sample_size=sample_size)
+input_video, input_video_mask, clip_image = get_video_to_video_latent(validation_video, video_length=video_length, fps=fps, validation_video_mask=validation_video_mask, sample_size=sample_size)
 
 with torch.no_grad():
     sample = pipeline(

predict_v2v_control.py

@@ -4,23 +4,26 @@ import numpy as np
 import torch
 from diffusers import (DDIMScheduler, DPMSolverMultistepScheduler,
                        EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
-                       PNDMScheduler)
+                       FlowMatchEulerDiscreteScheduler, PNDMScheduler)
 from omegaconf import OmegaConf
 from PIL import Image
-from transformers import (BertModel, BertTokenizer,
-                          CLIPImageProcessor, CLIPVisionModelWithProjection,
-                          Qwen2Tokenizer, Qwen2VLForConditionalGeneration,
-                          T5EncoderModel, T5Tokenizer)
+from transformers import (BertModel, BertTokenizer, CLIPImageProcessor,
+                          CLIPVisionModelWithProjection, Qwen2Tokenizer,
+                          Qwen2VLForConditionalGeneration, T5EncoderModel,
+                          T5Tokenizer)
 
 from easyanimate.data.dataset_image_video import process_pose_file
 from easyanimate.models import (name_to_autoencoder_magvit,
                                 name_to_transformer3d)
+from easyanimate.models.transformer3d import get_teacache_coefficients
 from easyanimate.pipeline.pipeline_easyanimate_control import \
     EasyAnimateControlPipeline
+from easyanimate.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                                convert_weight_dtype_wrapper)
 from easyanimate.utils.lora_utils import merge_lora, unmerge_lora
-from easyanimate.utils.utils import get_video_to_video_latent, save_videos_grid, get_image_latent
-from easyanimate.utils.fp8_optimization import convert_weight_dtype_wrapper
-from diffusers import FlowMatchEulerDiscreteScheduler
+from easyanimate.utils.utils import (get_image_latent,
+                                     get_video_to_video_latent,
+                                     save_videos_grid)
 
 # GPU memory mode, which can be choosen in [model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
@@ -31,9 +34,13 @@ from diffusers import FlowMatchEulerDiscreteScheduler
 # sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, 
 # resulting in slower speeds but saving a large amount of GPU memory.
 # 
-# EasyAnimateV5 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload"
-# EasyAnimateV5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" 
+# EasyAnimateV5 and V5.1 support "model_cpu_offload" "model_cpu_offload_and_qfloat8" "sequential_cpu_offload"
 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.08
 
 # Config and model path
 config_path         = "config/easyanimate_video_v5.1_magvit_qwen.yaml"
@@ -67,7 +74,7 @@ ref_image               = None
 
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
-prompt                  = "一位穿着合身的白色连衣裙，带着细肩带的女人站在一个铺着木地板的房间里。她有一头深色的长发。背景是一个放着各种瓶子的架子。灯光温暖，背景似乎在室内。"
+prompt                  = "在这个阳光明媚的户外花园里，美女身穿一袭及膝的白色无袖连衣裙，裙摆在她轻盈的舞姿中轻柔地摆动，宛如一只翩翩起舞的蝴蝶。阳光透过树叶间洒下斑驳的光影，映衬出她柔和的脸庞和清澈的眼眸，显得格外优雅。仿佛每一个动作都在诉说着青春与活力，她在草地上旋转，裙摆随之飞扬，仿佛整个花园都因她的舞动而欢愉。周围五彩缤纷的花朵在微风中摇曳，玫瑰、菊花、百合，各自释放出阵阵香气，营造出一种轻松而愉快的氛围。"
 negative_prompt         = "扭曲的身体，肢体残缺，文本字幕，漫画，静止，丑陋，错误，乱码。"
 # 
 # Using longer neg prompt such as "Blurring, mutation, deformation, distortion, dark and solid, comics, text subtitles, line art." can increase stability
@@ -88,7 +95,7 @@ Choosen_Transformer3DModel = name_to_transformer3d[
 ]
 
 transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
-if weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     transformer_additional_kwargs["upcast_attention"] = True
 
 transformer = Choosen_Transformer3DModel.from_pretrained_2d(
@@ -132,7 +139,7 @@ vae = Choosen_AutoencoderKL.from_pretrained(
     subfolder="vae",
     vae_additional_kwargs=OmegaConf.to_container(config['vae_kwargs'])
 ).to(weight_dtype)
-if config['vae_kwargs'].get('vae_type', 'AutoencoderKL') == 'AutoencoderKLMagvit' and weight_dtype == torch.float16:
+if weight_dtype == torch.float16 and "v5.1" not in model_name.lower():
     vae.upcast_vae = True
 
 if vae_path is not None:
@@ -221,13 +228,29 @@ pipeline = EasyAnimateControlPipeline(
 )
 
 if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline._manual_cpu_offload_in_sequential_cpu_offload = []
+    for name, _text_encoder in zip(["text_encoder", "text_encoder_2"], [pipeline.text_encoder, pipeline.text_encoder_2]):
+        if isinstance(_text_encoder, Qwen2VLForConditionalGeneration):
+            if hasattr(_text_encoder, "visual"):
+                del _text_encoder.visual
+            convert_model_weight_to_float8(_text_encoder)
+            convert_weight_dtype_wrapper(_text_encoder, weight_dtype)
+            pipeline._manual_cpu_offload_in_sequential_cpu_offload = [name]
     pipeline.enable_sequential_cpu_offload()
 elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    for _text_encoder in [pipeline.text_encoder, pipeline.text_encoder_2]:
+        if hasattr(_text_encoder, "visual"):
+            del _text_encoder.visual
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
     pipeline.enable_model_cpu_offload()
-    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
 else:
     pipeline.enable_model_cpu_offload()
 
+coefficients = get_teacache_coefficients(model_name)
+if coefficients is not None and enable_teacache:
+    print(f"Enable TeaCache with threshold: {teacache_threshold}.")
+    pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients)
+
 generator = torch.Generator(device="cuda").manual_seed(seed)
 
 if lora_path is not None:

scripts/README_TRAIN_REWARD.md

@@ -228,6 +228,9 @@ if GPU_memory_mode == "sequential_cpu_offload":
     pipeline.enable_sequential_cpu_offload()
 elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
     pipeline.enable_model_cpu_offload()
+    for _text_encoder in [pipeline.text_encoder, pipeline.text_encoder_2]:
+        if hasattr(_text_encoder, "visual"):
+            del _text_encoder.visual
     convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
 else:
     pipeline.enable_model_cpu_offload()

scripts/train.py

@@ -186,8 +186,6 @@ def encode_prompt(
             texts.append(text)
         text_inputs = tokenizer(
             text=texts,
-            images=None,
-            videos=None,
             padding="max_length",
             max_length=max_length,
             truncation=True,

scripts/train_control.py

@@ -182,8 +182,6 @@ def encode_prompt(
             texts.append(text)
         text_inputs = tokenizer(
             text=texts,
-            images=None,
-            videos=None,
             padding="max_length",
             max_length=max_length,
             truncation=True,

scripts/train_lora.py

@@ -186,8 +186,6 @@ def encode_prompt(
             texts.append(text)
         text_inputs = tokenizer(
             text=texts,
-            images=None,
-            videos=None,
             padding="max_length",
             max_length=max_length,
             truncation=True,

scripts/train_reward_lora.py

@@ -1414,193 +1414,192 @@ def main():
             ]
 
             with accelerator.accumulate(transformer3d):
-                with accelerator.autocast():
-                    latents = torch.randn(*latent_shape, device=accelerator.device, dtype=weight_dtype)
+                latents = torch.randn(*latent_shape, device=accelerator.device, dtype=weight_dtype)
 
-                    if hasattr(noise_scheduler, "init_noise_sigma"):
-                        latents = latents * noise_scheduler.init_noise_sigma
+                if hasattr(noise_scheduler, "init_noise_sigma"):
+                    latents = latents * noise_scheduler.init_noise_sigma
 
-                    # Prepare inpaint latents if it needs.
-                    # Use zero latents if we want to t2v.
-                    mask_latents = torch.zeros_like(latents)[:, :1].to(latents.device, latents.dtype)
-                    masked_video_latents = torch.zeros_like(latents).to(latents.device, latents.dtype)
+                # Prepare inpaint latents if it needs.
+                # Use zero latents if we want to t2v.
+                mask_latents = torch.zeros_like(latents)[:, :1].to(latents.device, latents.dtype)
+                masked_video_latents = torch.zeros_like(latents).to(latents.device, latents.dtype)
 
-                    mask_input = torch.cat([mask_latents] * 2) if do_classifier_free_guidance else mask_latents
-                    masked_video_latents_input = (
-                        torch.cat([masked_video_latents] * 2) if do_classifier_free_guidance else masked_video_latents
-                    )
-                    inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(latents.dtype)
+                mask_input = torch.cat([mask_latents] * 2) if do_classifier_free_guidance else mask_latents
+                masked_video_latents_input = (
+                    torch.cat([masked_video_latents] * 2) if do_classifier_free_guidance else masked_video_latents
+                )
+                inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(latents.dtype)
 
-                    # Check that sizes of mask, masked image and latents match
-                    if num_channels_transformer != num_channels_latents:
-                        num_channels_mask = mask_latents.shape[1]
-                        num_channels_masked_image = masked_video_latents.shape[1]
-                        if num_channels_latents + num_channels_mask + num_channels_masked_image != transformer3d.config.in_channels:
-                            raise ValueError(
-                                f"Incorrect configuration settings! The config of `pipeline.transformer`: {transformer3d.config} expects"
-                                f" {transformer3d.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
-                                f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}"
-                                f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of"
-                                " `pipeline.transformer` or your `mask_image` or `image` input."
-                            )
+                # Check that sizes of mask, masked image and latents match
+                if num_channels_transformer != num_channels_latents:
+                    num_channels_mask = mask_latents.shape[1]
+                    num_channels_masked_image = masked_video_latents.shape[1]
+                    if num_channels_latents + num_channels_mask + num_channels_masked_image != transformer3d.config.in_channels:
+                        raise ValueError(
+                            f"Incorrect configuration settings! The config of `pipeline.transformer`: {transformer3d.config} expects"
+                            f" {transformer3d.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
+                            f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}"
+                            f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of"
+                            " `pipeline.transformer` or your `mask_image` or `image` input."
+                        )
 
-                    generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
-                    # Prepare extra step kwargs.
-                    extra_step_kwargs = prepare_extra_step_kwargs(noise_scheduler, generator, args.eta)
+                generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
+                # Prepare extra step kwargs.
+                extra_step_kwargs = prepare_extra_step_kwargs(noise_scheduler, generator, args.eta)
 
-                    # Create image_rotary_emb, style embedding & time ids
-                    grid_height = height // 8 // transformer3d.config.patch_size
-                    grid_width = width // 8 // transformer3d.config.patch_size
-                    base_size_width = 720 // 8 // transformer3d.config.patch_size
-                    base_size_height = 480 // 8 // transformer3d.config.patch_size
-                    grid_crops_coords = get_resize_crop_region_for_grid(
-                        (grid_height, grid_width), base_size_width, base_size_height
-                    )
-                    image_rotary_emb = get_3d_rotary_pos_embed(
-                        transformer3d.config.attention_head_dim, grid_crops_coords, grid_size=(grid_height, grid_width),
-                        temporal_size=latents.size(2), use_real=True,
-                    )
+                # Create image_rotary_emb, style embedding & time ids
+                grid_height = height // 8 // transformer3d.config.patch_size
+                grid_width = width // 8 // transformer3d.config.patch_size
+                base_size_width = 720 // 8 // transformer3d.config.patch_size
+                base_size_height = 480 // 8 // transformer3d.config.patch_size
+                grid_crops_coords = get_resize_crop_region_for_grid(
+                    (grid_height, grid_width), base_size_width, base_size_height
+                )
+                image_rotary_emb = get_3d_rotary_pos_embed(
+                    transformer3d.config.attention_head_dim, grid_crops_coords, grid_size=(grid_height, grid_width),
+                    temporal_size=latents.size(2), use_real=True,
+                )
 
-                    # Get other hunyuan params
-                    style = torch.tensor([0], device=accelerator.device)
+                # Get other hunyuan params
+                style = torch.tensor([0], device=accelerator.device)
 
-                    original_size = (1024, 1024)
-                    crops_coords_top_left = (0, 0)
-                    target_size = (height, width)
-                    add_time_ids = list(original_size + target_size + crops_coords_top_left)
-                    add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype)
+                original_size = (1024, 1024)
+                crops_coords_top_left = (0, 0)
+                target_size = (height, width)
+                add_time_ids = list(original_size + target_size + crops_coords_top_left)
+                add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype)
 
-                    if do_classifier_free_guidance:
-                        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-                        prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask])
-                        if prompt_embeds_2 is not None:
-                            prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2])
-                            prompt_attention_mask_2 = torch.cat([negative_prompt_attention_mask_2, prompt_attention_mask_2])
-                        add_time_ids = torch.cat([add_time_ids] * 2, dim=0)
-                        style = torch.cat([style] * 2, dim=0)
-                    
-                    prompt_embeds = prompt_embeds.to(device=accelerator.device)
-                    prompt_attention_mask = prompt_attention_mask.to(device=accelerator.device)
+                if do_classifier_free_guidance:
+                    prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+                    prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask])
                     if prompt_embeds_2 is not None:
-                        prompt_embeds_2 = prompt_embeds_2.to(device=accelerator.device)
-                        prompt_attention_mask_2 = prompt_attention_mask_2.to(device=accelerator.device)
-                    add_time_ids = add_time_ids.to(dtype=prompt_embeds.dtype, device=accelerator.device).repeat(
-                        args.train_batch_size, 1
+                        prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2])
+                        prompt_attention_mask_2 = torch.cat([negative_prompt_attention_mask_2, prompt_attention_mask_2])
+                    add_time_ids = torch.cat([add_time_ids] * 2, dim=0)
+                    style = torch.cat([style] * 2, dim=0)
+                
+                prompt_embeds = prompt_embeds.to(device=accelerator.device)
+                prompt_attention_mask = prompt_attention_mask.to(device=accelerator.device)
+                if prompt_embeds_2 is not None:
+                    prompt_embeds_2 = prompt_embeds_2.to(device=accelerator.device)
+                    prompt_attention_mask_2 = prompt_attention_mask_2.to(device=accelerator.device)
+                add_time_ids = add_time_ids.to(dtype=prompt_embeds.dtype, device=accelerator.device).repeat(
+                    args.train_batch_size, 1
+                )
+                style = style.to(device=accelerator.device).repeat(args.train_batch_size)
+
+                # Denoising loop
+                if args.backprop:
+                    if args.backprop_step_list is None:
+                        if args.backprop_strategy == "last":
+                            backprop_step_list = [args.num_inference_steps - 1]
+                        elif args.backprop_strategy == "tail":
+                            backprop_step_list = list(range(args.num_inference_steps))[-args.backprop_num_steps:]
+                        elif args.backprop_strategy == "uniform":
+                            interval = args.num_inference_steps // args.backprop_num_steps
+                            random_start = random.randint(0, interval)
+                            backprop_step_list = [random_start + i * interval for i in range(args.backprop_num_steps)]
+                        elif args.backprop_strategy == "random":
+                            backprop_step_list = random.sample(
+                                range(args.backprop_random_start_step, args.backprop_random_end_step + 1), args.backprop_num_steps
+                            )
+                        else:
+                            raise ValueError(f"Invalid backprop strategy: {args.backprop_strategy}.")
+                    else:
+                        backprop_step_list = args.backprop_step_list
+                
+                for i, t in enumerate(tqdm(timesteps)):
+                    # expand the latents if we are doing classifier free guidance
+                    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                    if hasattr(noise_scheduler, "scale_model_input"):
+                        latent_model_input = noise_scheduler.scale_model_input(latent_model_input, t)
+                    
+                    # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input
+                    t_expand = torch.tensor([t] * latent_model_input.shape[0], device=accelerator.device).to(
+                        dtype=latent_model_input.dtype
                     )
-                    style = style.to(device=accelerator.device).repeat(args.train_batch_size)
 
-                    # Denoising loop
-                    if args.backprop:
-                        if args.backprop_step_list is None:
-                            if args.backprop_strategy == "last":
-                                backprop_step_list = [args.num_inference_steps - 1]
-                            elif args.backprop_strategy == "tail":
-                                backprop_step_list = list(range(args.num_inference_steps))[-args.backprop_num_steps:]
-                            elif args.backprop_strategy == "uniform":
-                                interval = args.num_inference_steps // args.backprop_num_steps
-                                random_start = random.randint(0, interval)
-                                backprop_step_list = [random_start + i * interval for i in range(args.backprop_num_steps)]
-                            elif args.backprop_strategy == "random":
-                                backprop_step_list = random.sample(
-                                    range(args.backprop_random_start_step, args.backprop_random_end_step + 1), args.backprop_num_steps
-                                )
-                            else:
-                                raise ValueError(f"Invalid backprop strategy: {args.backprop_strategy}.")
-                        else:
-                            backprop_step_list = args.backprop_step_list
-                    
-                    for i, t in enumerate(tqdm(timesteps)):
-                        # expand the latents if we are doing classifier free guidance
-                        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                        if hasattr(noise_scheduler, "scale_model_input"):
-                            latent_model_input = noise_scheduler.scale_model_input(latent_model_input, t)
-                        
-                        # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input
-                        t_expand = torch.tensor([t] * latent_model_input.shape[0], device=accelerator.device).to(
-                            dtype=latent_model_input.dtype
-                        )
+                    # predict the noise residual
+                    if args.stop_latent_model_input_gradient:
+                        # See https://arxiv.org/abs/2405.00760
+                        latent_model_input = latent_model_input.detach()
+                    noise_pred = transformer3d(
+                        latent_model_input,
+                        t_expand,
+                        encoder_hidden_states=prompt_embeds,
+                        text_embedding_mask=prompt_attention_mask,
+                        encoder_hidden_states_t5=prompt_embeds_2,
+                        text_embedding_mask_t5=prompt_attention_mask_2,
+                        image_meta_size=add_time_ids,
+                        style=style,
+                        image_rotary_emb=image_rotary_emb,
+                        inpaint_latents=inpaint_latents,
+                        clip_encoder_hidden_states=None,
+                        clip_attention_mask=None,
+                        return_dict=False,
+                    )[0]
 
-                        # predict the noise residual
-                        if args.stop_latent_model_input_gradient:
-                            # See https://arxiv.org/abs/2405.00760
-                            latent_model_input = latent_model_input.detach()
-                        noise_pred = transformer3d(
-                            latent_model_input,
-                            t_expand,
-                            encoder_hidden_states=prompt_embeds,
-                            text_embedding_mask=prompt_attention_mask,
-                            encoder_hidden_states_t5=prompt_embeds_2,
-                            text_embedding_mask_t5=prompt_attention_mask_2,
-                            image_meta_size=add_time_ids,
-                            style=style,
-                            image_rotary_emb=image_rotary_emb,
-                            inpaint_latents=inpaint_latents,
-                            clip_encoder_hidden_states=None,
-                            clip_attention_mask=None,
-                            return_dict=False,
-                        )[0]
+                    # Optimize the denoising results only for the specified steps.
+                    if i in backprop_step_list:
+                        noise_pred = noise_pred
+                    else:
+                        noise_pred = noise_pred.detach()
 
-                        # Optimize the denoising results only for the specified steps.
-                        if i in backprop_step_list:
-                            noise_pred = noise_pred
-                        else:
-                            noise_pred = noise_pred.detach()
+                    # perform guidance
+                    if do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + args.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
-                        # perform guidance
-                        if do_classifier_free_guidance:
-                            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                            noise_pred = noise_pred_uncond + args.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents = noise_scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
 
-                        # compute the previous noisy sample x_t -> x_t-1
-                        latents = noise_scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                # decode latents (tensor)
+                # latents = latents.permute(0, 2, 1, 3, 4)  # [B, C, T, H, W]
+                # Since the casual VAE decoding consumes a large amount of VRAM, and we need to keep the decoding 
+                # operation within the computational graph. Thus, we only decode the first args.num_decoded_latents 
+                # to calculate the reward.
+                # TODO: Decode all latents but keep a portion of the decoding operation within the computational graph.
+                sampled_latent_indices = list(range(args.num_decoded_latents))
+                sampled_latents = latents[:, :, sampled_latent_indices, :, :]
+                sampled_latents = 1 / vae.config.scaling_factor * sampled_latents
+                sampled_frames = vae.decode(sampled_latents)[0]
+                sampled_frames = sampled_frames.clamp(-1, 1)
+                sampled_frames = (sampled_frames / 2 + 0.5).clamp(0, 1)  # [-1, 1] -> [0, 1]
 
-                    # decode latents (tensor)
-                    # latents = latents.permute(0, 2, 1, 3, 4)  # [B, C, T, H, W]
-                    # Since the casual VAE decoding consumes a large amount of VRAM, and we need to keep the decoding 
-                    # operation within the computational graph. Thus, we only decode the first args.num_decoded_latents 
-                    # to calculate the reward.
-                    # TODO: Decode all latents but keep a portion of the decoding operation within the computational graph.
-                    sampled_latent_indices = list(range(args.num_decoded_latents))
-                    sampled_latents = latents[:, :, sampled_latent_indices, :, :]
-                    sampled_latents = 1 / vae.config.scaling_factor * sampled_latents
-                    sampled_frames = vae.decode(sampled_latents)[0]
-                    sampled_frames = sampled_frames.clamp(-1, 1)
-                    sampled_frames = (sampled_frames / 2 + 0.5).clamp(0, 1)  # [-1, 1] -> [0, 1]
+                if global_step % args.checkpointing_steps == 0:
+                    saved_file = f"sample-{global_step}-{accelerator.process_index}.mp4"
+                    save_videos_grid(
+                        sampled_frames.to(torch.float32).detach().cpu(),
+                        os.path.join(args.output_dir, "train_sample", saved_file),
+                        fps=8
+                    )
+                
+                if args.num_sampled_frames is not None:
+                    num_frames = sampled_frames.size(2) - 1
+                    sampled_frames_indices = torch.linspace(0, num_frames, steps=args.num_sampled_frames).long()
+                    sampled_frames = sampled_frames[:, :, sampled_frames_indices, :, :]
+                # compute loss and reward
+                loss, reward = loss_fn(sampled_frames, train_prompt)
 
-                    if global_step % args.checkpointing_steps == 0:
-                        saved_file = f"sample-{global_step}-{accelerator.process_index}.mp4"
-                        save_videos_grid(
-                            sampled_frames.to(torch.float32).detach().cpu(),
-                            os.path.join(args.output_dir, "train_sample", saved_file),
-                            fps=8
-                        )
-                    
-                    if args.num_sampled_frames is not None:
-                        num_frames = sampled_frames.size(2) - 1
-                        sampled_frames_indices = torch.linspace(0, num_frames, steps=args.num_sampled_frames).long()
-                        sampled_frames = sampled_frames[:, :, sampled_frames_indices, :, :]
-                    # compute loss and reward
-                    loss, reward = loss_fn(sampled_frames, train_prompt)
+                # Gather the losses and rewards across all processes for logging (if we use distributed training).
+                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                avg_reward = accelerator.gather(reward.repeat(args.train_batch_size)).mean()
+                train_loss += avg_loss.item() / args.gradient_accumulation_steps
+                train_reward += avg_reward.item() / args.gradient_accumulation_steps
 
-                    # Gather the losses and rewards across all processes for logging (if we use distributed training).
-                    avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
-                    avg_reward = accelerator.gather(reward.repeat(args.train_batch_size)).mean()
-                    train_loss += avg_loss.item() / args.gradient_accumulation_steps
-                    train_reward += avg_reward.item() / args.gradient_accumulation_steps
-
-                    # Backpropagate
-                    accelerator.backward(loss)
-                    if accelerator.sync_gradients:
-                        total_norm = accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
-                        # If use_deepspeed, `total_norm` cannot be logged by accelerator.
-                        if not args.use_deepspeed:
-                            accelerator.log({"total_norm": total_norm}, step=global_step)
-                        else:
-                            if hasattr(optimizer, "optimizer") and hasattr(optimizer.optimizer, "_global_grad_norm"):
-                                accelerator.log({"total_norm":  optimizer.optimizer._global_grad_norm}, step=global_step)
-                    optimizer.step()
-                    lr_scheduler.step()
-                    optimizer.zero_grad()
+                # Backpropagate
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    total_norm = accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm)
+                    # If use_deepspeed, `total_norm` cannot be logged by accelerator.
+                    if not args.use_deepspeed:
+                        accelerator.log({"total_norm": total_norm}, step=global_step)
+                    else:
+                        if hasattr(optimizer, "optimizer") and hasattr(optimizer.optimizer, "_global_grad_norm"):
+                            accelerator.log({"total_norm":  optimizer.optimizer._global_grad_norm}, step=global_step)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
             
             # Checks if the accelerator has performed an optimization step behind the scenes
             if accelerator.sync_gradients: