From 10478729cdb73cf6e3f260be12e11a0849f76f4b Mon Sep 17 00:00:00 2001 From: hkunzhe Date: Mon, 10 Feb 2025 19:52:29 +0800 Subject: [PATCH] update TeaCache --- easyanimate/models/transformer3d.py | 20 ++++++++++++++++---- easyanimate/ui/ui.py | 13 +++++++++---- predict_i2v.py | 6 ++++-- predict_t2v.py | 6 ++++-- predict_v2v.py | 6 ++++-- predict_v2v_control.py | 6 ++++-- 6 files changed, 41 insertions(+), 16 deletions(-) diff --git a/easyanimate/models/transformer3d.py b/easyanimate/models/transformer3d.py index c4e16ec..c1a86a6 100755 --- a/easyanimate/models/transformer3d.py +++ b/easyanimate/models/transformer3d.py @@ -121,6 +121,22 @@ class TeaCache(): 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. @@ -1472,10 +1488,6 @@ class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin): rel_l1_thresh: float, coefficients: list[float] = [-10.47857366, 8.33844143, -0.78477557, 0.68798618, 0.0136149] ): - # The coefficient was obtained by sampling videos from T2V CompBench using EasyAnimateV5.1-12b-zh-InP. - # This coefficient can be applied to both the EasyAnimateV5.1-12b-zh and EasyAnimateV5.1-12b-Control. - # The coefficients for EasyAnimateV5.1-7b-zh-InP should be: - # [-3.64204720e+03, 1.43764725e+03, -1.93045263e+02, 1.09596499e+01, -1.70663507e-01] self.teacache = TeaCache(coefficients, num_steps, rel_l1_thresh=rel_l1_thresh) def _set_gradient_checkpointing(self, module, value=False): diff --git a/easyanimate/ui/ui.py b/easyanimate/ui/ui.py index aa84a06..ab19634 100755 --- a/easyanimate/ui/ui.py +++ b/easyanimate/ui/ui.py @@ -30,6 +30,7 @@ from transformers import (BertModel, BertTokenizer, CLIPImageProcessor, from ..data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio 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 \ @@ -467,8 +468,10 @@ class EasyAnimateController: # lora part self.pipeline = merge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider) - if self.edition == "v5.1" and self.enable_teacache: - self.pipeline.transformer.enable_teacache(sample_step_slider, self.teacache_threshold) + 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": @@ -1020,6 +1023,7 @@ 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 @@ -1272,9 +1276,10 @@ class EasyAnimateController_Modelscope: # lora part self.pipeline = merge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider) - if self.edition == "v5.1" and self.enable_teacache: + 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) + self.pipeline.transformer.enable_teacache(sample_step_slider, self.teacache_threshold, coefficients=coefficients) try: if self.model_type == "Inpaint": diff --git a/predict_i2v.py b/predict_i2v.py index 3d214db..10bce44 100755 --- a/predict_i2v.py +++ b/predict_i2v.py @@ -14,6 +14,7 @@ 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 @@ -256,9 +257,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() -if "v5.1" in config_path and enable_teacache: +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) + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients) generator = torch.Generator(device="cuda").manual_seed(seed) diff --git a/predict_t2v.py b/predict_t2v.py index 6f202e7..dac7139 100755 --- a/predict_t2v.py +++ b/predict_t2v.py @@ -14,6 +14,7 @@ from transformers import (BertModel, BertTokenizer, from easyanimate.models import (name_to_autoencoder_magvit, name_to_transformer3d) +from easyanimate.models.transformer3d import get_teacache_coefficients from easyanimate.pipeline.pipeline_easyanimate import \ EasyAnimatePipeline from easyanimate.pipeline.pipeline_easyanimate_inpaint import \ @@ -264,9 +265,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() -if "v5.1" in config_path and enable_teacache: +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) + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients) generator = torch.Generator(device="cuda").manual_seed(seed) diff --git a/predict_v2v.py b/predict_v2v.py index 8c14659..a92fe66 100755 --- a/predict_v2v.py +++ b/predict_v2v.py @@ -14,6 +14,7 @@ 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 @@ -251,9 +252,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() -if "v5.1" in config_path and enable_teacache: +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) + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients) generator = torch.Generator(device="cuda").manual_seed(seed) diff --git a/predict_v2v_control.py b/predict_v2v_control.py index 19c8160..e2da163 100755 --- a/predict_v2v_control.py +++ b/predict_v2v_control.py @@ -15,6 +15,7 @@ from transformers import (BertModel, BertTokenizer, 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.lora_utils import merge_lora, unmerge_lora @@ -236,9 +237,10 @@ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": else: pipeline.enable_model_cpu_offload() -if "v5.1" in config_path and enable_teacache: +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) + pipeline.transformer.enable_teacache(num_inference_steps, teacache_threshold, coefficients=coefficients) generator = torch.Generator(device="cuda").manual_seed(seed)