Add TeaCache and Enhance-A-Video

README.md

@@ -8,6 +8,8 @@ Ruyi is an image-to-video model capable of generating cinematic-quality videos a
 
 ## News
 
+- Jan 14, 2025: [TeaCache](https://github.com/ali-vilab/TeaCache) and [Enhance-A-Video](https://github.com/NUS-HPC-AI-Lab/Enhance-A-Video) have been added as nodes for ComfyUI, offering faster generation speeds and improved video quality. Simply link these nodes before the sampler node to use them.
+
 - Jan 6, 2025: FP8 added for ComfyUI node. GPU memory decreases depending on the modes: bf16 default > fp8 lite > fp8 strong  > fp8 extreme.
 
 - Dec 24, 2024: The diffusion model is updated to fix the black lines when creating 3:4 or 4:5 videos.

README_CN.md

@@ -6,6 +6,8 @@ Ruyi 是一款图生视频模型，能够生成 **768 分辨率、每秒 24 帧
 
 ## 更新说明
 
+2025年1月14日：[TeaCache](https://github.com/ali-vilab/TeaCache) 和 [Enhance-A-Video](https://github.com/NUS-HPC-AI-Lab/Enhance-A-Video) 已作为 ComfyUI 的节点添加，提供更快的生成速度和更高的视频质量。只需在 Sampler 节点之前链接这些节点即可使用。
+
 2025年1月6日: 为ComfyUI加入了FP8的支持，有三种fp8的模式，显存消耗：bf16 default > fp8 lite > fp8 strong > fp8 extreme。
 
 2024年12月24日: 更新了transformer模型以解决3:4和4:5比例的视频会出现黑杠的问题。

comfyui/comfyui_nodes.py

@@ -214,6 +214,7 @@ class Ruyi_LoadModel:
             'model_type': model_type,
             'loras': [],
             'strength_model': [],
+            'plugins': {}
         }
         return (ruyi_model,)
 
@@ -243,12 +244,103 @@ class Ruyi_LoadLora:
                     'model_type': ruyi_model["model_type"],
                     'loras': ruyi_model.get("loras", []) + [folder_paths.get_full_path("loras", lora_name)],
                     'strength_model': ruyi_model.get("strength_model", []) + [strength_model],
+                    'plugins': ruyi_model["plugins"],
                 }, 
             )
         else:
             return (ruyi_model,)
 
 
+class Ruyi_TeaCache:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "ruyi_model": ("RUYI_MODEL",),
+                "enable": ("BOOLEAN", {"default": True, "tooltip": "Enable TeaCache."}),
+                "threshold": (
+                    "FLOAT", 
+                    {
+                        "default": 0.10, "min": 0.0, "max": 0.50, "step": 0.01, 
+                        "tooltip": "TeaCache threshold. 0.10 caches 6 ~ 8 steps, 0.15 caches 10 ~ 12 steps normally."
+                    }
+                ),
+                "skip_start_steps": ("INT", {"default": 3, "min": 1, "step": 1, "tooltip": "First n steps do not use TeaCache."}),
+                "skip_end_steps": ("INT", {"default": 1, "min": 1, "step": 1, "tooltip": "Last n steps do not use TeaCache."}),
+                "offload_cpu": ("BOOLEAN", {"default": True, "Tooltip": "Offload TeaCache tensors to cpu to save GPU memory."}),
+            }
+        }
+    RETURN_TYPES = ("RUYI_MODEL",)
+    RETURN_NAMES = ("ruyi_model",)
+    FUNCTION = "tea_cache"
+    CATEGORY = "Ruyi"
+
+    def tea_cache(self, ruyi_model, enable, threshold, skip_start_steps, skip_end_steps, offload_cpu):
+        tea_cache_parameters = {
+            "enable": enable,
+            "threshold": threshold,
+            "skip_start_steps": skip_start_steps,
+            "skip_end_steps": skip_end_steps,
+            "offload_cpu": offload_cpu
+        }
+
+        return (
+            {
+                'pipeline': ruyi_model["pipeline"], 
+                'dtype': ruyi_model["dtype"],
+                'model_path': ruyi_model["model_path"],
+                'model_type': ruyi_model["model_type"],
+                'loras': ruyi_model["loras"],
+                'strength_model': ruyi_model["strength_model"],
+                'plugins': {"tea_cache": tea_cache_parameters} | ruyi_model["plugins"],
+            }, 
+        )
+
+
+class Ruyi_EnhanceAVideo:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "ruyi_model": ("RUYI_MODEL",),
+                "enable": ("BOOLEAN", {"default": True, "tooltip": "Enable Enhance-A-Video."}),
+                "weight": (
+                    "FLOAT", 
+                    {
+                        "default": 1.0, "min": 0, "max": 100, 
+                        "tooltip": "The weight of Enhance-A-Video. Should be smaller than 5 or 10 to get better results."
+                    }
+                ),
+                "skip_start_steps": ("INT", {"default": 0, "min": 0, "step": 1, "tooltip": "First n steps do not use Enhance-A-Video."}),
+                "skip_end_steps": ("INT", {"default": 0, "min": 0, "step": 1, "tooltip": "First n steps do not use Enhance-A-Video."})
+            }
+        }
+    RETURN_TYPES = ("RUYI_MODEL",)
+    RETURN_NAMES = ("ruyi_model",)
+    FUNCTION = "enhance_a_video"
+    CATEGORY = "Ruyi"
+
+    def enhance_a_video(self, ruyi_model, enable, weight, skip_start_steps, skip_end_steps):
+        enhance_a_video_parameters = {
+            "enable": enable,
+            "weight": weight,
+            "skip_start_steps": skip_start_steps,
+            "skip_end_steps": skip_end_steps,
+        }
+
+        return (
+            {
+                'pipeline': ruyi_model["pipeline"], 
+                'dtype': ruyi_model["dtype"],
+                'model_path': ruyi_model["model_path"],
+                'model_type': ruyi_model["model_type"],
+                'loras': ruyi_model["loras"],
+                'strength_model': ruyi_model["strength_model"],
+                'plugins': {"enhance_a_video": enhance_a_video_parameters} | ruyi_model["plugins"],
+            }, 
+        )
+
+
 class Ruyi_I2VSampler:
     @classmethod
     def INPUT_TYPES(s):
