Coverage for wrapper/hunyuanframepack/wrapper_hunyuanframepack.py: 56%
108 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Wrapper class for HunyuanFramePack model generation using Hugging Face Diffusers and Xfuser.
3"""
4import sys
5import logging
6import time
7import random
8import math
9import torch
10import asyncio
12from torch import inference_mode
14from typing import Optional
16from wrapper_hunyuanframepack_base import HunyuanFramePackBase
18from PIL import Image
20from diffusers_helper.hunyuan import vae_decode
21from diffusers_helper.hunyuan import vae_encode
22from diffusers_helper.utils import soft_append_bcthw
24from xfuser.config import EngineConfig
27class HunyuanFramepackGeneration(HunyuanFramePackBase):
28 """Handle video generation using the HunyuanFramePack model."""
30 def __init__(
31 self,
32 engine_config: EngineConfig = None,
33 param_dtype: torch.dtype = torch.bfloat16,
34 enable_tiling: bool = False,
35 enable_slicing: bool = False,
36 ) -> None:
37 super().__init__(
38 model_name="hunyuanframepack",
39 framepack_model_name="lllyasviel/FramePackI2V_HY",
40 engine_config=engine_config,
41 param_dtype=param_dtype,
42 enable_tiling=enable_tiling,
43 enable_slicing=enable_slicing,
44 )
46 @inference_mode()
47 async def generate(
48 self,
49 img: Image.Image,
50 prompt: str,
51 neg_prompt: str = "",
52 height: int = 512,
53 width: int = 768,
54 num_frames: int = 1 + 80,
55 sampling_steps: int = 25, # 10
56 # latent frames for every Hunyuan Video inference window
57 # 9->36 pixel frames -> 1.2 seconds
58 latent_window_size: int = 9,
59 cfg: float = 1.0,
60 distilled_guidance_scale: float = 10.0,
61 guidance_rescale: int = 0,
62 save_intermediate: Optional[str] = None,
63 job_id: Optional[str] = None,
64 output_type: str = "tensor", # "tensor", "video_binary", "video_path"
65 ) -> Optional[torch.Tensor]:
66 """
67 Generate a video from an input image and a prompt.
68 Based on:
69 https://raw.githubusercontent.com/lllyasviel/FramePack/refs/heads/main/demo_gradio.py
70 TODO move more to the base class so we can unify regular and F1.
71 """
72 gen_timer = self._new_gen_timer(job_id)
74 self._assert_model_init()
75 self._assert_args(height, width)
77 video_seconds = num_frames / self.FPS
78 total_latent_sections = (video_seconds * self.FPS) / (latent_window_size * self.vae_stride[0])
79 total_latent_sections = int(math.ceil(total_latent_sections))
80 total_latent_sections = max(1, total_latent_sections)
82 num_frames_it = latent_window_size * self.vae_stride[0] - 3
84 if self.rank == 0:
85 logging.info(
86 f"[{self.rank}] Length:{video_seconds:.2f}secs Frames:{num_frames} "
87 f"Frames/iteration:{num_frames_it} Iterations:{total_latent_sections} "
88 f"FPS:{self.FPS} CFG:{cfg}.")
90 self.running = True # Mark running to avoid concurrent calls
92 try:
93 # Text encoder
94 llama_vec, llama_attention_mask, clip_l_pooler, \
95 llama_vec_n, llama_attention_mask_n, clip_l_pooler_n = self._encode_text(
96 gen_timer,
97 prompt,
98 neg_prompt,
99 cfg)
101 # Process input image
102 input_image_np, input_image_pt = self._process_image(
103 gen_timer,
104 img,
105 height,
106 width)
108 # CLIP Vision
109 image_encoder_last_hidden_state = self._clip_vision(
110 gen_timer,
111 input_image_np)
113 # VAE encoding
114 # 1,RGB,1,h,w -> 1,lat_channels,1,lat_h,lat_w ([1,3,1,544,704] -> [1,16,1,68,88])
115 gen_timer.start("vae_encoder")
116 t0_vae = time.time()
117 start_latent = vae_encode(input_image_pt, self.vae)
118 gen_timer.end("vae_encoder")
119 if self.rank == 0:
120 logging.info(f"[{self.rank}] VAE encoding time: {time.time() - t0_vae:.3f} seconds.")
122 # Prepare latent space
123 seed = self.base_seed if self.base_seed >= 0 else random.randint(0, sys.maxsize)
124 seed_g = torch.Generator(device=self.device)
125 seed_g.manual_seed(seed)
127 lat_h = height // self.vae_stride[1]
128 lat_w = width // self.vae_stride[2]
129 history_latents = torch.zeros(
130 # B, C, T, H, W
131 size=(1, self.LAT_CHANNELS, 1 + 2 + 16, lat_h, lat_w),
132 dtype=torch.float32
133 ).cpu()
134 history_pixels = None
135 total_generated_latent_frames = 0
137 latent_paddings: list[int] = list(reversed(range(total_latent_sections)))
139 if total_latent_sections > 4:
140 # In theory the latent_paddings should follow the above sequence, but it seems that duplicating some
141 # items looks better than expanding it when total_latent_sections > 4
142 # One can try to remove below trick and just
143 # use `latent_paddings = list(reversed(range(total_latent_sections)))` to compare
144 latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0]
146 # DiT sampling
147 # We do 2 loops:
148 # 1. Outer loop sticks together chunks into a final of num_frames
149 # 2. Inner loop generate chunks of N frames (9 latent frames = 36 video frames = 1.2 seconds)
150 for it, latent_padding in enumerate(latent_paddings):
151 logging.debug(f"Running step {it + 1}.")
153 self.check_interrupted()
155 gen_timer.start(f"dit_{it:03d}")
157 is_last_section = latent_padding == 0
158 latent_padding_size = latent_padding * latent_window_size
160 indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0)
161 (clean_latent_indices_pre, _, # blank_indices
162 latent_indices,
163 clean_latent_indices_post, clean_latent_2x_indices,
164 clean_latent_4x_indices) = indices.split(
165 [1, latent_padding_size, latent_window_size, 1, 2, 16], dim=1)
166 clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1)
168 clean_latents_pre = start_latent.to(history_latents)
169 clean_latents_post, clean_latents_2x, clean_latents_4x = \
170 history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2)
171 clean_latents = torch.cat([
172 clean_latents_pre,
173 clean_latents_post
174 ], dim=2)
176 if self.engine_config is not None and self.engine_config.runtime_config.use_teacache:
177 self.transformer.initialize_teacache(enable_teacache=True, num_steps=sampling_steps)
178 else:
179 self.transformer.initialize_teacache(enable_teacache=False)
181 # [B, C, T, H, W]
182 generated_latents = await asyncio.to_thread(
183 self._sample_hunyuan,
184 it0=it,
185 gen_timer=gen_timer,
186 width=width,
187 height=height,
188 frames=num_frames_it,
189 real_guidance_scale=cfg,
190 distilled_guidance_scale=distilled_guidance_scale,
191 guidance_rescale=guidance_rescale,
192 num_inference_steps=sampling_steps,
193 generator=seed_g,
194 prompt_embeds=llama_vec,
195 prompt_embeds_mask=llama_attention_mask,
196 prompt_poolers=clip_l_pooler,
197 negative_prompt_embeds=llama_vec_n,
198 negative_prompt_embeds_mask=llama_attention_mask_n,
199 negative_prompt_poolers=clip_l_pooler_n,
200 image_embeddings=image_encoder_last_hidden_state,
201 latent_indices=latent_indices,
202 clean_latents=clean_latents,
203 clean_latent_indices=clean_latent_indices,
204 clean_latents_2x=clean_latents_2x,
205 clean_latent_2x_indices=clean_latent_2x_indices,
206 clean_latents_4x=clean_latents_4x,
207 clean_latent_4x_indices=clean_latent_4x_indices,
208 )
210 if is_last_section:
211 generated_latents = torch.cat([
212 start_latent.to(generated_latents),
213 generated_latents
214 ], dim=2)
216 total_generated_latent_frames += int(generated_latents.shape[2])
217 history_latents = torch.cat([
218 generated_latents.to(history_latents),
219 history_latents
220 ], dim=2)
222 # [B, C, T, H, W]
223 real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :]
225 gen_timer.end(f"dit_{it:03d}")
227 if self.rank == 0:
228 lat_frames = real_history_latents.shape[2]
229 # Approximation, we should add and remove appropriately
230 cur_frames = lat_frames * self.vae_stride[0]
231 logging.info(
232 f"[{self.rank}] it:{it} lat_frames:{lat_frames} frames:{cur_frames} "
233 f"{cur_frames / self.FPS:.1f}/{video_seconds:.3f} seconds.")
234 if save_intermediate is not None:
235 torch.save(
236 real_history_latents, # We are saving the whole history latents
237 f"/tmp/{save_intermediate}_latents_{it:03d}.pt")
239 if self.rank != 0:
240 return None # other workers do not need to return anything or VAE decode
242 if output_type == "latent":
243 if history_pixels is None:
244 return real_history_latents
245 section_latent_frames = latent_window_size * 2
246 return real_history_latents[:, :, :section_latent_frames]
248 # VAE decode
249 # Can be done in the background yielding frames
250 # [1, 16, lat_frames, lat_h, lat_w] -> [1, 3, 1+2+16, h, w]
251 # ([1, 16, 37, 68, 88] -> [1, 3, 145, 544, 704])
252 gen_timer.start("vae_decoder")
253 if history_pixels is None:
254 # This is the common path
255 history_pixels = vae_decode(real_history_latents, self.vae).cpu()
256 else:
257 section_latent_frames = latent_window_size * 2
258 if is_last_section:
259 section_latent_frames = latent_window_size * 2 + 1
260 # vae.config.temporal_compression_ratio=4
261 overlapped_frames = latent_window_size * self.vae_stride[0] - 3
262 current_pixels = vae_decode(real_history_latents[:, :, :section_latent_frames], self.vae).cpu()
263 history_pixels = soft_append_bcthw(current_pixels, history_pixels, overlapped_frames)
264 gen_timer.end("vae_decoder")
266 if self.rank == 0:
267 logging.info(
268 f"[{self.rank}] VAE decode. "
269 f"Latent:{real_history_latents.shape}->Pixel:{history_pixels.shape}.")
271 out_num_frames = history_pixels.shape[2]
272 if out_num_frames < num_frames:
273 logging.warning(f"[{self.rank}] Output frames {out_num_frames} < requested {num_frames}.")
274 elif out_num_frames > num_frames:
275 logging.warning(f"[{self.rank}] Output frames {out_num_frames} > requested {num_frames}. Trimming.")
276 history_pixels = history_pixels[:, :, :num_frames, :, :]
278 return await self._output_video(
279 job_id,
280 gen_timer,
281 history_pixels,
282 output_type)
283 finally:
284 self.running = False
285 gen_timer.end("total")