diff --git a/README.md b/README.md index 6eed831..f217960 100644 --- a/README.md +++ b/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. diff --git a/README_CN.md b/README_CN.md index fbc0cc9..70c5dd7 100644 --- a/README_CN.md +++ b/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比例的视频会出现黑杠的问题。 diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py index 7159cc0..e775c5f 100644 --- a/comfyui/comfyui_nodes.py +++ b/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", } diff --git a/predict_i2v.py b/predict_i2v.py index 6ecb757..c276af0 100644 --- a/predict_i2v.py +++ b/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 != '': diff --git a/predict_i2v_80g.py b/predict_i2v_80g.py index b290d36..182b7fa 100644 --- a/predict_i2v_80g.py +++ b/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 != '': diff --git a/ruyi/models/attention.py b/ruyi/models/attention.py index 8630f0c..c844f3e 100644 --- a/ruyi/models/attention.py +++ b/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 diff --git a/ruyi/models/attention_processor.py b/ruyi/models/attention_processor.py new file mode 100644 index 0000000..1ad637d --- /dev/null +++ b/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 diff --git a/ruyi/models/transformer3d.py b/ruyi/models/transformer3d.py index 1e8b67b..07325a7 100644 --- a/ruyi/models/transformer3d.py +++ b/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)};")