@@ -379,6 +471,34 @@ class Ruyi_I2VSampler:
         
         return noise_scheduler
 
+    def initialize_plugins(self, pipeline, parameters, num_steps):
+        # TeaCache
+        if "tea_cache" in parameters:
+            tea_cache_parameters = parameters["tea_cache"]
+            pipeline.transformer.tea_cache.initialize(
+                tea_cache_parameters["enable"],
+                tea_cache_parameters["threshold"],
+                tea_cache_parameters["skip_start_steps"],
+                tea_cache_parameters["skip_end_steps"],
+                num_steps,
+                tea_cache_parameters["offload_cpu"],
+            )
+        else:
+            pipeline.transformer.tea_cache.disable()
+        
+        # Enhance-A-Video
+        if "enhance_a_video" in parameters:
+            enhance_a_video_parameters = parameters["enhance_a_video"]
+            pipeline.transformer.enhance_a_video.initialize(
+                enhance_a_video_parameters["enable"],
+                enhance_a_video_parameters["weight"],
+                enhance_a_video_parameters["skip_start_steps"],
+                enhance_a_video_parameters["skip_end_steps"],
+                num_steps,
+            )
+        else:
+            pipeline.transformer.enhance_a_video.disable()
+
     def process(
             self, ruyi_model, video_length, base_resolution, seed, steps, cfg, scheduler, 
             motion, camera_direction, GPU_memory_mode, GPU_offload_steps,
@@ -424,6 +544,10 @@ class Ruyi_I2VSampler:
         # Load control embeddings
         embeddings = self.get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction)
 
+        # Initialize plugins
+        plugin_parameters = ruyi_model["plugins"]
+        self.initialize_plugins(pipeline, plugin_parameters, steps)
+
         # Inference
         with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype):
             video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1
@@ -466,11 +590,15 @@ class Ruyi_I2VSampler:
 NODE_CLASS_MAPPINGS = {
     "Ruyi_LoadModel": Ruyi_LoadModel,
     "Ruyi_LoadLora": Ruyi_LoadLora,
+    "Ruyi_TeaCache": Ruyi_TeaCache,
+    "Ruyi_EnhanceAVideo": Ruyi_EnhanceAVideo,
     "Ruyi_I2VSampler": Ruyi_I2VSampler,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "Ruyi_LoadModel": "Load Model",
     "Ruyi_LoadLora": "Load LoRA",
+    "Ruyi_TeaCache": "TeaCache for Ruyi",
+    "Ruyi_EnhanceAVideo": "Enhance-A-Video for Ruyi",
     "Ruyi_I2VSampler": "Sampler for Image to Video",
 }

predict_i2v.py

@@ -61,6 +61,19 @@ lora_weight         = 1.0
 weight_dtype = torch.bfloat16
 device = torch.device("cuda")
 
