Coverage for wrapper/hunyuanavatar/sample_inference_audio.py: 18%
118 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1import math
2import time
3import torch
4from typing import Any
5from typing import Dict
6from typing import Optional
7from typing import Tuple
8from loguru import logger
9from einops import rearrange
10from hymm_sp.diffusion import load_diffusion_pipeline
11from hymm_sp.helpers import get_nd_rotary_pos_embed_new
12from hymm_sp.inference import Inference
13from hymm_sp.data_kits.audio_preprocessor import encode_audio, get_facemask
14import logging
17def align_to(value: float, alignment: int) -> int:
18 return int(math.ceil(value / alignment) * alignment)
21class HunyuanVideoSampler(Inference):
22 def __init__(
23 self,
24 args: Any,
25 vae: Any,
26 vae_kwargs: Any,
27 text_encoder: Any,
28 model: Any,
29 text_encoder_2: Any = None,
30 pipeline: Any = None,
31 device: int = 0,
32 logger: Any = None,
33 ) -> None:
34 super().__init__(args, vae, vae_kwargs, text_encoder, model, text_encoder_2=text_encoder_2,
35 pipeline=pipeline, device=device, logger=logger)
37 self.args = args
38 self.pipeline = load_diffusion_pipeline(
39 args, 0, self.vae, self.text_encoder, self.text_encoder_2, self.model,
40 device=self.device)
41 self.pipeline.scheduler.config.shift = 1.0
42 logging.info('Loaded HunyuanVideoSampler successfully... ')
44 def get_rotary_pos_embed(
45 self,
46 video_length: int,
47 height: int,
48 width: int,
49 concat_dict: Dict = {}
50 ) -> Tuple[torch.Tensor, torch.Tensor]:
51 target_ndim = 3
52 ndim = 5 - 2
53 if '884' in self.args.vae:
54 latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8]
55 else:
56 latents_size = [video_length, height // 8, width // 8]
58 if isinstance(self.model.patch_size, int):
59 assert all(s % self.model.patch_size == 0 for s in latents_size), \
60 f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), " \
61 f"but got {latents_size}."
62 rope_sizes = [s // self.model.patch_size for s in latents_size]
63 elif isinstance(self.model.patch_size, list):
64 assert all(s % self.model.patch_size[idx] == 0 for idx, s in enumerate(latents_size)), \
65 f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), " \
66 f"but got {latents_size}."
67 rope_sizes = [s // self.model.patch_size[idx] for idx, s in enumerate(latents_size)]
69 if len(rope_sizes) != target_ndim:
70 rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis
71 head_dim = self.model.hidden_size // self.model.num_heads
72 rope_dim_list = self.model.rope_dim_list
73 if rope_dim_list is None:
74 rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
75 assert sum(rope_dim_list) == head_dim, "sum(rope_dim_list) should equal to head_dim of attention layer"
76 freqs_cos, freqs_sin = get_nd_rotary_pos_embed_new(rope_dim_list,
77 rope_sizes,
78 theta=self.args.rope_theta,
79 use_real=True,
80 theta_rescale_factor=1,
81 concat_dict=concat_dict)
82 return freqs_cos, freqs_sin
84 @torch.no_grad()
85 def predict(
86 self,
87 args: Any,
88 batch: Dict[str, Any],
89 wav2vec: Any,
90 feature_extractor: Any,
91 align_instance: Any,
92 **kwargs: Any,
93 ) -> Optional[Dict[str, Any]]:
94 """
95 Predict the image from the given text.
97 Args:
98 prompt (str or List[str]): The input text.
99 kwargs:
100 size (int): The (height, width) of the output image/video. Default is (256, 256).
101 video_length (int): The frame number of the output video. Default is 1.
102 seed (int or List[str]): The random seed for the generation. Default is a random integer.
103 negative_prompt (str or List[str]): The negative text prompt. Default is an empty string.
104 infer_steps (int): The number of inference steps. Default is 100.
105 guidance_scale (float): The guidance scale for the generation. Default is 6.0.
106 num_videos_per_prompt (int): The number of videos per prompt. Default is 1.
107 verbose (int): 0 for no log, 1 for all log, 2 for fewer log. Default is 1.
108 output_type (str): The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
109 Default is 'pil'.
110 """
112 out_dict = dict()
114 prompt = batch['text_prompt']
115 audio_path = str(batch["audio_path"])
116 neg_prompt = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, "
117 neg_prompt += "bad teeth, bad eyes, bad limbs, distortion, blurring, Lens changes"
118 fps = batch["fps"].to(self.device)
119 # (hqiu) check if fps is a scalar or tensor
120 if fps.dim() == 0:
121 fps = fps.unsqueeze(0)
122 print(f"fps was a scalar, unsqueezed to: {fps}, dim: {fps.dim()}")
123 else:
124 print(f"fps was not a scalar, keeping as: {fps}, dim: {fps.dim()}")
125 audio_prompts = batch["audio_prompts"].to(self.device)
126 weight_dtype = audio_prompts.dtype
128 # audio_prompts = [
129 # encode_audio(wav2vec, audio_feat.to(dtype=wav2vec.dtype), fps.item(), num_frames=batch["audio_len"])
130 # for audio_feat in audio_prompts
131 # ]
132 audio_prompts = encode_audio(wav2vec, audio_prompts.to(dtype=wav2vec.dtype),
133 fps.item(), num_frames=batch["audio_len"])
134 # audio_prompts = torch.cat(audio_prompts, dim=0).to(device=self.device, dtype=weight_dtype)
135 audio_prompts = audio_prompts.to(device=self.device, dtype=weight_dtype)
136 if audio_prompts.shape[1] <= 129:
137 audio_prompts = torch.cat([audio_prompts, torch.zeros_like(
138 audio_prompts[:, :1]).repeat(1, 129 - audio_prompts.shape[1], 1, 1, 1)], dim=1)
139 else:
140 audio_prompts = torch.cat([audio_prompts, torch.zeros_like(
141 audio_prompts[:, :1]).repeat(1, 5, 1, 1, 1)], dim=1)
143 wav2vec.to("cpu")
144 torch.cuda.empty_cache()
146 uncond_audio_prompts = torch.zeros_like(audio_prompts[:, :129])
147 motion_exp = batch["motion_bucket_id_exps"].to(self.device)
148 motion_pose = batch["motion_bucket_id_heads"].to(self.device)
150 pixel_value_ref = batch['pixel_value_ref'].to(self.device) # (b f c h w) range from [0,255]
151 # convert pixel_value_ref from (b c h w) to (b f c h w) with f = 1
152 pixel_value_ref = rearrange(pixel_value_ref, "b c h w -> b 1 c h w") # adding the f dimension
153 face_masks = get_facemask(pixel_value_ref.clone(), align_instance, area=3.0)
155 pixel_value_ref = pixel_value_ref.clone().repeat(1, 129, 1, 1, 1)
156 uncond_pixel_value_ref = torch.zeros_like(pixel_value_ref)
157 pixel_value_ref = pixel_value_ref / 127.5 - 1.
158 uncond_pixel_value_ref = uncond_pixel_value_ref * 2 - 1
160 pixel_value_ref_for_vae = rearrange(pixel_value_ref, "b f c h w -> b c f h w")
161 uncond_uncond_pixel_value_ref = rearrange(uncond_pixel_value_ref, "b f c h w -> b c f h w")
163 pixel_value_llava = batch["pixel_value_ref_llava"].to(self.device)
164 pixel_value_llava = rearrange(pixel_value_llava, "b c h w -> b 1 c h w") # adding the f dimension
165 pixel_value_llava = rearrange(pixel_value_llava, "b f c h w -> (b f) c h w")
166 uncond_pixel_value_llava = pixel_value_llava.clone()
168 # ========== Encode reference latents ==========
169 vae_dtype = self.vae.dtype
170 with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_dtype != torch.float32):
172 if args.cpu_offload:
173 self.vae.to('cuda')
175 self.vae.enable_tiling()
176 ref_latents = self.vae.encode(pixel_value_ref_for_vae.clone()).latent_dist.sample()
177 uncond_ref_latents = self.vae.encode(uncond_uncond_pixel_value_ref).latent_dist.sample()
178 self.vae.disable_tiling()
179 if hasattr(self.vae.config, 'shift_factor') and self.vae.config.shift_factor:
180 ref_latents.sub_(self.vae.config.shift_factor).mul_(self.vae.config.scaling_factor)
181 uncond_ref_latents.sub_(self.vae.config.shift_factor).mul_(self.vae.config.scaling_factor)
182 else:
183 ref_latents.mul_(self.vae.config.scaling_factor)
184 uncond_ref_latents.mul_(self.vae.config.scaling_factor)
186 if args.cpu_offload:
187 self.vae.to('cpu')
188 torch.cuda.empty_cache()
190 face_masks = torch.nn.functional.interpolate(face_masks.float().squeeze(2),
191 (ref_latents.shape[-2], ref_latents.shape[-1]),
192 mode="bilinear").unsqueeze(2).to(dtype=ref_latents.dtype)
194 size = (batch['pixel_value_ref'].shape[-2], batch['pixel_value_ref'].shape[-1])
195 target_length = 129
196 target_height = align_to(size[0], 16)
197 target_width = align_to(size[1], 16)
198 concat_dict = {'mode': 'timecat', 'bias': -1}
199 freqs_cos, freqs_sin = self.get_rotary_pos_embed(
200 target_length,
201 target_height,
202 target_width,
203 concat_dict)
204 n_tokens = freqs_cos.shape[0]
206 generator = torch.Generator(device=self.device).manual_seed(args.seed)
208 debug_str = f"""
209 prompt: {prompt}
210 audio_path: {audio_path}
211 negative_prompt: {neg_prompt}
212 fps: {fps}
213 infer_steps: {args.infer_steps}
214 target_height: {target_height}
215 target_width: {target_width}
216 target_length: {target_length}
217 guidance_scale: {args.cfg_scale}
218 """
219 self.logger.info(debug_str)
220 pipeline_kwargs = {
221 "cpu_offload": args.cpu_offload
222 }
223 start_time = time.time()
224 samples = self.pipeline(prompt=prompt,
225 height=target_height,
226 width=target_width,
227 frame=target_length,
228 num_inference_steps=args.infer_steps,
229 guidance_scale=args.cfg_scale,
230 negative_prompt=neg_prompt,
231 num_images_per_prompt=args.num_images,
232 generator=generator,
233 prompt_embeds=None,
234 ref_latents=ref_latents, # [1, 16, 1, h//8, w//8]
235 uncond_ref_latents=uncond_ref_latents,
236 pixel_value_llava=pixel_value_llava, # [1, 3, 336, 336]
237 uncond_pixel_value_llava=uncond_pixel_value_llava,
238 face_masks=face_masks, # [b f h w]
239 audio_prompts=audio_prompts,
240 uncond_audio_prompts=uncond_audio_prompts,
241 motion_exp=motion_exp,
242 motion_pose=motion_pose,
243 fps=fps,
244 num_videos_per_prompt=1,
245 attention_mask=None,
246 negative_prompt_embeds=None,
247 negative_attention_mask=None,
248 output_type="pil",
249 freqs_cis=(freqs_cos, freqs_sin),
250 n_tokens=n_tokens,
251 data_type='video',
252 is_progress_bar=True,
253 vae_ver=self.args.vae,
254 enable_tiling=self.args.vae_tiling,
255 **pipeline_kwargs
256 )[0]
257 if samples is None:
258 return None
259 out_dict['samples'] = samples
260 gen_time = time.time() - start_time
261 logger.info(f"Success, time: {gen_time}")
263 wav2vec.to(self.device)
265 return out_dict