Coverage for wrapper/hunyuanframepackf1/wrapper_hunyuanframepackf1.py: 59%

1""" 

2Hunyuan FramePack F1 wrapper class. 

3""" 

4import sys 

5import logging 

6import time 

7import random 

8import math 

9import torch 

10import asyncio 

11 

12from torch import inference_mode 

13 

14from typing import override 

15from typing import Optional 

16 

17from wrapper_hunyuanframepack_base import HunyuanFramePackBase 

18 

19from PIL import Image 

20 

21from diffusers_helper.hunyuan import vae_decode 

22from diffusers_helper.hunyuan import vae_encode 

23from diffusers_helper.utils import soft_append_bcthw 

24 

25from xfuser.config import EngineConfig 

26 

27 

28class HunyuanFramepackF1Generation(HunyuanFramePackBase): 

29 """Handle video generation using the Hunyuan Framepack F1 model.""" 

30 def __init__( 

31 self, 

32 engine_config: EngineConfig = None, 

33 param_dtype: torch.dtype = torch.bfloat16, 

34 ) -> None: 

35 super().__init__( 

36 model_name="hunyuanframepackf1", 

37 framepack_model_name="lllyasviel/FramePack_F1_I2V_HY_20250503", 

38 engine_config=engine_config, 

39 param_dtype=param_dtype, 

40 ) 

41 

42 @override 

43 @inference_mode() 

44 async def generate( 

45 self, 

46 img: Image.Image, 

47 prompt: str, 

48 neg_prompt: str = "", 

49 height: int = 512, 

50 width: int = 768, 

51 num_frames: int = 1 + 80, 

52 sampling_steps: int = 25, # 10 

53 latent_window_size: int = 9, # latent frames for every inference window: 9->36 pixel frames -> 1.2 seconds 

54 cfg: float = 1.0, # prompt guidance scale 

55 distilled_guidance_scale: float = 10.0, 

56 guidance_rescale: int = 0, 

57 save_intermediate: Optional[str] = None, 

58 job_id: Optional[str] = None, 

59 output_type: str = "tensor", # "tensor", "latent", "video_binary", "video_path" 

60 ) -> Optional[torch.Tensor]: 

61 """ 

62 Generate a video from an input image and a prompt. 

63 Based on wrapper_hunyuanframepack.py and: 

64 https://raw.githubusercontent.com/lllyasviel/FramePack/refs/heads/main/demo_gradio_f1.py 

65 TODO move more to the base class so we can unify regular and F1. 

66 """ 

67 gen_timer = self._new_gen_timer(job_id) 

68 

69 self._assert_model_init() 

70 self._assert_args(height, width) 

71 

72 video_seconds = num_frames / self.FPS 

73 total_latent_sections = (video_seconds * self.FPS) / (latent_window_size * self.vae_stride[0]) 

74 total_latent_sections = int(math.ceil(total_latent_sections)) 

75 total_latent_sections = max(1, total_latent_sections) 

76 

77 num_frames_it = latent_window_size * self.vae_stride[0] - 3 

78 

79 if self.rank == 0: 

80 logging.info( 

81 f"[{self.rank}] Length:{video_seconds:.3f} seconds " 

82 f"Frames:{num_frames} Frames/iteration:{num_frames_it} " 

83 f"Iterations:{total_latent_sections} FPS:{self.FPS} CFG:{cfg}.") 

84 

85 self.running = True # Mark running to avoid concurrent calls 

86 

87 try: 

88 # Text encoder 

89 llama_vec, llama_attention_mask, clip_l_pooler, llama_vec_n, \ 

90 llama_attention_mask_n, clip_l_pooler_n = self._encode_text( 

91 gen_timer, 

92 prompt, 

93 neg_prompt, 

94 cfg) 

95 

96 # Process input image 

97 input_image_np, input_image_pt = self._process_image( 

98 gen_timer, 

99 img, 

100 height, 

101 width 

102 ) 

103 

104 # CLIP Vision 

105 image_encoder_last_hidden_state = self._clip_vision( 

106 gen_timer, 

107 input_image_np) 

108 

109 # VAE encoding 

110 # 1,RGB,1,h,w -> 1,lat_channels,1,lat_h,lat_w ([1,3,1,544,704] -> [1,16,1,68,88]) 

111 gen_timer.start("vae_encoder") 

112 t0_vae = time.time() 

113 start_latent = vae_encode(input_image_pt, self.vae) 

114 gen_timer.end("vae_encoder") 

115 if self.rank == 0: 

116 logging.info(f"[{self.rank}] VAE encoding time: {time.time() - t0_vae:.3f} seconds.") 

117 

118 # Prepare latent space 

119 seed = self.base_seed if self.base_seed >= 0 else random.randint(0, sys.maxsize) 

120 seed_g = torch.Generator(device=self.device) 

121 seed_g.manual_seed(seed) 

122 

123 lat_h = height // self.vae_stride[1] 

124 lat_w = width // self.vae_stride[2] 

125 history_latents = torch.zeros( 

126 # B, C, T, H, W 

127 size=(1, self.LAT_CHANNELS, 1 + 2 + 16, lat_h, lat_w), 

128 dtype=torch.float32 

129 ).cpu() 

130 history_pixels = None 

131 history_latents = torch.cat([ 

132 history_latents, 

133 start_latent.to(history_latents) 

134 ], dim=2) 

135 total_generated_latent_frames = 1 

136 

137 # DiT sampling 

138 # We do 2 loops: 

139 # 1. Outer loop sticks together chunks into a final of num_frames 

140 # 2. Inner loop generate chunks of N frames (9 latent frames = 36 video frames = 1.2 seconds) 