+# TeaCache settings
+tea_cache_enabled                  = False
+tea_cache_threshold                = 0.10   # A smaller threshold results in fewer cached steps. 0.10 caches 6 ~ 8 steps, 0.15 caches 10 ~ 12 steps normally.
+tea_cache_skip_start_steps         = 3      # First n steps do not use TeaCache, should be >= 1
+tea_cache_skip_end_steps           = 1      # Last n steps do not use TeaCache, should be >= 1
+tea_cache_offload_cpu              = True   # Offload TeaCache tensors to cpu, which could save some GPU memory
+
+# Enhance-A-Video settings
+enhance_a_video_enabled            = False
+enhance_a_video_weight             = 1.0    # Should be smaller than 10. For smaller video length and lower resolution, should be smaller than 5.
+enhance_a_video_skip_start_steps   = 0      # First n steps do not use Enhance-A-Video, should be >= 0
+enhance_a_video_skip_end_steps     = 0      # Last n steps do not use Enhance-A-Video, should be >= 0
+
 
 def get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction):
     # Default keys
@@ -250,6 +263,13 @@ generator= torch.Generator(device).manual_seed(seed)
 # Load control embeddings
 embeddings = get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction)
 
+# Initialize TeaCache
+pipeline.transformer.tea_cache.initialize(tea_cache_enabled, tea_cache_threshold, tea_cache_skip_start_steps, tea_cache_skip_end_steps, steps, tea_cache_offload_cpu)
+
+# Initialize Enhance-A-Video
+pipeline.transformer.enhance_a_video.initialize(enhance_a_video_enabled, enhance_a_video_weight, enhance_a_video_skip_start_steps, enhance_a_video_skip_end_steps, steps)
+
+# Generate video
 with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype):
     video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1
     input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width))
@@ -283,6 +303,10 @@ with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.d
     for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])):
         pipeline = unmerge_lora(pipeline, _lora_path, _lora_weight)
 
+# Log information
+if tea_cache_enabled:
+    print("TeaCache cached steps:", pipeline.transformer.tea_cache.skip_count)
+
 # Save the video
 output_folder = os.path.dirname(output_video_path)
 if output_folder != '':

predict_i2v_80g.py

@@ -61,6 +61,19 @@ lora_weight         = 1.0
 weight_dtype = torch.bfloat16
 device = torch.device("cuda")
 
+# TeaCache settings
+tea_cache_enabled                  = False
+tea_cache_threshold                = 0.10   # A smaller threshold results in fewer cached steps. 0.10 caches 6 ~ 8 steps, 0.15 caches 10 ~ 12 steps normally.
+tea_cache_skip_start_steps         = 3      # First n steps do not use TeaCache, should be >= 1
+tea_cache_skip_end_steps           = 1      # Last n steps do not use TeaCache, should be >= 1
+tea_cache_offload_cpu              = False  # Offload TeaCache tensors to cpu, which could save some GPU memory
+
+# Enhance-A-Video settings
+enhance_a_video_enabled            = False
+enhance_a_video_weight             = 1.0    # Should be smaller than 10. For smaller video length and lower resolution, should be smaller than 5.
+enhance_a_video_skip_start_steps   = 0      # First n steps do not use Enhance-A-Video, should be >= 0
+enhance_a_video_skip_end_steps     = 0      # Last n steps do not use Enhance-A-Video, should be >= 0
+
 
 def get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction):
     # Default keys
@@ -250,6 +263,13 @@ generator= torch.Generator(device).manual_seed(seed)
 # Load control embeddings
 embeddings = get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction)
 
+# Initialize TeaCache
+pipeline.transformer.tea_cache.initialize(tea_cache_enabled, tea_cache_threshold, tea_cache_skip_start_steps, tea_cache_skip_end_steps, steps, tea_cache_offload_cpu)
+
+# Initialize Enhance-A-Video
+pipeline.transformer.enhance_a_video.initialize(enhance_a_video_enabled, enhance_a_video_weight, enhance_a_video_skip_start_steps, enhance_a_video_skip_end_steps, steps)
+
+# Generate video
 with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype):
     video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1
     input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width))
@@ -283,6 +303,10 @@ with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.d
     for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])):
         pipeline = unmerge_lora(pipeline, _lora_path, _lora_weight)
 
+# Log information
+if tea_cache_enabled:
+    print("TeaCache cached steps:", pipeline.transformer.tea_cache.skip_count)
+
 # Save the video
 output_folder = os.path.dirname(output_video_path)
 if output_folder != '':

ruyi/models/attention.py

