Coverage for wrapper/wan/wrapper_wan21.py: 43%

1import sys 

2import os 

3import logging 

4import math 

5import random 

6import types 

7 

8from typing import Generator 

9from typing import Union 

10from typing import List 

11from typing import Optional 

12 

13from PIL import Image 

14 

15from contextlib import contextmanager 

16 

17import torch 

18import torch.amp as amp 

19import torch.distributed as dist 

20from torch import inference_mode 

21 

22import torchvision.transforms.functional as TF 

23 

24from functools import partial 

25 

26from wrapper_wan import WanVideoGeneration 

27 

28from wan.modules.model import WanModel 

29from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler 

30from wan.distributed.fsdp import shard_model 

31from wan.modules.clip import CLIPModel 

32from wan.distributed.xdit_context_parallel import usp_dit_forward 

33from wan.distributed.xdit_context_parallel import usp_attn_forward 

34from wan.utils.utils import cache_video 

35 

36from xfuser.config import EngineConfig 

37from xfuser.core.distributed import get_sequence_parallel_world_size 

38 

39 

40class Wan21VideoGeneration(WanVideoGeneration): 

41 """Handle video generation using the Wan 2.1 model.""" 

42 

43 interrupted: bool 

44 

45 def __init__( 

46 self, 

47 model_name: str = "wan", 

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

49 engine_config: EngineConfig = None, 

50 param_dtype: torch.dtype = torch.bfloat16, 

51 ) -> None: 

52 super().__init__( 

53 model_name=model_name, 

54 engine_config=engine_config, 

55 param_dtype=param_dtype, 

56 ) 

57 

58 self.ckpt_dir = ckpt_dir 

59 

60 # Model components 

61 self.image_encoder: Optional[CLIPModel] = None 

62 self.model: Optional[WanModel] = None 

63 

64 # Model features 

65 self.sp_size = 1 

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

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

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

69 

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

71 self.patch_size = (1, 2, 2) 

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

73 

74 def __del__(self) -> None: 

75 # Clean models 

76 if self.image_encoder is not None: 

77 del self.image_encoder 

78 if self.model is not None: 

79 del self.model 

80 super().__del__() 

81 

82 def load_model(self) -> None: 

83 """Load the Wan 2.1 model and its components.""" 

84 super().load_model() 

85 

86 prev_memory = torch.cuda.memory_allocated() 

87 self.load_timer.start("image_encoder") 

