From 35bc2af7eeb39738c90d511eb3591ae1597fdb00 Mon Sep 17 00:00:00 2001 From: Yuan-Man <68322456+Yuan-ManX@users.noreply.github.com> Date: Tue, 13 May 2025 14:35:54 +0800 Subject: [PATCH] Create teacache_forward.py --- teacache_forward.py | 353 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 353 insertions(+) create mode 100644 teacache_forward.py diff --git a/teacache_forward.py b/teacache_forward.py new file mode 100644 index 0000000..462f035 --- /dev/null +++ b/teacache_forward.py @@ -0,0 +1,353 @@ +# teacache +import torch +import numpy as np +from typing import Optional, Union, Dict, Any + +from matrixgame.model_variants.matrixgame_dit_src.modulate_layers import modulate +from matrixgame.model_variants.matrixgame_dit_src.attenion import attention, get_cu_seqlens +from diffusers.models.modeling_outputs import Transformer2DModelOutput + + +def teacache_forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, # Should be in range(0, 1000). + encoder_hidden_states: torch.Tensor = None, + encoder_attention_mask: torch.Tensor = None, # Now we don't use it. + guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000. + mouse_condition = None, + keyboard_condition = None, + return_dict: bool = True, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + x = hidden_states + t = timestep + text_states, text_states_2 = encoder_hidden_states + text_mask, test_mask_2 = encoder_attention_mask + out = {} + img = x + txt = text_states + _, _, ot, oh, ow = x.shape + freqs_cos, freqs_sin = self.get_rotary_pos_embed(ot, oh, ow) + tt, th, tw = ( + ot // self.patch_size[0], + oh // self.patch_size[1], + ow // self.patch_size[2], + ) + + # Prepare modulation vectors. + vec = self.time_in(t) + if self.i2v_condition_type == "token_replace": + token_replace_t = torch.zeros_like(t) + token_replace_vec = self.time_in(token_replace_t) + frist_frame_token_num = th * tw + else: + token_replace_vec = None + frist_frame_token_num = None + # text modulation + #vec_2 = self.vector_in(text_states_2) + #vec = vec + vec_2 + #if self.i2v_condition_type == "token_replace": + # token_replace_vec = token_replace_vec + vec_2 + + # guidance modulation + if self.guidance_embed: + if guidance is None: + raise ValueError( + "Didn't get guidance strength for guidance distilled model." + ) + + # our timestep_embedding is merged into guidance_in(TimestepEmbedder) + vec = vec + self.guidance_in(guidance) + + # Embed image and text. + img = self.img_in(img) + if self.text_projection == "linear": + txt = self.txt_in(txt) + elif self.text_projection == "single_refiner": + txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None) + else: + raise NotImplementedError( + f"Unsupported text_projection: {self.text_projection}" + ) + + txt_seq_len = txt.shape[1] + img_seq_len = img.shape[1] + + # Compute cu_squlens and max_seqlen for flash attention + cu_seqlens_q = get_cu_seqlens(text_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 + + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + # teacache + if self.enable_teacache: + inp = img.clone() + vec_ = vec.clone() + txt_ = txt.clone() + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1) + normed_inp = self.double_blocks[0].img_norm1(inp) + modulated_inp = modulate( + normed_inp, shift=img_mod1_shift, scale=img_mod1_scale + ) + if self.cnt == 0 or self.cnt == self.num_steps-1: + should_calc = True + self.accumulated_rel_l1_distance = 0 + else: + coefficients = [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02] + #coefficients = [-296.53, 191.67, -39.037, 3.705, -0.0383] + rescale_func = np.poly1d(coefficients) + self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()) + if self.accumulated_rel_l1_distance < self.rel_l1_thresh: + should_calc = False + else: + should_calc = True + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = modulated_inp + self.cnt += 1 + if self.cnt == self.num_steps: + self.cnt = 0 + + if self.enable_teacache: + if not should_calc: + img += self.previous_residual + else: + ori_img = img.clone() + # --------------------- Pass through DiT blocks ------------------------ + for _, block in enumerate(self.double_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + img, txt = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + img, + txt, + vec, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + freqs_cis, + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + **ckpt_kwargs, + ) + else: + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + double_block_args = [ + img, + txt, + vec, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + freqs_cis, + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + ] + + img, txt = block(*double_block_args) + + # Merge txt and img to pass through single stream blocks. + x = torch.cat((img, txt), 1) + if len(self.single_blocks) > 0: + for _, block in enumerate(self.single_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + vec, + txt_seq_len, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + (freqs_cos, freqs_sin), + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + **ckpt_kwargs, + ) + else: + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + single_block_args = [ + x, + vec, + txt_seq_len, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + (freqs_cos, freqs_sin), + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + ] + + x = block(*single_block_args) + + img = x[:, :img_seq_len, ...] + self.previous_residual = img - ori_img + else: + # --------------------- Pass through DiT blocks ------------------------ + for _, block in enumerate(self.double_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + img, txt = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + img, + txt, + vec, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + freqs_cis, + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + **ckpt_kwargs, + ) + else: + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + double_block_args = [ + img, + txt, + vec, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + freqs_cis, + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + ] + + img, txt = block(*double_block_args) + + # Merge txt and img to pass through single stream blocks. + x = torch.cat((img, txt), 1) + if len(self.single_blocks) > 0: + for _, block in enumerate(self.single_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + vec, + txt_seq_len, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + (freqs_cos, freqs_sin), + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + **ckpt_kwargs, + ) + else: + image_kwargs: Dict[str, Any] = {"tt":hidden_states.shape[2] // self.patch_size[0], + "th":hidden_states.shape[3] // self.patch_size[1], + "tw":hidden_states.shape[4] // self.patch_size[2]} + single_block_args = [ + x, + vec, + txt_seq_len, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + (freqs_cos, freqs_sin), + image_kwargs, + mouse_condition, + keyboard_condition, + self.i2v_condition_type, + token_replace_vec, + frist_frame_token_num, + ] + + x = block(*single_block_args) + + img = x[:, :img_seq_len, ...] + + # ---------------------------- Final layer ------------------------------ + img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + + img = self.unpatchify(img, tt, th, tw) + if return_dict: + out["x"] = img + return out + return (img,) +