Coverage for wrapper/wan/wrapper_wan.py: 88%

1""" 

2Wrapper for Wan video generation model. 

3""" 

4import os 

5import logging 

6import tempfile 

7import aiofiles 

8 

9from typing import Optional 

10from typing import Dict 

11from typing import Any 

12from typing import Union 

13from typing import Tuple 

14from typing import List 

15 

16from abc import abstractmethod 

17 

18from PIL import Image 

19 

20import torch 

21 

22from torch import inference_mode 

23 

24import torchvision 

25 

26from functools import partial 

27 

28from model_timing import GenTimer 

29from wrapper_usp import USPGeneration 

30from image_utils import base64_to_img 

31 

32from wan.modules.t5 import T5EncoderModel 

33from wan.modules.vae import WanVAE 

34from wan.distributed.fsdp import shard_model 

35 

36from xfuser.config import EngineConfig 

37 

38 

39class WanVideoGeneration(USPGeneration): 

40 """Generic class handle video generation using the Wan model.""" 

41 

42 FPS = 16 

43 NUM_HEADS = 40 

44 interrupted: bool 

45 

46 def __init__( 

47 self, 

48 model_name: str = "wan", 

49 ckpt_dir: str = "./Wan2.1-I2V-14B-480P", 

50 engine_config: EngineConfig = None, 

51 param_dtype: torch.dtype = torch.bfloat16, 

52 ) -> None: 

53 super().__init__( 

54 model_name=model_name, 

55 engine_config=engine_config, 

56 param_dtype=param_dtype, 

57 ) 

58 

59 self.ckpt_dir = ckpt_dir 

60 

61 # Model components 

62 self.text_encoder: Optional[T5EncoderModel] = None 

63 self.vae: Optional[WanVAE] = None 

64 

65 # Model features 

66 self.sp_size = 1 

67 # https://replicate.com/blog/wan-21-parameter-sweep 

68 self.shift = 3.0 # Low values -> Less movement and smoother 

69 self.guide_scale = 5.0 # Higher values -> Follow the prompt closer (maybe also the image???) 

70 

71 # https://github.com/Wan-Video/Wan2.1/blob/main/wan/configs/wan_i2v_14B.py 

72 self.patch_size: Tuple[int, int, int] = (1, 2, 2) 

73 self.vae_stride = (4, 8, 8) # time, height, width 

74 

75 def __del__(self) -> None: 

76 """Clean models.""" 

77 if self.text_encoder is not None: 

78 del self.text_encoder 

79 if self.vae is not None: 

80 del self.vae 

81 super().__del__() 

82 

83 def init_vae(self) -> None: 

84 """Initialize the VAE model.""" 

85 assert torch.cuda.is_available() 

86 

87 prev_memory = torch.cuda.memory_allocated() 

88 self.load_timer.start("vae") 

89 self.vae = WanVAE( 

90 vae_pth=os.path.join(self.ckpt_dir, "Wan2.1_VAE.pth"), 

91 # dtype=torch.float, 

92 device=self.device, # device_img_encoder, 

93 ) 

94 self.load_timer.end("vae") 

95 diff_memory = torch.cuda.memory_allocated() - prev_memory 

96 logging.info(f"[{self.rank}] VAE memory allocated: {diff_memory / 1024 / 1024 ** 2:.2f} GB.") 

97 

98 def load_model(self) -> None: 

99 """Load the model into memory.""" 

100 assert torch.cuda.is_available() 

101 

102 # Load across GPUs 

103 shard_fn = None 

104 if self.world_size > 1: 

105 shard_fn = partial(shard_model, device_id=self.device_id) 

106 

107 prev_memory = torch.cuda.memory_allocated() 

108 self.load_timer.start("text_encoder") 

109 self.text_encoder = T5EncoderModel( 

110 text_len=512, 

111 dtype=torch.bfloat16, 

112 device=torch.device('cpu'), # device_txt_encoder, 

113 checkpoint_path=os.path.join(self.ckpt_dir, 'models_t5_umt5-xxl-enc-bf16.pth'), 

114 tokenizer_path=os.path.join(self.ckpt_dir, 'google/umt5-xxl'), 

115 shard_fn=shard_fn, # Load across GPUs 

116 ) 

