From 6eecfc9c106f592c8ea0cd77eb4f9a79c5c4ec47 Mon Sep 17 00:00:00 2001 From: WarpedAnimation Date: Thu, 3 Jul 2025 14:09:38 -0500 Subject: [PATCH] v1.7.0 --- .../models/hunyuan_video_packed.py | 188 +++++++++++++++++- 1 file changed, 187 insertions(+), 1 deletion(-) diff --git a/diffusers_helper/models/hunyuan_video_packed.py b/diffusers_helper/models/hunyuan_video_packed.py index 03424a1..c5da846 100644 --- a/diffusers_helper/models/hunyuan_video_packed.py +++ b/diffusers_helper/models/hunyuan_video_packed.py @@ -774,6 +774,9 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, self.inner_dim = inner_dim self.use_gradient_checkpointing = False self.enable_teacache = False + self.enable_magcache = False + self.verbose_messaging = False + self.orig_forward = self.forward if has_image_proj: self.install_image_projection(image_proj_dim) @@ -800,7 +803,7 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, self.use_gradient_checkpointing = False print('self.use_gradient_checkpointing = False') - def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_l1_thresh=0.15): + def initialize_teacache(self, enable_teacache=True, num_steps=25, rel_l1_thresh=0.15, verbose_messaging=False): self.enable_teacache = enable_teacache self.cnt = 0 self.num_steps = num_steps @@ -809,6 +812,32 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, self.previous_modulated_input = None self.previous_residual = None self.teacache_rescale_func = np.poly1d([7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02]) + self.verbose_messaging = verbose_messaging + + def nearest_interp(self, src_array, target_length): + src_length = len(src_array) + if target_length == 1: + return np.array([src_array[-1]]) + + scale = (src_length - 1) / (target_length - 1) + mapped_indices = np.round(np.arange(target_length) * scale).astype(int) + return src_array[mapped_indices] + + def initialize_magcache(self, enable_magcache=True, num_steps=25, magcache_thresh=0.1, K=2, retention_ratio=0.2, verbose_messaging=False): + self.enable_magcache = enable_magcache + self.cnt = 0 + self.num_steps = num_steps + self.magcache_thresh = magcache_thresh + self.K = K + self.retention_ratio = retention_ratio + self.mag_ratios = np.array([1.0]+[1.25781, 1.08594, 1.02344, 1.00781, 1.02344, 1.00781, 1.02344, 1.05469, 0.99609, 1.03906, 1.00781, 1.01562, 1.00781, 1.02344, 1.01562, 0.98047, 1.05469, 0.98047, 0.96875, 1.03125, 0.97266, 0.9375, 0.96484, 0.78516]) + + # Nearest interpolation when the num_steps is different from the length of mag_ratios + if len(self.mag_ratios) != num_steps: + interpolated_mag_ratios = self.nearest_interp(self.mag_ratios, num_steps) + self.mag_ratios = interpolated_mag_ratios + + self.verbose_messaging = verbose_messaging def gradient_checkpointing_method(self, block, *args): if self.use_gradient_checkpointing: @@ -876,6 +905,162 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, return hidden_states, rope_freqs + def magcache_framepack_forward( + self, + hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, pooled_projections, guidance, + latent_indices=None, + clean_latents=None, clean_latent_indices=None, + clean_latents_2x=None, clean_latent_2x_indices=None, + clean_latents_4x=None, clean_latent_4x_indices=None, + image_embeddings=None, + attention_kwargs=None, return_dict=True + ): + + if self.verbose_messaging: + print("\nMagCache Forward Entry...") + print("self.enable_magcache: {}".format(self.enable_magcache)) + print("self.cnt: {}".format(self.cnt)) + print("self.accumulated_ratio: {}".format(self.accumulated_ratio)) + print("self.accumulated_steps: {}".format(self.accumulated_steps)) + print("self.accumulated_err: {}".format(self.accumulated_err)) + + if attention_kwargs is None: + attention_kwargs = {} + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p, p_t = self.config['patch_size'], self.config['patch_size_t'] + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p + post_patch_width = width // p + original_context_length = post_patch_num_frames * post_patch_height * post_patch_width + + hidden_states, rope_freqs = self.process_input_hidden_states(hidden_states, latent_indices, clean_latents, clean_latent_indices, clean_latents_2x, clean_latent_2x_indices, clean_latents_4x, clean_latent_4x_indices) + + temb = self.gradient_checkpointing_method(self.time_text_embed, timestep, guidance, pooled_projections) + encoder_hidden_states = self.gradient_checkpointing_method(self.context_embedder, encoder_hidden_states, timestep, encoder_attention_mask) + + if self.image_projection is not None: + assert image_embeddings is not None, 'You must use image embeddings!' + extra_encoder_hidden_states = self.gradient_checkpointing_method(self.image_projection, image_embeddings) + extra_attention_mask = torch.ones((batch_size, extra_encoder_hidden_states.shape[1]), dtype=encoder_attention_mask.dtype, device=encoder_attention_mask.device) + + # must cat before (not after) encoder_hidden_states, due to attn masking + encoder_hidden_states = torch.cat([extra_encoder_hidden_states, encoder_hidden_states], dim=1) + encoder_attention_mask = torch.cat([extra_attention_mask, encoder_attention_mask], dim=1) + + if batch_size == 1: + # When batch size is 1, we do not need any masks or var-len funcs since cropping is mathematically same to what we want + # If they are not same, then their impls are wrong. Ours are always the correct one. + text_len = encoder_attention_mask.sum().item() + encoder_hidden_states = encoder_hidden_states[:, :text_len] + attention_mask = None, None, None, None + else: + img_seq_len = hidden_states.shape[1] + txt_seq_len = encoder_hidden_states.shape[1] + + cu_seqlens_q = get_cu_seqlens(encoder_attention_mask, img_seq_len) + cu_seqlens_kv = cu_seqlens_q + max_seqlen_q = img_seq_len + txt_seq_len + max_seqlen_kv = max_seqlen_q + + attention_mask = cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv + + if self.enable_magcache: + if self.cnt == 0: # initialize MagCache + self.accumulated_ratio = 1.0 + self.accumulated_steps = 0 + self.accumulated_err = 0 + + skip_forward = False + if self.cnt>=int(self.retention_ratio*self.num_steps) and self.cnt>=1: # keep first retention_ratio steps + cur_mag_ratio = self.mag_ratios[self.cnt] + self.accumulated_ratio = self.accumulated_ratio*cur_mag_ratio + cur_skip_err = np.abs(1-self.accumulated_ratio) + self.accumulated_err += cur_skip_err + self.accumulated_steps += 1 + if self.accumulated_err<=self.magcache_thresh and self.accumulated_steps<=self.K and np.abs(1-cur_mag_ratio)<=0.06: + skip_forward = True + else: + self.accumulated_ratio = 1.0 + self.accumulated_steps = 0 + self.accumulated_err = 0 + + if skip_forward: + hidden_states = hidden_states + self.previous_residual + else: + ori_hidden_states = hidden_states.clone() + + for block_id, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + for block_id, block in enumerate(self.single_transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + self.previous_residual = hidden_states - ori_hidden_states + self.cnt += 1 + if self.cnt == self.num_steps: + self.cnt = 0 + else: + for block_id, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + for block_id, block in enumerate(self.single_transformer_blocks): + hidden_states, encoder_hidden_states = self.gradient_checkpointing_method( + block, + hidden_states, + encoder_hidden_states, + temb, + attention_mask, + rope_freqs + ) + + hidden_states = self.gradient_checkpointing_method(self.norm_out, hidden_states, temb) + + hidden_states = hidden_states[:, -original_context_length:, :] + + if self.high_quality_fp32_output_for_inference: + hidden_states = hidden_states.to(dtype=torch.float32) + if self.proj_out.weight.dtype != torch.float32: + self.proj_out.to(dtype=torch.float32) + + hidden_states = self.gradient_checkpointing_method(self.proj_out, hidden_states) + + hidden_states = einops.rearrange(hidden_states, 'b (t h w) (c pt ph pw) -> b c (t pt) (h ph) (w pw)', + t=post_patch_num_frames, h=post_patch_height, w=post_patch_width, + pt=p_t, ph=p, pw=p) + + if return_dict: + if self.verbose_messaging: + print("MagCache Forward Exit.") + + return Transformer2DModelOutput(sample=hidden_states) + + if self.verbose_messaging: + print("MagCache Forward Exit.") + + return hidden_states, + def forward( self, hidden_states, timestep, encoder_hidden_states, encoder_attention_mask, pooled_projections, guidance, @@ -1014,3 +1199,4 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, return Transformer2DModelOutput(sample=hidden_states) return hidden_states, +