@@ -39,6 +39,7 @@ from torch import nn
 
 from .motion_module import PositionalEncoding, get_motion_module
 from .norm import FP32LayerNorm, AdaLayerNormShift
+from .attention_processor import HunyuanAttnProcessor2_0_EnhanceAVideo
 
 if is_xformers_available():
     import xformers
@@ -1438,7 +1439,7 @@ class HunyuanDiTBlock(nn.Module):
             qk_norm="layer_norm" if qk_norm else None,
             eps=1e-6,
             bias=True,
-            processor=HunyuanAttnProcessor2_0(),
+            processor=HunyuanAttnProcessor2_0_EnhanceAVideo(),
         )
 
         # 2. Cross-Attn
@@ -1549,6 +1550,10 @@ class HunyuanDiTBlock(nn.Module):
         height: int = 32,
         width: int = 32,
         clip_encoder_hidden_states: Optional[torch.Tensor] = None,
+        # ========== Enhance-A-Video ==========
+        enhance_a_video_enabled: bool = False,
+        enhance_a_video_weight: float = 0.0,
+        # ========== Enhance-A-Video ==========
     ) -> torch.Tensor:
         # Notice that normalization is always applied before the real computation in the following blocks.
         # 0. Long Skip Connection
@@ -1577,6 +1582,11 @@ class HunyuanDiTBlock(nn.Module):
             attn_output = self.attn1(
                 norm_hidden_states_1,
                 image_rotary_emb=attn1_image_rotary_emb,
+                # ========== Enhance-A-Video ==========
+                enhance_a_video_enabled=enhance_a_video_enabled,
+                enhance_a_video_weight=enhance_a_video_weight,
+                num_frames=num_frames,
+                # ========== Enhance-A-Video ==========
             )
             attn_output = rearrange(attn_output, "(b f) (p d) c -> b (f p) d c", p = 2, f = num_frames // 2)
 
@@ -1586,6 +1596,11 @@ class HunyuanDiTBlock(nn.Module):
             attn_output_2 = self.attn1(
                 norm_hidden_states_2,
                 image_rotary_emb=attn1_image_rotary_emb,
+                # ========== Enhance-A-Video ==========
+                enhance_a_video_enabled=enhance_a_video_enabled,
+                enhance_a_video_weight=enhance_a_video_weight,
+                num_frames=num_frames,
+                # ========== Enhance-A-Video ==========
             )
             attn_output_2 = rearrange(attn_output_2, "(b f) (p d) c -> b (f p) d c", p = 2, f = local_attention_frames_num)
             attn_output[:, 1:-1] = (attn_output[:, 1:-1] + attn_output_2) / 2
@@ -1595,6 +1610,11 @@ class HunyuanDiTBlock(nn.Module):
             attn_output = self.attn1(
                 norm_hidden_states,
                 image_rotary_emb=image_rotary_emb,
+                # ========== Enhance-A-Video ==========
+                enhance_a_video_enabled=enhance_a_video_enabled,
+                enhance_a_video_weight=enhance_a_video_weight,
+                num_frames=num_frames,
+                # ========== Enhance-A-Video ==========
             )
         hidden_states = hidden_states + attn_output
 

ruyi/models/attention_processor.py

@@ -0,0 +1,164 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from diffusers.models.attention import Attention
+from diffusers.models.embeddings import apply_rotary_emb
+
+
+def enhance_score(query_image, key_image, head_dim, num_frames, enhance_weight, is_return_attention: bool = False):
+    scale = head_dim**-0.5
+    query_image = query_image * scale
+    attn_temp = query_image @ key_image.transpose(-2, -1)  # translate attn to float32
+    attn_temp = attn_temp.to(torch.float32)
+    attn_temp = attn_temp.softmax(dim=-1)
+
+    # Reshape to [batch_size * num_tokens, num_frames, num_frames]
+    attn_temp = attn_temp.reshape(-1, num_frames, num_frames)
+
+    # Create a mask for diagonal elements
+    diag_mask = torch.eye(num_frames, device=attn_temp.device).bool()
+    diag_mask = diag_mask.unsqueeze(0).expand(attn_temp.shape[0], -1, -1)
+
+    # Zero out diagonal elements
+    attn_wo_diag = attn_temp.masked_fill(diag_mask, 0)
+
+    # Calculate mean for each token's attention matrix
+    # Number of off-diagonal elements per matrix is n*n - n
+    num_off_diag = num_frames * num_frames - num_frames
+    mean_scores = attn_wo_diag.sum(dim=(1, 2)) / num_off_diag
+
+    enhance_scores = mean_scores.mean() * (num_frames + enhance_weight)
+    enhance_scores = enhance_scores.clamp(min=1)
+
+    # if enhance_scores > 1:
+    #     print(f"ENHANCE! {enhance_scores}")
+
+    if is_return_attention:
+        return enhance_scores, attn_temp
+    else:
+        return enhance_scores
+
+
+class HunyuanAttnProcessor2_0_EnhanceAVideo:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
+    used in the HunyuanDiT model. It applies a s normalization layer and rotary embedding on query and key vector.
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def get_enhance_scores(self, attn, query, key, num_frames, enhance_weight, is_return_attention=False):
+        batch_size, num_heads, ST, head_dim = query.shape
+        spatial_dim = ST / num_frames
+        spatial_dim = int(spatial_dim)
+
+        query_image = rearrange(
+            query, "B N (T S) C -> (B S) N T C", T=num_frames, S=spatial_dim, N=num_heads, C=head_dim
+        )
+        key_image = rearrange(key, "B N (T S) C -> (B S) N T C", T=num_frames, S=spatial_dim, N=num_heads, C=head_dim)
+
+        return enhance_score(query_image, key_image, head_dim, num_frames, enhance_weight, is_return_attention)
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        # ========== Enhance-A-Video ==========
+        enhance_a_video_enabled: bool = False,
+        enhance_a_video_weight: float = 0.0,
+        num_frames: int = 1,
+        # ========== Enhance-A-Video ==========
+    ) -> torch.Tensor:
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # ========== Enhance-A-Video ==========
+        if enhance_a_video_enabled:
+            enhance_scores = self.get_enhance_scores(attn, query, key, num_frames, enhance_a_video_weight)
+        # ========== Enhance-A-Video ==========
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb)
+            if not attn.is_cross_attention:
+                key = apply_rotary_emb(key, image_rotary_emb)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        # ========== Enhance-A-Video ==========
+        if enhance_a_video_enabled:
+            hidden_states = hidden_states * enhance_scores
+        # ========== Enhance-A-Video ==========
+
+        return hidden_states