117 self.load_timer.end("text_encoder") 

118 diff_memory = torch.cuda.memory_allocated() - prev_memory 

119 logging.info(f"[{self.rank}] Text encoder memory allocated: {diff_memory / 1024 / 1024 ** 2:.2f} GB.") 

120 

121 self.init_vae() 

122 

123 def _assert_model_init(self) -> None: 

124 super()._assert_model_init() 

125 if self.text_encoder is None: 

126 raise ValueError("Text encoder not initialized.") 

127 if self.vae is None: 

128 raise ValueError("VAE not initialized.") 

129 

130 def _assert_args( 

131 self, 

132 height: int, 

133 width: int, 

134 num_frames: int, 

135 start_frames: int = 1, 

136 ) -> None: 

137 if not self.vae_stride: 

138 raise ValueError("VAE stride not set.") 

139 if height % self.vae_stride[1] != 0: 

140 raise ValueError(f"Height {height} should be divisible by VAE factor {self.vae_stride[1]}") 

141 if width % self.vae_stride[2] != 0: 

142 raise ValueError(f"Width {width} should be divisible by VAE factor {self.vae_stride[2]}") 

143 # TODO check world size divisible 

144 

145 # Latent space is the first frame (input image) + empty frames//4 

146 if (num_frames - start_frames) % self.vae_stride[0] != 0: 

147 raise ValueError(f"num_frames {num_frames} should be {self.vae_stride[0]}*n") 

148 

149 # The number of frames should be less than 1+80 

150 # Over 1+80 frames, it triggers weird video effects 

151 if num_frames < 1 or num_frames > start_frames + 80: 

152 raise ValueError(f"num_frames {num_frames} should be between 1 and 1+80") 

153 

154 def _get_mask( 

155 self, 

156 lat_h: int, 

157 lat_w: int, 

158 total_frames: int, 

159 start_frames: int, 

160 end_frames: int, 

161 ) -> torch.Tensor: 

162 """ 

163 Create a mask for the latent frames. 

164 Converts from raw frame mask [1, total_frames, H, W] to latent mask [4, latent_frames, H, W] 

165 

166 latent_frames = start_frames + (middle_frames // 4) + end_frames 

167 where middle_frames = total_frames - start_frames - end_frames 

168 

169 Start and end frames are repeated 4x; middle frames are grouped into latent frames (1 per 4 frames). 

170 """ 

171 assert self.vae_stride is not None 

172 assert (total_frames - start_frames 

173 - end_frames) % self.vae_stride[0] == 0, "Middle frames must be divisible by {f}" 

174 

175 # Step 1: Create base mask 

176 mask = torch.ones(1, total_frames, lat_h, lat_w, device=self.device) 

177 mask[:, start_frames:total_frames - end_frames] = 0 # zero out middle frames 

178 

179 # Step 2: Expand start and end frames (repeated 4x each) 

180 start_frames_repeated = torch.repeat_interleave( 

181 mask[:, 0:start_frames], 

182 repeats=4, 

183 dim=1 

184 ) 

185 

186 end_frames_repeated = torch.repeat_interleave( 

187 mask[:, total_frames - end_frames:], 

188 repeats=4, 

189 dim=1 

190 ) if end_frames > 0 else torch.zeros(1, 0, lat_h, lat_w, device=self.device) 

191 

192 # Step 3: Combine into final mask 

193 middle_mask = mask[:, start_frames:total_frames - end_frames] 

194 mask = torch.cat([ 

195 start_frames_repeated, 

196 middle_mask, 

197 end_frames_repeated 

198 ], dim=1) 

199 

200 # Step 4: Reshape to latent format: [1, latent_frames, 4, H, W] -> [4, latent_frames, H, W] 

