v1.7.0
This commit is contained in:
@@ -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,
|
||||
|
||||
|
||||
Reference in New Issue
Block a user