ruyi/models/transformer3d.py

@@ -962,6 +962,194 @@ class HunyuanTransformer3DModel(ModelMixin, ConfigMixin):
 
         self.gradient_checkpointing = False
         self.hidden_cache_size = 0
+
+        # plugins
+        self.tea_cache = HunyuanTransformer3DModel.TeaCacheParameters()
+        self.enhance_a_video = HunyuanTransformer3DModel.EnhanceAVideoParameters()
+    
+    class TeaCacheParameters():
+        coefficients = [-7.51185620e+03, 2.32931511e+03, -2.25610663e+02, 9.39902237e+00, -7.43067848e-02]
+
+        def __init__(self):
+            # basic
+            self.enabled = False
+            self.rel_l1_threshold = 0.0
+            self.skip_start_steps = 1
+            self.skip_end_steps = 1
+
+            # count
+            self.count = 0
+            self.skip_count = 0
+            
+            # internal
+            self.num_steps = 0
+            self.accumulated_rel_l1_distance = 0
+            self.previous_modulated_input = None
+            self.previous_residual = None
+            self.should_calculate = False
+            self.hidden_states_input = None
+
+            # offload to cpu
+            self.offload_cpu = False
+
+            # rescale func
+            self.rescale_func = np.poly1d(self.coefficients)
+        
+        def initialize(self, is_enabled, threshold, skip_start_steps, skip_end_steps, num_steps, offload_cpu=True):
+            self.enabled = is_enabled
+            self.rel_l1_threshold = threshold
+            self.skip_start_steps = max(1, skip_start_steps)
+            self.skip_end_steps = max(1, skip_end_steps)
+
+            self.count = 0
+            self.skip_count = 0
+            
+            self.num_steps = num_steps
+            self.accumulated_rel_l1_distance = 0
+            self.previous_modulated_input = None
+            self.previous_residual = None
+            self.should_calculate = False
+            self.hidden_states_input = None
+
+            self.offload_cpu = offload_cpu
+        
+        def disable(self):
+            self.enabled = False
+        
+        def clean(self):
+            self.previous_modulated_input = None
+            self.previous_residual = None
+            self.hidden_states_input = None
+        
+        def update(self, blocks_first, temb, hidden_states, video_length, height, width):
+            if self.count < self.skip_start_steps - 1:
+                # skip start steps
+                self.should_calculate = True
+
+            elif self.count == self.skip_start_steps - 1:
+                # start TeaCache in next step
+                self.should_calculate = True
+                self.accumulated_rel_l1_distance = 0
+
+                tea_cache_modulated_inp = self.compute_timestep_modulated_input(blocks_first, temb, hidden_states, video_length, height, width)
+                self.store_previouse_modulated_input(tea_cache_modulated_inp)
+                self.store_hidden_states_input(hidden_states)
+
+            elif self.count < self.num_steps - self.skip_end_steps:
+                # using TeaCache
+                tea_cache_modulated_inp = self.compute_timestep_modulated_input(blocks_first, temb, hidden_states, video_length, height, width)
+                self.accumulated_rel_l1_distance += self.compute_timestep_modulated_input_rel_l1_distance(tea_cache_modulated_inp)
+                if self.accumulated_rel_l1_distance < self.rel_l1_threshold:
+                    self.should_calculate = False
+                else:
+                    self.should_calculate = True
+                    self.accumulated_rel_l1_distance = 0
+
+                    self.store_previouse_modulated_input(tea_cache_modulated_inp)
+                    self.store_hidden_states_input(hidden_states.clone())
+
+            else:
+                # skip end steps
+                self.should_calculate = True
+
+            # Update step count
+            self.count += 1
+            if self.count == self.num_steps:
+                self.count = 0
+                self.clean()
+        
+        def compute_timestep_modulated_input(self, blocks_first, temb, hidden_states, video_length, height, width):
+            # Compute timestep embedding modulated input
+            tea_cache_temb = temb.clone()
+            tea_cache_inp = hidden_states.clone()
+
+            if video_length != 1:
+                # add time embedding
+                tea_cache_inp = rearrange(tea_cache_inp, "b (f d) c -> (b d) f c", f=video_length)
+                if blocks_first.t_embed is not None:
+                    tea_cache_inp = blocks_first.t_embed(tea_cache_inp)
+                tea_cache_inp = rearrange(tea_cache_inp, "(b d) f c -> b (f d) c", d=height * width)
+
+            tea_cache_modulated_inp = blocks_first.norm1(tea_cache_inp, tea_cache_temb)
+            return tea_cache_modulated_inp
+        
+        def compute_timestep_modulated_input_rel_l1_distance(self, current_modulated_inp):
+            if self.offload_cpu:
+                device = current_modulated_inp.device
+                previous_modulated_input = self.previous_modulated_input.to(device)
+            else:
+                previous_modulated_input = self.previous_modulated_input
+
+            return self.rescale_func(((current_modulated_inp - previous_modulated_input).abs().mean() / previous_modulated_input.abs().mean()).cpu().item())
+
+        def store_previouse_modulated_input(self, timestep_modulated_input):
+            if self.offload_cpu:
+                self.previous_modulated_input = timestep_modulated_input.cpu()
+            else:
+                self.previous_modulated_input = timestep_modulated_input
+
+        def should_update_residual(self):
+            return self.enabled and self.should_calculate and self.hidden_states_input is not None
+
+        def store_hidden_states_input(self, hidden_states):
+            if self.offload_cpu:
+                self.hidden_states_input = hidden_states.cpu()
+            else:
+                self.hidden_states_input = hidden_states
+        
+        def update_residual(self, hidden_states):
+            if self.offload_cpu:
+                self.previous_residual = hidden_states - self.hidden_states_input.to(hidden_states.device)
+                self.previous_residual = self.previous_residual.cpu()
+            else:
+                self.previous_residual = hidden_states - self.hidden_states_input
+            self.hidden_states_input = None
+        
+        def should_use_cached_residual(self):
+            return self.enabled and not self.should_calculate
+        
+        def apply_cached_residual(self, hidden_states):
+            self.skip_count += 1
+
+            if self.offload_cpu:
+                return hidden_states + self.previous_residual.to(hidden_states.device)
+            else:
+                return hidden_states + self.previous_residual
+    
+    class EnhanceAVideoParameters():
+        def __init__(self):
+            self.enabled = False
+            self.weight = 0.0
+
+            self.is_current_step_enabled = False
+
+            self.count = 0
+            self.num_steps = 0
+            self.skip_start_steps = 0
+            self.skip_end_steps = 0
+        
+        def initialize(self, is_enabled, weight, skip_start_steps, skip_end_steps, num_steps):
+            self.enabled = is_enabled
+            self.weight = weight
+
+            self.is_current_step_enabled = False
+
+            self.count = 0
+            self.num_steps = num_steps
+            self.skip_start_steps = max(0, skip_start_steps)
+            self.skip_end_steps = max(0, skip_end_steps)
+        
+        def disable(self):
+            self.enabled = False
+            self.is_current_step_enabled = False
+        
+        def update(self):
+            if self.enabled and self.count >= self.skip_start_steps and self.count < (self.num_steps - self.skip_end_steps):
+                self.is_current_step_enabled = True
+            else:
+                self.is_current_step_enabled = False
+
+            self.count += 1
     
     def _set_gradient_checkpointing(self, module, value=False):
         if hasattr(module, "gradient_checkpointing"):
