Coverage for wrapper/hunyuanavatar/wrapper_hunyuanavatar.py: 84%
214 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Wrapper for Hunyuan-Avatar video generation.
3"""
4import os
5import einops
6import librosa
7import logging
8import datetime
9import tempfile
10import aiofiles
11import numpy as np
13from typing_extensions import override # for Python 3.11 compatibility
14from typing import Union
15from typing import List
16from typing import Optional
17from typing import Dict
18from typing import Any
20import torch
21import torch.distributed as dist
22from torch import inference_mode
24from hymm_sp.config import parse_args
25from hymm_sp.data_kits.face_align import AlignImage
26from hymm_sp.modules.parallel_states import nccl_info
27from hymm_sp.modules.parallel_states import initialize_sequence_parallel_state
28from sample_inference_audio import HunyuanVideoSampler
29from encode_data import VideoAudioTextLoaderVal
31from transformers import WhisperModel
32from transformers import AutoFeatureExtractor
34from PIL import Image
36from xfuser.config import EngineConfig
38from model_timing import GenTimer
40from wrapper_usp import USPGeneration
42from image_utils import base64_to_img
43from media_utils import base64_to_audio_file
44from media_utils import empty_audio_file
45from media_utils import save_video_audio
48class HunyuanAvatarGeneration(USPGeneration):
49 """
50 Hunyuan-Avatar video generation wrapper.
51 """
53 def __init__(
54 self,
55 engine_config: EngineConfig = None,
56 param_dtype: torch.dtype = torch.bfloat16,
57 ) -> None:
58 super().__init__(
59 "hunyuanavatar",
60 engine_config,
61 param_dtype)
63 # Model components
64 self.hunyuan_video_sampler: Optional[HunyuanVideoSampler] = None
65 self.wav2vec: Optional[WhisperModel] = None
66 self.align_instance: Optional[AlignImage] = None
67 self.feature_extractor: Optional[AutoFeatureExtractor] = None
68 self.text_encoder: Optional[torch.nn.Module] = None
69 self.text_encoder_2: Optional[torch.nn.Module] = None
71 # Model features
72 self.image_size = 704
73 self.vae_stride = (4, 8, 8) # time, height, width
74 self.MAX_FRAMES = 1 + 80
75 self.FPS = 12.5 # only support either 12.5 or 25 fps
77 # Load models
78 self.models_root_path = '/hunyuanavatar/weights/ckpts'
79 os.environ['MODEL_BASE'] = '/hunyuanavatar/weights'
81 self.args = parse_args()
82 self.args.ckpt = f"{self.models_root_path}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt"
83 self.args.prompt_template_video = None
85 def __del__(self) -> None:
86 # Ensure the model is properly cleaned up
87 if self.hunyuan_video_sampler is not None:
88 self.hunyuan_video_sampler = None
89 if self.wav2vec is not None:
90 self.wav2vec = None
91 if self.align_instance is not None:
92 self.align_instance = None
93 if self.feature_extractor is not None:
94 self.feature_extractor = None
95 if self.text_encoder is not None:
96 self.text_encoder = None
97 if self.text_encoder_2 is not None:
98 self.text_encoder_2 = None
100 super().__del__()
102 def init_parallelism(self) -> None:
103 self.load_timer.start("torch_dist")
105 logging.info("Initializing distributed environment...")
107 if "MASTER_ADDR" not in os.environ:
108 logging.info("MASTER_ADDR not set, skipping distributed initialization.")
109 return
111 self.rank = int(os.getenv("RANK", 0))
112 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
113 self.world_size = int(os.getenv("WORLD_SIZE", 1))
115 self.device_id = self.local_rank
116 self.device = torch.device(f"cuda:{self.device_id}")
118 torch.cuda.set_device(self.local_rank)
120 if self.world_size <= 1:
121 logging.info("World size is 1, skipping distributed initialization.")
122 self.load_timer.end("torch_dist")
123 return
125 if not dist.is_initialized():
126 logging.info(
127 f"Initializing distributed environment with local_rank: {self.rank}, world_size: {self.world_size}.")
128 dist.init_process_group(
129 backend="nccl",
130 init_method="env://",
131 timeout=datetime.timedelta(seconds=2**31 - 1),
132 world_size=self.world_size,
133 rank=self.rank)
134 else:
135 logging.info("Distributed environment already initialized.")
137 torch.manual_seed(self.base_seed)
138 torch.cuda.manual_seed_all(self.base_seed)
140 logging.info(
141 f"Distributed environment initialized with rank: {dist.get_rank()}, world_size: {dist.get_world_size()}.")
143 initialize_sequence_parallel_state(self.world_size)
145 self.load_timer.end("torch_dist")
147 def load_model(self) -> None:
148 self.rank = 0
149 self.vae_dtype = torch.float16
150 self.device = torch.device("cuda")
151 if nccl_info.sp_size > 1:
152 self.device = torch.device(f"cuda:{dist.get_rank()}")
153 self.rank = dist.get_rank()
155 self.load_timer.start("hunyuan_video_sampler")
156 self.hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(
157 f"{self.models_root_path}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
158 args=self.args,
159 device=self.device)
160 assert self.hunyuan_video_sampler is not None
161 self.load_timer.end("hunyuan_video_sampler")
163 # Get the updated args
164 self.args = self.hunyuan_video_sampler.args
166 # Load the wav2vec model
167 self.load_timer.start("wav2vec")
168 self.wav2vec = WhisperModel.from_pretrained(
169 f"{self.models_root_path}/whisper-tiny/" # nosec B615 - local path
170 )
171 assert self.wav2vec is not None
172 self.wav2vec = self.wav2vec.to(device=self.device, dtype=torch.float32) # type: ignore[call-arg]
173 self.wav2vec.requires_grad_(False)
174 self.load_timer.end("wav2vec")
176 # Load the align instance
177 self.load_timer.start("align_instance")
178 det_path = f"{self.models_root_path}/det_align/detface.pt"
179 self.align_instance = AlignImage("cuda", det_path=det_path)
180 self.load_timer.end("align_instance")
182 # Load the feature extractor
183 self.load_timer.start("feature_extractor")
184 self.feature_extractor = AutoFeatureExtractor.from_pretrained(
185 f"{self.models_root_path}/whisper-tiny/") # nosec B615 - local path
186 self.load_timer.end("feature_extractor")
188 self.text_encoder = self.hunyuan_video_sampler.text_encoder
189 self.text_encoder_2 = self.hunyuan_video_sampler.text_encoder_2
191 self.data_loader = VideoAudioTextLoaderVal(
192 image_size=self.image_size,
193 text_encoder=self.text_encoder,
194 text_encoder_2=self.text_encoder_2,
195 feature_extractor=self.feature_extractor,
196 )
198 def _assert_model_init(self) -> None:
199 assert self.hunyuan_video_sampler is not None, "HunyuanVideoSampler is not initialized."
200 assert self.wav2vec is not None, "Wav2Vec model is not initialized."
201 assert self.align_instance is not None, "AlignImage instance is not initialized."
202 assert self.feature_extractor is not None, "Feature extractor is not initialized."
203 assert self.text_encoder is not None, "Text encoder is not initialized."
204 assert self.text_encoder_2 is not None, "Text encoder 2 is not initialized."
206 @inference_mode()
207 async def warmup(self) -> None:
208 logging.info(f"[{self.rank}] Warmup for Hunyuan Avatar generation.")
209 empty_img = Image.new("RGB", (640, 480), color=(255, 255, 255))
210 empty_audio_path = empty_audio_file(0.5) # 0.5 seconds of silence
211 await self.generate(
212 img=empty_img,
213 audio_path=empty_audio_path,
214 prompt="A warmup generation",
215 width=640,
216 height=480,
217 sampling_steps=2)
218 if os.path.exists(empty_audio_path):
219 os.remove(empty_audio_path)
221 @override
222 @inference_mode()
223 async def generate(
224 self,
225 img: Image.Image,
226 audio_path: str,
227 prompt: str,
228 width: int = 1280,
229 height: int = 720,
230 sampling_steps: int = 30, # 10
231 audio_scale: float = 1.0, # 1.0 for audio, 0.0 for no audio influence TODO not used
232 cfg_scale: float = 5.0, # prompt
233 audio_cfg_scale: float = 5.0, # TODO not used
234 job_id: Optional[str] = None,
235 output_type: str = "pil", # "pil", "video_binary", "video_path"
236 ) -> Union[List[Image.Image], np.ndarray, str, bytes, None]:
237 """
238 Generate a video from an image, a piece of audio, and a text prompt.
239 Args:
240 img (Image.Image): Input image.
241 audio_path (str): Path to the audio file.
242 prompt (str): Text prompt for video generation.
243 neg_prompt (str, optional): Negative text prompt. Defaults to "".
244 width (int, optional): Width of the output video. Defaults to 1280.
245 height (int, optional): Height of the output video. Defaults to 720.
246 sampling_steps (int, optional): Number of sampling steps. Defaults to 30.
247 """
248 gen_timer = self._new_gen_timer(job_id)
250 self.running = True # Mark running to avoid concurrent calls
252 try:
253 self._assert_model_init()
255 audio_duration = librosa.get_duration(path=audio_path) # TODO figure audio
256 num_frames = int(self.FPS * audio_duration // self.vae_stride[0]) * self.vae_stride[0] + 5
257 if num_frames > self.MAX_FRAMES:
258 raise ValueError(f"Audio {audio_duration:.2f}s exceeds {self.MAX_FRAMES} frames")
259 # num_frames = min(num_frames, self.MAX_FRAMES)
260 lat_num_frames = (num_frames - 1) // self.vae_stride[0] + 1
262 if self.rank == 0:
263 logging.info(
264 f"[{self.rank}] Audio:{audio_duration:.2f}s #frames:{num_frames} #lat_frames:{lat_num_frames}.")
266 # Prepare the data
267 gen_timer.start("encoding_inputs")
268 img_resized = img.resize((width, height), Image.Resampling.LANCZOS)
269 results = self.data_loader.encode_data(
270 ref_image=img_resized,
271 audio_path=audio_path,
272 prompt=prompt,
273 fps=self.FPS,
274 )
275 gen_timer.end("encoding_inputs")
277 self.args.cfg_scale = cfg_scale
278 self.args.infer_steps = sampling_steps
279 gen_timer.start("hunyuanavatar_generation")
280 assert self.hunyuan_video_sampler is not None
281 samples = self.hunyuan_video_sampler.predict(
282 self.args,
283 results,
284 self.wav2vec,
285 self.feature_extractor,
286 self.align_instance)
287 gen_timer.end("hunyuanavatar_generation")
289 sample = samples['samples'][0].unsqueeze(0) # de-noised latent, (bs, 16, t//4, h//8, w//8)
290 # sample = sample[:, :, :results["audio_len"][0]]
291 sample = sample[:, :, :results["audio_len"]]
293 logging.info(f"[{self.rank}] Sample shape after slicing: {sample.shape}.")
294 video = einops.rearrange(sample[0], "c f h w -> f h w c")
295 video = (video * 255.).data.cpu().numpy().astype(np.uint8) # （f h w c)
297 torch.cuda.empty_cache()
299 final_frames = []
300 for frame in video:
301 final_frames.append(frame)
302 final_frames = np.stack(final_frames, axis=0)
304 return await self._output_video(
305 job_id,
306 gen_timer,
307 audio_path,
308 final_frames,
309 output_type)
310 finally:
311 self.running = False
312 gen_timer.end("total")
314 async def _output_video(
315 self,
316 job_id: Optional[str],
317 gen_timer: GenTimer,
318 audio_path: str,
319 video_frames: np.ndarray,
320 output_type: str = "pil", # "pil", "video_binary", "video_path"
321 ) -> Union[List[Image.Image], np.ndarray, str, bytes, None]:
322 gen_timer.start("output")
323 try:
324 if output_type == "pil":
325 return video_frames
327 if output_type in ("video_binary", "video_path"):
328 if not job_id:
329 video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
330 else:
331 video_path = f"/tmp/{job_id}.mp4"
332 video_path = await save_video_audio(
333 video_content=video_frames,
334 audio_path=audio_path,
335 out_video_path=video_path,
336 fps=self.FPS)
337 if output_type == "video_path":
338 return video_path
340 # video_binary
341 async with aiofiles.open(video_path, "rb") as file:
342 video_binary = await file.read()
343 return video_binary
345 logging.error(f"Unknown output type: {output_type}")
346 return None
347 finally:
348 gen_timer.end("output")
350 async def get_rest_args(self, data_json: Dict[str, str]) -> Dict[str, Any]:
351 if data_json is None:
352 raise ValueError("Missing JSON body")
354 job_id = data_json.get("job_id", None)
356 img_base64 = data_json.get("img", None)
357 if img_base64 is None:
358 raise ValueError("Missing 'img' parameter")
359 img = base64_to_img(img_base64)
361 audio_base64 = data_json.get("audio", None)
362 if audio_base64 is None:
363 raise ValueError("Missing 'audio' parameter")
364 audio_path = None
365 if not job_id:
366 audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
367 else:
368 audio_path = f"/tmp/{job_id}.wav"
369 audio_path = await base64_to_audio_file(
370 audio_base64,
371 audio_path=audio_path)
373 prompt = data_json.get("prompt", None)
374 if prompt is None:
375 raise ValueError("Missing 'prompt' parameter")
377 height = int(data_json.get("height", 480))
378 width = int(data_json.get("width", 640))
379 steps = int(data_json.get("sampling_steps", 10))
380 audio_scale = float(data_json.get("audio_scale", 1.0))
381 cfg_scale = float(data_json.get("cfg_scale", 5.0))
382 audio_cfg_scale = float(data_json.get("audio_cfg_scale", 5.0))
383 return {
384 "task": self.model_name,
385 "args": {
386 "job_id": job_id,
387 "img": img,
388 "prompt": prompt,
389 "audio_path": audio_path,
390 "width": width,
391 "height": height,
392 "sampling_steps": steps,
393 "audio_scale": audio_scale,
394 "cfg_scale": cfg_scale,
395 "audio_cfg_scale": audio_cfg_scale,
396 }
397 }