201 mask = mask.view(1, mask.shape[1] // self.vae_stride[0], 4, lat_h, lat_w) 

202 mask = mask.transpose(1, 2)[0] # remove batch dim 

203 

204 # Step 5: Sanity check 

205 num_lat_frames = mask.shape[1] 

206 expected_lat_frames = start_frames + (total_frames - start_frames 

207 - end_frames) // self.vae_stride[0] + end_frames 

208 assert num_lat_frames == expected_lat_frames, f"Latent frames {num_lat_frames} != {expected_lat_frames}" 

209 

210 return mask 

211 

212 @inference_mode() 

213 async def warmup(self) -> None: 

214 logging.info(f"[{self.rank}] Warmup for Wan generation.") 

215 await self.generate( 

216 img=Image.new("RGB", (1280, 800), (255, 255, 255)), 

217 prompt="Warmup prompt", 

218 neg_prompt="", 

219 width=1280, 

220 height=720, 

221 num_frames=1 + 4, 

222 sampling_steps=2) 

223 

224 @inference_mode() 

225 async def generate( 

226 self, 

227 img: Image.Image, 

228 prompt: str, 

229 neg_prompt: str = "", 

230 width: int = 640, 

231 height: int = 480, 

232 num_frames: int = 1 + 80, 

233 sampling_steps: int = 50, 

234 job_id: Optional[str] = None, 

235 output_type: str = "tensor" 

236 ) -> Union[List[Image.Image], str, bytes, torch.Tensor]: 

237 raise NotImplementedError("Method should be implemented in subclasses.") 

238 

239 @inference_mode() 

240 def vae_decode( 

241 self, 

242 latents: torch.Tensor, 

243 job_id: Optional[str] = None, 

244 ) -> torch.Tensor: 

245 """Latent -> Pixels.""" 

246 gen_timer = self._new_gen_timer(job_id) 

247 

248 assert self.vae is not None 

249 # Assert arguments 

250 assert latents is not None 

251 assert isinstance(latents, torch.Tensor) 

252 assert latents.dim() == 4 # C, T, H, W 

253 assert latents.shape[0] == 20 

254 

255 try: 

256 gen_timer.start("vae_decoder") 

257 latents = latents.to(self.device, dtype=self.param_dtype) 

258 pixels = self.vae.decode([latents])[0] 

259 gen_timer.end("vae_decoder") 

260 return pixels 

261 finally: 

262 gen_timer.end("total") 

263 

264 @inference_mode() 

265 def vae_encode( 

266 self, 

267 pixels: torch.Tensor, 

268 job_id: Optional[str] = None, 

269 ) -> torch.Tensor: 

270 """Pixels -> Latent.""" 

271 gen_timer = self._new_gen_timer(job_id) 

272 

273 assert self.vae is not None 

274 # Assert arguments 

275 assert pixels is not None 

276 assert isinstance(pixels, torch.Tensor) 

277 assert pixels.dim() == 4 # C, T, H, W 

278 assert pixels.shape[1] == 3 # RGB 

279 

280 try: 

281 gen_timer.start("vae_encoder") 

282 pixels = pixels.to(self.device, dtype=self.param_dtype) 

283 latents = self.vae.encode([pixels])[0] 

284 gen_timer.end("vae_encoder") 

285 return latents 

286 finally: 

287 gen_timer.end("total") 

288 

289 async def _output_video( 

290 self, 

291 job_id: Optional[str], 

292 gen_timer: GenTimer, 

293 video_tensor: torch.Tensor, # C, T, H, W 

294 output_type: str = "tensor", # "tensor", "pil", "video_binary", "video_path" 

295 ) -> Union[List[Image.Image], str, bytes, torch.Tensor, None]: 

296 gen_timer.start("output") 

297 try: 

298 if output_type == "tensor": 

299 return video_tensor 

300 

301 if output_type == "pil": 

302 return self._tensor_to_pil(video_tensor) 

303 

304 if output_type in ("video_binary", "video_path"): 

305 if not job_id: 

306 video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name 

307 else: 

308 video_path = f"/tmp/{job_id}.mp4" 

309 self._save_video( 

310 video_tensor=video_tensor, 

311 video_path=video_path) 

312 if output_type == "video_path": 

313 return video_path 

314 

315 # video_binary 

316 async with aiofiles.open(video_path, "rb") as file: 

317 video_binary = await file.read() 

318 return video_binary 

319 

320 logging.error(f"Unknown output type: {output_type}") 

321 return None 

322 finally: 

323 gen_timer.end("output") 

324 

325 @abstractmethod 

326 def _save_video( 

327 self, 

328 video_tensor: torch.Tensor, # C, T, H, W (not B, C, T, H, W) 

329 video_path: str 

330 ) -> Optional[str]: 

331 raise NotImplementedError("Method should be implemented in subclasses.") 

332 

333 def _tensor_to_pil( 

334 self, 

335 tensor: torch.Tensor, # C, T, H, W 

336 nrow: int = 8, 

337 normalize: bool = True, 

338 value_range: Union[tuple, List] = (-1, 1), 

339 ) -> List[Image.Image]: 

340 assert tensor is not None 

341 assert isinstance(tensor, torch.Tensor) 

342 assert tensor.dim() == 4 # C, T, H, W 

343 

344 tensor = tensor.clamp(min(value_range), max(value_range)) 

345 tensor = torch.stack([ 

346 torchvision.utils.make_grid(u, nrow=nrow, normalize=normalize, value_range=value_range) 

347 for u in tensor.unbind(2) 

348 ], dim=1).permute(1, 2, 3, 0) 

349 tensor = (tensor * 255).type(torch.uint8).cpu() 

350 

351 return [Image.fromarray(frame) for frame in tensor.numpy()] 

352 

353 async def get_rest_args( 

354 self, 

355 data_json: Dict[str, Union[str, int, float]] 

356 ) -> Dict[str, Any]: 

357 if data_json is None: 

358 raise ValueError("Missing JSON body") 

359 

360 job_id = data_json.get("job_id", None) 

361 

362 img_base64 = data_json.get("img", None) 

363 if not img_base64: 

364 raise ValueError("Missing 'img' parameter") 

365 if not isinstance(img_base64, str): 

366 raise ValueError("'img' parameter must be a base64-encoded string") 

367 img = base64_to_img(img_base64) 

368 

369 prompt = data_json.get("prompt", None) 

370 if prompt is None: 

371 raise ValueError("Missing 'prompt' parameter") 

372 neg_prompt = data_json.get("neg_prompt", "") 

373 

374 width = int(data_json.get("width", 640)) 

375 height = int(data_json.get("height", 480)) 

376 sampling_steps = int(data_json.get("sampling_steps", 5)) or int(data_json.get("steps", 5)) 

377 output_type = data_json.get("output_type", "tensor") 

378 num_frames = int(data_json.get("num_frames", 1 + 16)) 

379 

380 if height <= 0: 

381 raise ValueError(f"height {height} must be positive.") 

382 if width <= 0: 

383 raise ValueError(f"width {width} must be positive.") 

384 if sampling_steps <= 0: 

385 raise ValueError(f"sampling_steps {sampling_steps} must be positive.") 

386 

387 video_seconds = data_json.get("video_seconds", 0.0) 

388 if video_seconds: 

389 if float(video_seconds) <= 0: 

390 raise ValueError(f"video_seconds {video_seconds} must be positive.") 

391 if not self.vae_stride: 

392 raise ValueError("VAE stride not set.") 

393 vae_frames = self.vae_stride[0] 

394 num_frames = int(video_seconds * self.FPS) 

395 num_frames = 1 + ((num_frames - 1) // vae_frames) * vae_frames # 4n + 1 

396 

397 if num_frames <= 0: 

398 raise ValueError(f"num_frames {num_frames} must be positive.") 

399 

400 return { 

401 "task": self.model_name, 

402 "args": { 

403 "job_id": job_id, 

404 "img": img, 

405 "prompt": prompt, 

406 "neg_prompt": neg_prompt, 

407 "width": width, 

408 "height": height, 

409 "num_frames": num_frames, 

410 "sampling_steps": sampling_steps, 

411 "output_type": output_type 

412 } 

413 }