@@ -1053,120 +1241,145 @@ class HunyuanTransformer3DModel(ModelMixin, ConfigMixin):
         if skips_cache_size <= 0:
             skips_cache_size = self.config.num_layers
         
-        for layer, block in enumerate(self.blocks):
-            if layer > self.config.num_layers // 2:
-                if skips_cache_size == 1:
-                    skip_cpu = skips.pop()
-                    skip = skip_cpu.to(hidden_states.device)
-                else:
-                    if len(skips) == 0:
-                        skips_cache_block = skips_cpu_cache.pop()
-                        for si in range(len(skips_cache_block)):
-                            skips.append(skips_cache_block[si].to(hidden_states.device))
-                    skip = skips.pop()
+        # Enhance-A-Video
+        if self.enhance_a_video.enabled:
+            self.enhance_a_video.update()
+        
+        # TeaCache
+        if self.tea_cache.enabled:
+            self.tea_cache.update(self.blocks[0], temb, hidden_states, video_length, height, width)
+        
+        if self.tea_cache.should_use_cached_residual():
+            hidden_states = self.tea_cache.apply_cached_residual(hidden_states)
+        else:
+            # Blocks
+            for layer, block in enumerate(self.blocks):
+                if layer > self.config.num_layers // 2:
+                    if skips_cache_size == 1:
+                        skip_cpu = skips.pop()
+                        skip = skip_cpu.to(hidden_states.device)
+                    else:
+                        if len(skips) == 0:
+                            skips_cache_block = skips_cpu_cache.pop()
+                            for si in range(len(skips_cache_block)):
+                                skips.append(skips_cache_block[si].to(hidden_states.device))
+                        skip = skips.pop()
 
