From 797cb2f659e04ba470b92ffd142ed1f9e09463b9 Mon Sep 17 00:00:00 2001 From: rlsu9 Date: Thu, 13 Mar 2025 17:08:02 +0000 Subject: [PATCH] recover sample hunyuan --- fastvideo/sample/sample_t2v_hunyuan.py | 195 +------------------------ 1 file changed, 4 insertions(+), 191 deletions(-) diff --git a/fastvideo/sample/sample_t2v_hunyuan.py b/fastvideo/sample/sample_t2v_hunyuan.py index 3c75fdb20..b5836b358 100644 --- a/fastvideo/sample/sample_t2v_hunyuan.py +++ b/fastvideo/sample/sample_t2v_hunyuan.py @@ -1,175 +1,17 @@ import argparse import os from pathlib import Path -import json + import imageio import numpy as np import torch import torch.distributed as dist import torchvision from einops import rearrange -from fastvideo.models.hunyuan.modules.modulate_layers import modulate + from fastvideo.models.hunyuan.inference import HunyuanVideoSampler from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info -from typing import Any, Dict, Optional, Union -def teacache_forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_attention_mask: torch.Tensor, - mask_strategy=None, - output_features=False, - output_features_stride=8, - attention_kwargs: Optional[Dict[str, Any]] = None, - return_dict: bool = False, - guidance=None, -) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if guidance is None: - guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16) - img = x = hidden_states - text_mask = encoder_attention_mask - t = timestep - txt = encoder_hidden_states[:, 1:] - text_states_2 = encoder_hidden_states[:, 0, :self.config.text_states_dim_2] - _, _, ot, oh, ow = x.shape # codespell:ignore - tt, th, tw = ( - ot // self.patch_size[0], # codespell:ignore - oh // self.patch_size[1], # codespell:ignore - ow // self.patch_size[2], # codespell:ignore - ) - original_tt = nccl_info.sp_size * tt - freqs_cos, freqs_sin = self.get_rotary_pos_embed((original_tt, th, tw)) - # Prepare modulation vectors. - vec = self.time_in(t) - - # text modulation - vec = vec + self.vector_in(text_states_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] - - freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None - - if self.enable_teacache: - inp = img.clone() - vec_ = vec.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).to("cpu") - del inp, vec_, img_mod1_shift, img_mod1_scale, normed_inp - - 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] - 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.to(img.device) - self.previous_residual = self.previous_residual.to(img.device) - else: - ori_img = img.clone().to("cpu") - # --------------------- Pass through DiT blocks ------------------------ - for index, block in enumerate(self.double_blocks): - double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]] - img, txt = block(*double_block_args) - - # Merge txt and img to pass through single stream blocks. - x = torch.cat((img, txt), 1) - if output_features: - features_list = [] - if len(self.single_blocks) > 0: - for index, block in enumerate(self.single_blocks): - single_block_args = [ - x, - vec, - txt_seq_len, - (freqs_cos, freqs_sin), - text_mask, - mask_strategy[index + len(self.double_blocks)], - ] - x = block(*single_block_args) - if output_features and _ % output_features_stride == 0: - features_list.append(x[:, :img_seq_len, ...]) - - img = x[:, :img_seq_len, ...] - self.previous_residual = (img.clone().to("cpu") - ori_img).to("cpu") - del ori_img - else: - # --------------------- Pass through DiT blocks ------------------------ - for index, block in enumerate(self.double_blocks): - double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]] - img, txt = block(*double_block_args) - # Merge txt and img to pass through single stream blocks. - x = torch.cat((img, txt), 1) - if output_features: - features_list = [] - if len(self.single_blocks) > 0: - for index, block in enumerate(self.single_blocks): - single_block_args = [ - x, - vec, - txt_seq_len, - (freqs_cos, freqs_sin), - text_mask, - mask_strategy[index + len(self.double_blocks)], - ] - x = block(*single_block_args) - if output_features and _ % output_features_stride == 0: - features_list.append(x[:, :img_seq_len, ...]) - - 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) - assert not return_dict, "return_dict is not supported." - if output_features: - features_list = torch.stack(features_list, dim=0) - else: - features_list = None - - return (img, features_list) def initialize_distributed(): local_rank = int(os.getenv("RANK", 0)) @@ -186,12 +28,6 @@ def main(args): print(args) models_root_path = Path(args.model_path) - - if os.path.exists(args.mask_strategy_file_path): - with open(args.mask_strategy_file_path, 'r') as f: - mask_strategy = json.load(f) - else: - mask_strategy = None if not models_root_path.exists(): raise ValueError(f"`models_root` not exists: {models_root_path}") @@ -204,15 +40,6 @@ def main(args): # Get the updated args args = hunyuan_video_sampler.args - - hunyuan_video_sampler.pipeline.transformer.__class__.enable_teacache = args.enable_teacache - hunyuan_video_sampler.pipeline.transformer.__class__.cnt = 0 - hunyuan_video_sampler.pipeline.transformer.__class__.num_steps = args.num_inference_steps - hunyuan_video_sampler.pipeline.transformer.__class__.rel_l1_thresh = args.rel_l1_thresh # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup - hunyuan_video_sampler.pipeline.transformer.__class__.accumulated_rel_l1_distance = 0 - hunyuan_video_sampler.pipeline.transformer.__class__.previous_modulated_input = None - hunyuan_video_sampler.pipeline.transformer.__class__.previous_residual = None - hunyuan_video_sampler.pipeline.transformer.__class__.forward = teacache_forward if args.prompt.endswith('.txt'): with open(args.prompt) as f: @@ -234,7 +61,6 @@ def main(args): flow_shift=args.flow_shift, batch_size=args.batch_size, embedded_guidance_scale=args.embedded_cfg_scale, - mask_strategy=mask_strategy, ) videos = rearrange(outputs["samples"], "b c t h w -> t b c h w") outputs = [] @@ -336,7 +162,6 @@ if __name__ == "__main__": ) # Model parameters - parser.add_argument("--use-fp8", action='store_true') parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill") parser.add_argument("--latent-channels", type=int, default=16) parser.add_argument("--precision", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) @@ -377,21 +202,9 @@ if __name__ == "__main__": parser.add_argument("--text-states-dim-2", type=int, default=768) parser.add_argument("--tokenizer-2", type=str, default="clipL") parser.add_argument("--text-len-2", type=int, default=77) - parser.add_argument("--vae_tiling", action='store_true') - parser.add_argument("--mask_strategy_file_path", type=str, default="assets/mask_strategy.json") - parser.add_argument( - "--rel_l1_thresh", - type=float, - default=0.15, - help="0.1 for 1.6x speedup, 0.15 for 2.1x speedup", - ) - parser.add_argument( - "--enable_teacache", - action="store_true", - help="Use teacache for speeding up inference", - ) + args = parser.parse_args() # process for vae sequence parallel if args.vae_sp and not args.vae_tiling: raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") - main(args) + main(args) \ No newline at end of file