This commit is contained in:
WarpedAnimation
2025-07-03 14:09:38 -05:00
committed by GitHub
parent 1b69f315a4
commit 6eecfc9c10
+187 -1
View File
@@ -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,