141 for it in range(total_latent_sections): 

142 logging.debug(f"Running step {it + 1}.") 

143 

144 self.check_interrupted() 

145 

146 gen_timer.start(f"dit_{it:03d}") 

147 

148 indices = torch.arange(0, sum([1, 16, 2, 1, latent_window_size])).unsqueeze(0) 

149 clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, \ 

150 clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, latent_window_size], dim=1) 

151 clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1) 

152 

153 clean_latents_4x, clean_latents_2x, clean_latents_1x = \ 

154 history_latents[:, :, -sum([16, 2, 1]):, :, :].split([16, 2, 1], dim=2) 

155 clean_latents = torch.cat([ 

156 start_latent.to(history_latents), 

157 clean_latents_1x 

158 ], dim=2) 

159 

160 if self.engine_config is not None and self.engine_config.runtime_config.use_teacache: 

161 self.transformer.initialize_teacache(enable_teacache=True, num_steps=sampling_steps) 

162 else: 

163 self.transformer.initialize_teacache(enable_teacache=False) 

164 

165 # [B, C, T, H, W] 

166 generated_latents = await asyncio.to_thread( 

167 self._sample_hunyuan, 

168 it0=it, 

169 gen_timer=gen_timer, 

170 width=width, 

171 height=height, 

172 frames=num_frames_it, 

173 real_guidance_scale=cfg, 

174 distilled_guidance_scale=distilled_guidance_scale, 

175 guidance_rescale=guidance_rescale, 

176 num_inference_steps=sampling_steps, 

177 generator=seed_g, 

178 prompt_embeds=llama_vec, 

179 prompt_embeds_mask=llama_attention_mask, 

180 prompt_poolers=clip_l_pooler, 

181 negative_prompt_embeds=llama_vec_n, 

182 negative_prompt_embeds_mask=llama_attention_mask_n, 

183 negative_prompt_poolers=clip_l_pooler_n, 

184 image_embeddings=image_encoder_last_hidden_state, 

185 latent_indices=latent_indices, 

186 clean_latents=clean_latents, 

187 clean_latent_indices=clean_latent_indices, 

188 clean_latents_2x=clean_latents_2x, 

189 clean_latent_2x_indices=clean_latent_2x_indices, 

190 clean_latents_4x=clean_latents_4x, 

191 clean_latent_4x_indices=clean_latent_4x_indices, 

192 ) 

193 

194 total_generated_latent_frames += int(generated_latents.shape[2]) 

195 history_latents = torch.cat([ 

196 history_latents, 

197 generated_latents.to(history_latents) 

198 ], dim=2) 

199 

200 # [B, C, T, H, W] 

201 real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :] 

202 

203 gen_timer.end(f"dit_{it:03d}") 

204 

205 if self.rank == 0: 

206 lat_frames = real_history_latents.shape[2] 

207 # Approximation, we should add and remove appropriately 

208 cur_frames = lat_frames * self.vae_stride[0] 

209 logging.info( 

210 f"[{self.rank}] lat_frames:{lat_frames} frames:{cur_frames} " 

211 f"{cur_frames / self.FPS:.1f}/{video_seconds:.3f} seconds.") 

212 if save_intermediate is not None: 

213 torch.save( 

214 real_history_latents, # We are saving the whole history latents 

215 f"/tmp/{save_intermediate}_latents_{it:03d}.pt") 

216 

217 if self.rank != 0: 

218 return None # other workers do not need to return anything or VAE decode 

219 

220 if output_type == "latent": 

221 if history_pixels is None: 

222 return real_history_latents 

223 section_latent_frames = latent_window_size * 2 

224 return real_history_latents[:, :, :section_latent_frames] 

225 

226 # VAE decode 

227 # Can be done in the background yielding frames 

228 # [1, 16, lat_frames, lat_h, lat_w] -> [1, 3, 1+2+16, h, w] ([1, 16, 37, 68, 88]->[1, 3, 145, 544, 704]) 

229 gen_timer.start("vae_decoder") 

230 if history_pixels is None: 

231 # This is the common path 

232 history_pixels = await asyncio.to_thread( 

233 vae_decode, 

234 real_history_latents, 

235 self.vae 

236 ) 

237 history_pixels = history_pixels.cpu() 

238 else: 

239 section_latent_frames = latent_window_size * 2 

240 # vae.config.temporal_compression_ratio=4 

241 overlapped_frames = latent_window_size * self.vae_stride[0] - 3 

242 current_pixels = await asyncio.to_thread( 

243 vae_decode, 

244 real_history_latents[:, :, -section_latent_frames:], 

245 self.vae) 

246 current_pixels = current_pixels.cpu() 

247 history_pixels = soft_append_bcthw(history_pixels, current_pixels, overlapped_frames) 

248 gen_timer.end("vae_decoder") 

249 

250 if self.rank == 0: 

251 logging.info( 

252 f"[{self.rank}] VAE decode. Latent:{real_history_latents.shape}->Pixel:{history_pixels.shape}.") 

253 

254 out_num_frames = history_pixels.shape[2] 

255 if out_num_frames < num_frames: 

256 logging.warning(f"[{self.rank}] Output frames {out_num_frames} < requested {num_frames}.") 

257 elif out_num_frames > num_frames: 

258 logging.warning(f"[{self.rank}] Output frames {out_num_frames} > requested {num_frames}. Trimming.") 

259 history_pixels = history_pixels[:, :, :num_frames, :, :] 

260 

261 return await self._output_video( 

262 job_id, 

263 gen_timer, 

264 history_pixels, 

265 output_type) 

266 finally: 

267 self.running = False 

268 torch.cuda.empty_cache() 

269 gen_timer.end("total")