88 self.image_encoder = CLIPModel( 

89 dtype=torch.float16, # torch.float32 

90 device=self.device, # device_img_encoder, 

91 checkpoint_path=os.path.join(self.ckpt_dir, 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'), 

92 tokenizer_path=os.path.join(self.ckpt_dir, 'xlm-roberta-large') 

93 ) 

94 self.load_timer.end("image_encoder") 

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

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

97 

98 prev_memory = torch.cuda.memory_allocated() 

99 self.load_timer.start("dit") 

100 self.model = WanModel.from_pretrained( 

101 self.ckpt_dir, 

102 torch_dtype=self.param_dtype, 

103 # torch_dtype=torch.uint8, # does not work 

104 # torch_dtype=torch.bfloat16, # ~31 GB 

105 # torch_dtype=torch.float32, # ~61 GB 

106 ) 

107 if not self.model: 

108 raise ValueError("Failed to load Wan model") 

109 self.model.eval() 

110 self.model.requires_grad_(False) 

111 self.load_timer.end("dit") 

112 

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

114 logging.info(f"[{self.rank}] DiT Memory allocated: {diff_memory / 1024 / 1024 ** 2:.2f} GB.") 

115 

116 def init_model_parallelism(self) -> None: 

117 """Initialize model parallelism for Wan 2.1.""" 

118 if not self.model: 

119 raise ValueError("Model not loaded") 

120 

121 if not dist.is_initialized() or self.world_size <= 1: 

122 self.model = self.model.to(self.device) 

123 return 

124 

125 self.sp_size = 1 

126 self.load_timer.start("dit_parallel") 

127 for block in self.model.blocks: 

128 block.self_attn.forward = types.MethodType(usp_attn_forward, block.self_attn) 

129 self.model.forward = types.MethodType(usp_dit_forward, self.model) 

130 self.sp_size = get_sequence_parallel_world_size() 

131 self.load_timer.end("dit_parallel") 

132 

133 if dist.is_initialized(): 

134 dist.barrier() 

135 

136 # Load the DiT model across GPUs 

137 if self.world_size > 1: 

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

139 self.model = shard_fn(self.model) 

140 if not self.model: 

141 raise ValueError("Model sharding failed") 

142 self.model = self.model.to(self.device) 

143 

144 def model_compile(self) -> None: 

145 """Compile the Wan 2.1 model with torch.compile().""" 

146 if not self.torch_compile: 

147 return 

148 logging.info(f"[{self.rank}] Compiling DiT with torch.compile().") 

149 self.load_timer.start("dit_compile") 

150 self.model = torch.compile( 

151 self.model, 

152 mode="max-autotune-no-cudagraphs", 

153 ) 

154 self.load_timer.end("dit_compile") 

155 

156 def _assert_model_init(self) -> None: 

157 super()._assert_model_init() 

158 if not self.image_encoder: 

159 raise ValueError("Image encoder model not initialized") 

160 if not self.image_encoder.model: 

161 raise ValueError("Image encoder model not initialized") 

162 if not self.model: 

163 raise ValueError("DiT model not initialized") 

164 

165 @inference_mode() 

166 async def generate( 

167 self, 

168 img: Image.Image, 

169 prompt: str, 

170 neg_prompt: str = "", 

171 width: int = 640, 

172 height: int = 480, 

173 num_frames: int = 1 + 80, 

174 sampling_steps: int = 50, 

175 job_id: Optional[str] = None, 

176 output_type: str = "tensor" 

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

178 """ 

179 Generate a video from an image and a prompt. 

180 """ 

181 gen_timer = self._new_gen_timer(job_id) 

182 

183 start_frames = 1 

184 

185 self._assert_model_init() 

186 self._assert_args(height, width, num_frames, start_frames) 

187 

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

189 

190 try: 

191 # Convert image to normalized tensor 

192 img_resized = img.resize((width, height), Image.Resampling.LANCZOS) 

193 img_tensor_norm = TF.to_tensor(img_resized).sub_(0.5).div_(0.5).to(self.device) 

194 # img is PIL.Image.Image image mode=RGB size=640x480 at 0x7FA09029B1D0 

195 # img_tensor_norm: [3, 480, 640] [RGB, height, width] 

196 

197 h, w = img_tensor_norm.shape[1:] 

198 if not self.vae_stride: 

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

200 lat_h = h // self.vae_stride[1] 

201 lat_w = w // self.vae_stride[2] 

202 

203 vae_t = self.vae_stride[0] 

204 max_seq_len = ((num_frames - start_frames) 

205 // vae_t + start_frames) * lat_h * lat_w // (self.patch_size[1] * self.patch_size[2]) 

206 max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size 

207 if self.rank == 0: 

208 logging.info(f"[{self.rank}] size:{w}x{h}, lat_size:{lat_w}x{lat_h}, " 

209 f"#frames:{num_frames}, #start_frames:{start_frames}, max_seq_len:{max_seq_len}, " 

210 f"sp_size:{self.sp_size}, patch:{self.patch_size}, stride:{self.vae_stride}.") 

211 

212 seed = random.randint(0, sys.maxsize) 

213 if self.base_seed is not None and self.base_seed >= 0: 

214 seed = self.base_seed 

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

216 seed_g.manual_seed(seed) 

217 

218 # Adjusted from https://github.com/Wan-Video/Wan2.1/pull/100 -> 80+1 -> 20+1 

219 latent_num_frames = (num_frames - start_frames) // self.vae_stride[0] + start_frames 

220 noise = torch.randn( 

221 16, # latent channel 

222 latent_num_frames, 

223 lat_h, 

224 lat_w, 

225 dtype=torch.float32, 

226 generator=seed_g, 

227 device=self.device, 

228 ) # [lat_channel, #lat_frames, lat_h, lat_w] [16, 21, 68, 90] 

229 

230 # Preprocess 

231 # Text encoder 

232 gen_timer.start("text_encoder") 

233 if not self.text_encoder or not self.text_encoder.model: 

234 raise ValueError("Text encoder model not initialized") 

235 self.text_encoder.model.to(self.device) 

236 context = self.text_encoder([prompt], self.device) 

237 context_null = self.text_encoder([neg_prompt], self.device) 

238 gen_timer.end("text_encoder") 

239 

240 # Image encoder 

241 gen_timer.start("image_encoder") 

242 if not self.image_encoder or not self.image_encoder.model: 

243 raise ValueError("Image encoder model not initialized") 

244 self.image_encoder.model.to(self.device) 

245 clip_context = self.image_encoder.visual([img_tensor_norm[:, None, :, :]]) 

246 gen_timer.end("image_encoder") 

247 

248 # VAE encoder: Pixels -> Latent 

249 gen_timer.start("vae_encoder") 

250 

251 # Mask+image (in latent space) with 1s for the first frame (input image) and empty (0s) for the rest 

252 y_1_frame = None 

253 if img is not None: 

254 msk_1_frame = self._get_mask(lat_h, lat_w, num_frames, 1, 0) 

255 img_frame = torch.nn.functional.interpolate( 

256 img_tensor_norm[None].cpu(), 

257 size=(h, w), 

258 mode='bicubic' 

259 ).transpose(0, 1) # Shape: [RGB, 1, h, w] 

260 

261 num_empty_frames = num_frames - 1 

262 empty_frames = torch.zeros(3, num_empty_frames, h, w) # Shape: [RGB, #EmptyFrames, h, w] 

263 vid_frames = [ 

264 # img_frame * prev_frame_weight, # TODO check how to use this 

265 img_frame, 

266 empty_frames 

267 ] 

268 

269 # Create mask + y: pixels (img/video) to latent 

270 if not self.vae or not self.vae.model: 

271 raise ValueError("VAE model not initialized") 

272 y_1_frame = self.vae.encode([ 

273 torch.concat(vid_frames, dim=1).to(self.device) 

274 ])[0] 

275 y_1_frame = torch.concat([ 

276 msk_1_frame, 

277 y_1_frame 

278 ]) # [latent channels, #latent frames, lat_h, lat_w] [20, 21, 68, 90] 

279 

280 gen_timer.end("vae_encoder") 

281 

282 @contextmanager 

283 def noop_no_sync() -> Generator[None, None, None]: 

284 yield 

285 

286 no_sync = getattr(self.model, 'no_sync', noop_no_sync) 

287 

288 # DiT sampling 

289 x0 = [] 

290 with amp.autocast('cuda', dtype=self.param_dtype), torch.no_grad(), no_sync(): 

291 # Setup scheduler 

292 gen_timer.start("scheduler_setup") 

293 

294 sample_scheduler = FlowUniPCMultistepScheduler( 

295 num_train_timesteps=1000, 

296 shift=1, 

297 use_dynamic_shifting=False 

298 ) 

299 

300 sample_scheduler.set_timesteps(sampling_steps, device=self.device, shift=self.shift) 

301 timesteps = sample_scheduler.timesteps 

302 gen_timer.end("scheduler_setup") 

303 

304 # Sample videos 

305 latent = noise 

306 

307 if not self.model: 

308 raise ValueError("DiT model not initialized") 

309 self.model.to(self.device) 

310 

311 for it, t in enumerate(timesteps): 

312 logging.debug(f"[{self.rank}] Running step {it + 1}/{len(timesteps)}.") 

313 

314 self.check_interrupted() 

315 

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

317 

318 # TODO test batching 

319 latent_model_input = [latent.to(self.device)] 

320 timestep_list = [t] 

321 

322 timestep = torch.stack(timestep_list).to(self.device) 

323 

324 # Choose the mask depending on the previous video and stage 

325 y = y_1_frame # We use the mask starting from the image 

326 

327 arg_c = { 

328 "t": timestep, 

329 "context": [context[0]], 

330 "clip_fea": clip_context, 

331 "seq_len": max_seq_len, 

332 "y": [y], 

333 } 

334 arg_null = { 

335 "t": timestep, 

336 "context": context_null, 

337 "clip_fea": clip_context, 

338 "seq_len": max_seq_len, 

339 "y": [y], 

340 } 

341 

342 # TODO we can make this parallel 

343 # TODO we could batch the two of them 

344 # This takes a lot of memory because of the KV cache 

345 gen_timer.start(f"dit_cond_{it:03d}") 

346 noise_pred_cond = self.model( 

347 latent_model_input, 

348 **arg_c 

349 )[0].to(self.device) 

350 gen_timer.end(f"dit_cond_{it:03d}") 

351 

352 gen_timer.start(f"dit_uncond_{it:03d}") 

353 noise_pred_uncond = self.model( 

354 latent_model_input, 

355 **arg_null 

356 )[0].to(self.device) 

357 gen_timer.end(f"dit_uncond_{it:03d}") 

358 

359 noise_pred = noise_pred_uncond + self.guide_scale * (noise_pred_cond - noise_pred_uncond) 

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

361 

362 gen_timer.start(f"scheduler_{it:03d}") 

363 latent = latent.to(self.device) 

364 temp_x0 = sample_scheduler.step( 

365 noise_pred.unsqueeze(0), 

366 t, 

367 latent.unsqueeze(0), 

368 return_dict=False, 

369 generator=seed_g 

370 )[0] 

371 latent = temp_x0.squeeze(0) 

372 gen_timer.end(f"scheduler_{it:03d}") 

373 

374 x0 = [latent.to(self.device)] 

375 del latent_model_input, timestep 

376 

377 videos = None 

378 if self.rank == 0 and output_type != "latent": 

379 # Latent -> Pixels (video) 

380 gen_timer.start("vae_decoder") 

381 if not self.vae: 

382 raise ValueError("VAE model not initialized") 

383 videos = self.vae.decode(x0, start_frames=1, end_frames=0) 

384 gen_timer.end("vae_decoder") 

385 

386 del noise 

387 del sample_scheduler 

388 if dist.is_initialized(): 

389 dist.barrier() 

390 

391 if self.rank != 0: 

392 return None 

393 

394 if output_type == "latent": 

395 if not x0: 

396 raise ValueError("No latent generated") 

397 x0 = x0[0] 

398 return x0 

399 

400 if not videos: 

401 raise ValueError("No videos generated") 

402 video_tensor = videos[0] 

403 return await self._output_video( 

404 job_id, 

405 gen_timer, 

406 video_tensor, # C, T, H, W 

407 output_type) 

408 finally: 

409 self.running = False 

410 gen_timer.end("total") 

411 

412 def _save_video( 

413 self, 

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

415 video_path: str, 

416 ) -> str: 

417 assert video_tensor is not None 

418 assert isinstance(video_tensor, torch.Tensor) 

419 assert video_tensor.dim() == 4 

420 assert video_tensor.shape[0] == 3 # Channels 

421 return cache_video( 

422 tensor=video_tensor[None], # C, T, H, W -> B, C, T, H, W 

423 save_file=video_path, 

424 fps=self.FPS, 

425 nrow=1)