-                if self.training and self.gradient_checkpointing:
+                    if self.training 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)
+                        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
+                            return custom_forward
 
-                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                    args = {
-                        "basic": [video_length, height, width, clip_encoder_hidden_states],
-                        "hybrid_attention": [video_length, height, width, clip_encoder_hidden_states],
-                        "motionmodule": [video_length, height, width],
-                        "global_motionmodule": [video_length, height, width],
-                    }[self.basic_block_type]
-                    
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(block),
-                        hidden_states,
-                        encoder_hidden_states,
-                        temb,
-                        image_rotary_emb,
-                        skip,
-                        *args,
-                        **ckpt_kwargs,
-                    )
-                else:
-                    kwargs = {
-                        "basic": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
-                        "hybrid_attention": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
-                        "motionmodule": {"num_frames":video_length, "height":height, "width":width},
-                        "global_motionmodule": {"num_frames":video_length, "height":height, "width":width},
-                    }[self.basic_block_type]
-                    hidden_states = block(
-                        hidden_states,
-                        temb=temb,
-                        encoder_hidden_states=encoder_hidden_states,
-                        image_rotary_emb=image_rotary_emb,  
-                        skip=skip,                  
-                        **kwargs
-                    )  # (N, L, D)
-            else:
-                if self.training 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 {}
-                    args = {
-                        "basic": [None, video_length, height, width, clip_encoder_hidden_states],
-                        "hybrid_attention": [None, video_length, height, width, clip_encoder_hidden_states],
-                        "motionmodule": [None, video_length, height, width],
-                        "global_motionmodule": [None, video_length, height, width],
-                    }[self.basic_block_type]
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(block),
-                        hidden_states,
-                        encoder_hidden_states,
-                        temb,
-                        image_rotary_emb, 
-                        *args,
-                        **ckpt_kwargs,
-                    )
-                else:
-                    kwargs = {
-                        "basic": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
-                        "hybrid_attention": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
-                        "motionmodule": {"num_frames":video_length, "height":height, "width":width},
-                        "global_motionmodule": {"num_frames":video_length, "height":height, "width":width},
-                    }[self.basic_block_type]
-                    hidden_states = block(
-                        hidden_states,
-                        temb=temb,
-                        encoder_hidden_states=encoder_hidden_states,
-                        image_rotary_emb=image_rotary_emb,  
-                        **kwargs
-                    )  # (N, L, D)
-
-            if layer < (self.config.num_layers // 2 - 1):
-                if skips_cache_size == 1:
-                    skips.append(hidden_states.to("cpu"))
-                else:
-                    skips.append(hidden_states)
-                    if len(skips) >= skips_cache_size:
-                        skips_cache_block = []
-                        for si in range(len(skips)):
-                            skips_cache_block.append(skips[si].to("cpu"))
+                        ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                        args = {
+                            "basic": [video_length, height, width, clip_encoder_hidden_states],
+                            "hybrid_attention": [video_length, height, width, clip_encoder_hidden_states],
+                            "motionmodule": [video_length, height, width],
+                            "global_motionmodule": [video_length, height, width],
+                        }[self.basic_block_type]
                         
-                        skips = []
-                        skips_cpu_cache.append(skips_cache_block)
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            hidden_states,
+                            encoder_hidden_states,
+                            temb,
+                            image_rotary_emb,
+                            skip,
+                            *args,
+                            **ckpt_kwargs,
+                        )
+                    else:
+                        kwargs = {
+                            "basic": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
+                            "hybrid_attention": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
+                            "motionmodule": {"num_frames":video_length, "height":height, "width":width},
+                            "global_motionmodule": {"num_frames":video_length, "height":height, "width":width},
+                        }[self.basic_block_type]
+                        hidden_states = block(
+                            hidden_states,
+                            temb=temb,
+                            encoder_hidden_states=encoder_hidden_states,
+                            image_rotary_emb=image_rotary_emb,  
+                            skip=skip,                  
+
+                            enhance_a_video_enabled=self.enhance_a_video.is_current_step_enabled,
+                            enhance_a_video_weight=self.enhance_a_video.weight,
+
+                            **kwargs
+                        )  # (N, L, D)
+                else:
+                    if self.training 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 {}
+                        args = {
+                            "basic": [None, video_length, height, width, clip_encoder_hidden_states],
+                            "hybrid_attention": [None, video_length, height, width, clip_encoder_hidden_states],
+                            "motionmodule": [None, video_length, height, width],
+                            "global_motionmodule": [None, video_length, height, width],
+                        }[self.basic_block_type]
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            hidden_states,
+                            encoder_hidden_states,
+                            temb,
+                            image_rotary_emb, 
+                            *args,
+                            **ckpt_kwargs,
+                        )
+                    else:
+                        kwargs = {
+                            "basic": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
+                            "hybrid_attention": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states},
+                            "motionmodule": {"num_frames":video_length, "height":height, "width":width},
+                            "global_motionmodule": {"num_frames":video_length, "height":height, "width":width},
+                        }[self.basic_block_type]
+                        hidden_states = block(
+                            hidden_states,
+                            temb=temb,
+                            encoder_hidden_states=encoder_hidden_states,
+                            image_rotary_emb=image_rotary_emb,  
+
+                            enhance_a_video_enabled=self.enhance_a_video.is_current_step_enabled,
+                            enhance_a_video_weight=self.enhance_a_video.weight,
+                            
+                            **kwargs
+                        )  # (N, L, D)
+
+                if layer < (self.config.num_layers // 2 - 1):
+                    if skips_cache_size == 1:
+                        skips.append(hidden_states.to("cpu"))
+                    else:
+                        skips.append(hidden_states)
+                        if len(skips) >= skips_cache_size:
+                            skips_cache_block = []
+                            for si in range(len(skips)):
+                                skips_cache_block.append(skips[si].to("cpu"))
+                            
+                            skips = []
+                            skips_cpu_cache.append(skips_cache_block)
+
+            # final layer
+            hidden_states = self.norm_out(hidden_states, temb.to(torch.float32))
+
+        # update TeaCache
+        if self.tea_cache.should_update_residual():
+            self.tea_cache.update_residual(hidden_states)
 
-        # final layer
-        hidden_states = self.norm_out(hidden_states, temb.to(torch.float32))
         hidden_states = self.proj_out(hidden_states)
         # (N, L, patch_size ** 2 * out_channels)
 
@@ -1217,11 +1430,11 @@ class HunyuanTransformer3DModel(ModelMixin, ConfigMixin):
                 model.state_dict()['proj_out.bias'][:state_dict['proj_out.bias'].size()[0]] = state_dict['proj_out.bias']
                 state_dict['proj_out.bias'] = model.state_dict()['proj_out.bias']
 
-        tmp_state_dict = {} 
-        for key in state_dict:
-            if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size():
-                tmp_state_dict[key] = state_dict[key]
-        state_dict = tmp_state_dict
+        # tmp_state_dict = {} 
+        # for key in state_dict:
+        #     if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size():
+        #         tmp_state_dict[key] = state_dict[key]
+        # state_dict = tmp_state_dict
 
         m, u = model.load_state_dict(state_dict, strict=False)
         # print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")