Coverage for wrapper/fantasytalking/wrapper_fantasytalking.py: 60%

1""" 

2Wrapper for Fantasy Talking video generation. 

3""" 

4import os 

5import sys 

6import types 

7import logging 

8import tempfile 

9import math 

10import asyncio 

11import aiofiles 

12import aiofiles.os 

13 

14from typing import override 

15from typing import List 

16from typing import Union 

17from typing import Optional 

18from typing import Dict 

19from typing import Any 

20from typing import Tuple 

21 

22from PIL import Image 

23 

24import torchvision.transforms.functional as TF 

25 

26import torch 

27import torch.distributed as dist 

28from torch import inference_mode 

29 

30import numpy as np 

31 

32from functools import partial 

33 

34from model_timing import GenTimer 

35from wrapper_model import ModelGeneration 

36from wrapper_usp import USPGeneration 

37 

38from image_utils import base64_to_img 

39from media_utils import base64_to_video_frames 

40from media_utils import base64_to_audio_file 

41from media_utils import empty_audio_file 

42from media_utils import save_video_audio 

43 

44import librosa 

45 

46from transformers import Wav2Vec2Model 

47from transformers import Wav2Vec2Processor 

48 

49from diffsynth import ModelManager 

50from diffsynth import WanVideoPipeline 

51from model import FantasyTalkingAudioConditionModel 

52 

53from utils import get_audio_features 

54 

55from xfuser.config import EngineConfig 

56 

57sys.path.append("/wan") # noqa: E402 

58from fantasytalking_xfuser import usp_fantasytalking_forward 

59from wan.distributed.xdit_context_parallel import usp_attn_forward 

60from wan.distributed.fsdp import shard_model 

61 

62 

63def resample_frames( 

64 video_frames: List[Image.Image], 

65 src_fps: float, 

66 dst_fps: float, 

67 audio_duration: Optional[float] = None 

68) -> List[Image.Image]: 

69 """ 

70 Instead of simply trimming or padding, we resample the frames to match the desired number of frames. 

71 # if video_num_frames < num_frames: 

72 # logging.warning( 

73 # f"[{self.rank}] Video {video_num_frames} < Output {num_frames}, " 

74 # "padding video with last frame.") 

75 # video += [video[-1]] * (num_frames - video_num_frames) 

76 """ 

77 

78 if src_fps > dst_fps and audio_duration: 

79 # Truncate the video to the first audio_duration seconds 

80 num_frames = int(audio_duration * src_fps) 

81 if num_frames < len(video_frames): 

82 logging.info(f"Truncating video from {len(video_frames)} to {num_frames} frames to match audio duration.") 

83 video_frames = video_frames[:num_frames] 

84 

85 duration = len(video_frames) / src_fps 

86 num_dst_frames = int(round(duration * dst_fps)) 

87 idxs = np.linspace(0, len(video_frames) - 1, num_dst_frames) 

88 idxs = idxs.astype(int) 

89 return [ 

90 video_frames[i] 

91 for i in idxs 

92 ] 

93 

94 

95def resample_and_normalize_frames( 

96 video_frames: List[Image.Image], 

97 src_fps: float, 

98 dst_fps: float, 

99 target_num_frames: int, 

100 audio_duration: Optional[float] = None 

101) -> List[Image.Image]: 

102 """Resample video frames from src_fps to dst_fps and normalize the result to exactly 

103 target_num_frames. 

104 

105 resample_frames() uses floating-point rounding which can produce a count that differs 

106 from target_num_frames by one frame. That off-by-one translates directly into a 

107 latent-tensor shape mismatch (e.g. 14 vs 15 frames) when the frames are later encoded 

108 by the VAE and fed to scheduler.add_noise(). This wrapper ensures the output always 

109 has exactly target_num_frames frames by padding with the last frame or trimming. 

110 """ 

111 resampled = resample_frames(video_frames, src_fps, dst_fps, audio_duration) 

112 if len(resampled) == 0: 

113 raise ValueError("Resampled video is empty; cannot normalize to target frame count.") 

114 if len(resampled) < target_num_frames: 

115 resampled = resampled + [resampled[-1]] * (target_num_frames - len(resampled)) 

116 elif len(resampled) > target_num_frames: 

117 resampled = resampled[:target_num_frames] 

118 return resampled 

119 

120 

121class FantasyTalking(USPGeneration): 

122 """ 

123 Fantasy Talking video generation wrapper. 

124 """ 

125 

126 # Weird FPS, but it is what Fantasy Talking uses 

127 FPS = 23.0 

128 SRC_FPS = 30.0 # from HunyuanVideo 

129 # Original limit for Wan at 16 FPS -> 5.1 seconds 

130 # self.MAX_FRAMES = 1 + 80 

131 # Increased the limit from 1+80 to 1+116 which is ~5.1 seconds 

132 MAX_FRAMES = 1 + 116 

133 # Wan values for number of attention heads 

134 NUM_HEADS = 40 

135 

136 def __init__( 

137 self, 

138 model_name: str = "fantasytalking", 

139 engine_config: EngineConfig = None, 

140 param_dtype: torch.dtype = torch.bfloat16, 

141 ) -> None: 

142 super().__init__( 

143 model_name, 

144 engine_config, 

145 param_dtype) 

146 

147 # Model components 

148 self.pipeline_wan: Optional[WanVideoPipeline] = None 

149 self.fantasytalking: Optional[FantasyTalkingAudioConditionModel] = None 

150 self.wav2vec_processor: Optional[Wav2Vec2Processor] = None 

151 self.wav2vec: Optional[Wav2Vec2Model] = None 

152 

153 # Model features 

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

155 

156 def __del__(self) -> None: 

157 # Clean models 

158 if self.pipeline_wan is not None: 

159 self.pipeline_wan = None 

160 if self.fantasytalking is not None: 

161 self.fantasytalking = None 

162 if self.wav2vec_processor is not None: 

163 self.wav2vec_processor = None 

164 if self.wav2vec is not None: 

165 self.wav2vec = None 

166 super().__del__() 

167 

168 def load_model(self) -> None: 

169 assert torch.cuda.is_available() 

170 

171 self.load_timer.start("wan") 

172 model_manager = ModelManager(device="cpu") 

173 BASE_MODELS_FOLDER = "/fantasytalking" 

174 NUM_MODEL_CHUNKS = 7 

175 wan_model_path = f"{BASE_MODELS_FOLDER}/Wan2.1-I2V-14B-720P" 

176 model_manager.load_models( 

177 [ 

178 [ 

179 f"{wan_model_path}/diffusion_pytorch_model-{idx:05d}-of-00007.safetensors" 

180 for idx in range(1, NUM_MODEL_CHUNKS + 1) 

181 ], 

182 f"{wan_model_path}/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", 

183 f"{wan_model_path}/models_t5_umt5-xxl-enc-bf16.pth", 

184 f"{wan_model_path}/Wan2.1_VAE.pth", 

185 ], 

186 torch_dtype=self.param_dtype 

187 ) 

188 # self.pipeline_wan = WanVideoPipeline.from_model_manager( 

189 self.pipeline_wan = CustomWanVideoPipeline( 

190 device=self.device, 

191 torch_dtype=self.param_dtype) 

192 if not self.pipeline_wan: 

193 raise ValueError("Wan model not initialized") 

194 self.pipeline_wan.fetch_models(model_manager) 

195 # self.pipeline_wan.enable_vram_management(num_persistent_param_in_dit=None) 

196 self.load_timer.end("wan") 

197 

198 # Load FantasyTalking weights 

199 self.load_timer.start("fantasytalking") 

200 self.fantasytalking = FantasyTalkingAudioConditionModel( 

201 self.pipeline_wan.dit, 

202 audio_in_dim=768, 

203 audio_proj_dim=2048 

204 ).to(self.device) 

205 if self.fantasytalking is None: 

206 raise ValueError("Fantasy Talking model not initialized") 

207 self.fantasytalking.load_audio_processor( 

208 f"{BASE_MODELS_FOLDER}/fantasytalking_model.ckpt", 

209 self.pipeline_wan.dit 

210 ) 

211 self.load_timer.end("fantasytalking") 

212 

213 # Load wav2vec models 

214 self.load_timer.start("wav2vec") 

215 self.wav2vec_processor = Wav2Vec2Processor.from_pretrained( 

216 f"{BASE_MODELS_FOLDER}/wav2vec2-base-960h" # nosec B615 - local path 

217 ) 

218 self.wav2vec = Wav2Vec2Model.from_pretrained( 

219 f"{BASE_MODELS_FOLDER}/wav2vec2-base-960h" # nosec B615 - local path 

220 ).to(self.device) 

221 self.load_timer.end("wav2vec") 

222 

223 def init_model_parallelism(self) -> None: 

224 if self.pipeline_wan is None: 

225 raise RuntimeError("Pipeline WAN not initialized") 

226 

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

228 self.pipeline_wan.to(self.device) 

229 return 

230 

231 self.load_timer.start("dit_parallel") 

232 for block in self.pipeline_wan.dit.blocks: 

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

234 self.pipeline_wan.dit.forward = types.MethodType(usp_fantasytalking_forward, self.pipeline_wan.dit) 

235 

236 # Load across GPUs 

237 shard_fn = None 

238 if self.world_size > 1: 

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

240 self.pipeline_wan = self.pipeline_wan.to(self.param_dtype) 

241 if not self.pipeline_wan: 

242 raise ValueError("Wan not initialized for sharding") 

243 self.pipeline_wan.dit = shard_fn(self.pipeline_wan.dit) 

244 self.load_timer.end("dit_parallel") 

245 

246 self.pipeline_wan.to(self.device) 

247 

248 def model_compile(self) -> None: 

249 if not self.torch_compile: 

250 return 

251 if not self.pipeline_wan: 

252 raise RuntimeError("Pipeline Wan not initialized") 

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

254 self.load_timer.start("compile") 

255 self.pipeline_wan.dit = torch.compile( 

256 self.pipeline_wan.dit, 

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

258 ) 

259 # This may cause issues in the video processing, and doesn't really give much speedup 

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

261 self.pipeline_wan.vae = torch.compile( 

262 self.pipeline_wan.vae, 

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

264 ) 

265 logging.info(f"[{self.rank}] Compiling Fantasy Talking with torch.compile().") 

266 self.fantasytalking = torch.compile( 

267 self.fantasytalking, 

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

269 ) 

270 self.load_timer.end("compile") 

271 

272 def _assert_model_init(self) -> None: 

273 super()._assert_model_init() 

274 assert self.pipeline_wan is not None, "Pipeline WAN not initialized" 

275 assert self.fantasytalking is not None, "Fantasy Talking model not initialized" 

276 assert self.wav2vec_processor is not None, "Wav2Vec processor not initialized" 

277 assert self.wav2vec is not None, "Wav2Vec model not initialized" 

278 

279 def _assert_args(self, height: int, width: int) -> None: 

280 if not self.vae_stride: 

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

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

283 raise ValueError(f"Height {height} is not divisible by VAE scale factor {self.vae_stride}") 

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

285 raise ValueError(f"Width {width} is not divisible by VAE scale factor {self.vae_stride}") 

286 # TODO add check for sizes based on world_size similar to what we do in Flux 

287 """ 

288 # Check if the image size is supported for the current parallelism setting 

289 # https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux.py 

290 height_latent = height // self.vae_stride[1] 

291 width_latent = width // self.vae_stride[2] 

292 img_latent_shape = (height_latent // 2) * (width_latent // 2) 

293 if img_latent_shape % self.world_size != 0: 

294 raise ValueError(f"{height}x{width} not supported for {self.world_size} GPUs.") 

295 """ 

296 

297 @inference_mode() 

298 async def warmup(self) -> None: 

299 logging.info(f"[{self.rank}] Warmup for Fantasy Talking generation.") 

300 empty_img = Image.new("RGB", (640, 480), color=(255, 255, 255)) 

301 empty_audio_path = empty_audio_file(0.5) # 0.5 seconds of silence 

302 await self.generate( 

303 job_id="warmup", 

304 img=empty_img, 

305 video=None, 

306 audio_path=empty_audio_path, 

307 prompt="A warmup generation", 

308 neg_prompt="", 

309 width=640, 

310 height=480, 

311 sampling_steps=2) 

312 if os.path.exists(empty_audio_path): 

313 os.remove(empty_audio_path) 

314 

315 def _to_num_latent_frames(self, num_frames: int) -> int: 

316 if not self.vae_stride: 

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

318 return (num_frames - 1) // self.vae_stride[0] + 1 

319 

320 def _to_num_frames(self, num_latent_frames: int) -> int: 

321 if not self.vae_stride: 

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

323 return num_latent_frames * self.vae_stride[0] + 1 

324 

325 def _get_audio_num_frames( 

326 self, 

327 audio_path: str 

328 ) -> Tuple[float, int, int]: 

329 # We need to align the audio frame number with (1 + 4n) frames 

330 if not self.vae_stride: 

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

332 audio_duration = librosa.get_duration(path=audio_path) 

333 audio_num_frames = int(math.ceil(self.FPS * audio_duration)) 

334 video_num_frames = int(1 + math.ceil((audio_num_frames - 1) / self.vae_stride[0]) * self.vae_stride[0]) 

335 return audio_duration, audio_num_frames, video_num_frames 

336 

337 @override 

338 @inference_mode() 

339 async def generate( 

340 self, 

341 img: Image.Image, 

342 video: Optional[List[Image.Image]], 

343 audio_path: str, 

344 prompt: str, 

345 neg_prompt: str = "", 

346 width: int = 1280, 

347 height: int = 720, 

348 sampling_steps: int = 30, # 10 

349 audio_scale: float = 1.0, # 1.0 for audio, 0.0 for no audio influence 

350 cfg_scale: float = 5.0, # how much does video follow prompt 

351 audio_cfg_scale: float = 5.0, # how much does video follow audio 

352 # We use this to avoid first frame: https://github.com/Fantasy-AMAP/fantasy-talking/issues/52 

353 end_percent: float = 0.9, 

354 adjust_durations: bool = True, 

355 job_id: Optional[str] = None, 

356 output_type: str = "pil", # "pil", "video_binary", "video_path" 

357 ) -> Union[List[Image.Image], str, bytes, None]: 

358 """ 

359 Generate a video from an image, and audio, and a prompt. 

360 Args: 

361 img (Image.Image): Input image. 

362 audio_path (str): Path to the audio file. 

363 prompt (str): Text prompt for video generation. 

364 neg_prompt (str, optional): Negative text prompt. Defaults to "". 

365 width (int, optional): Width of the output video. Defaults to 1280. 

366 height (int, optional): Height of the output video. Defaults to 720. 

367 sampling_steps (int, optional): Number of sampling steps. Defaults to 30. 

368 audio_scale (float, optional): Scale for audio influence. Defaults to 1.0. 

369 cfg_scale (float, optional): Scale for prompt influence. Defaults to 5.0. 

370 audio_cfg_scale (float, optional): Scale for audio prompt influence. Defaults to 5.0. 

371 end_percent (float, optional): Percentage of de-noising steps considering audio inputs. Defaults to 0.9. 

372 adjust_durations (bool, optional): Whether to adjust video durations to match audio. Defaults to True. 

373 job_id (str, optional): Job ID for tracking. Defaults to None. 

374 output_type (str, optional): Output type. Can be "pil", "video_binary", or "video_path". Defaults to "pil". 

375 """ 

376 gen_timer = self._new_gen_timer(job_id) 

377 

378 logging.info(f"[{self.rank}] Generating video for '{job_id if job_id else prompt[0:80]}'.") 

379 

380 self._assert_model_init() 

381 self._assert_args(height, width) 

382 

383 if not await aiofiles.os.path.exists(audio_path): 

384 raise ValueError(f"Audio file '{audio_path}' does not exist") 

385 

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

387 

388 try: 

389 audio_duration, audio_num_frames, num_frames = self._get_audio_num_frames(audio_path) 

390 

391 if num_frames > self.MAX_FRAMES: 

392 raise ValueError( 

393 f"Audio {audio_duration:.3f}s exceeds maximum frames {self.MAX_FRAMES} at {self.FPS} FPS " 

394 f"-> {num_frames} frames") 

395 

396 lat_num_frames = self._to_num_latent_frames(num_frames) 

397 if self.rank == 0: 

398 logging.info(f"[{self.rank}] Audio:{audio_duration:.3f}s #audio_frames:{audio_num_frames} " 

399 f"#video_frames:{num_frames} #lat_frames:{lat_num_frames} FPS:{self.FPS}.") 

400 

401 if not img and not video: 

402 raise ValueError("No image or video provided") 

403 

404 if video is not None and len(video) > 0: 

405 gen_timer.start("video_preprocess") 

406 video = [frame.resize((width, height), Image.Resampling.LANCZOS) for frame in video] 

407 

408 # We assume the input video has the same FPS 

409 video_num_frames = len(video) 

410 video_duration = num_frames / self.FPS 

411 if video_num_frames != num_frames: 

412 if not adjust_durations: 

413 raise ValueError(f"Video frames {video_num_frames} != Output frames {num_frames}") 

414 logging.warning( 

415 f"[{self.rank}] Resampling video from {video_num_frames} (source) to {num_frames} " 

416 "(destination) frames to match output.") 

417 video = resample_and_normalize_frames( 

418 video, 

419 src_fps=self.SRC_FPS, dst_fps=self.FPS, 

420 target_num_frames=num_frames, 

421 audio_duration=audio_duration) 

422 

423 if img is None: 

424 if self.rank == 0: 

425 logging.info(f"[{self.rank}] Using first video frame as starting image.") 

426 img = video[0] 

427 if self.rank == 0: 

428 video_num_frames = len(video) 

429 logging.info( 

430 f"[{self.rank}] Input video:{video_duration:.3f}s #frames:{video_num_frames} FPS:{self.FPS}.") 

431 gen_timer.end("video_preprocess") 

432 

433 if img.size != (width, height): 

434 gen_timer.start("img_preprocess") 

435 if self.rank == 0: 

436 logging.info(f"[{self.rank}] Image:{img.size}->{(width, height)}.") 

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

438 gen_timer.end("img_preprocess") 

439 

440 gen_timer.start("audio_encoder") 

441 audio_wav2vec_fea = get_audio_features( 

442 self.wav2vec, 

443 self.wav2vec_processor, 

444 audio_path, 

445 self.FPS, 

446 num_frames 

447 ) 

448 if self.fantasytalking is None: 

449 raise ValueError("Fantasy Talking model not initialized") 

450 audio_proj_fea = self.fantasytalking.get_proj_fea(audio_wav2vec_fea) 

451 pos_idx_ranges = self.fantasytalking.split_audio_sequence( 

452 audio_proj_fea.size(1), 

453 num_frames=num_frames 

454 ) 

455 audio_proj_split, audio_context_lens = self.fantasytalking.split_tensor_with_padding( 

456 audio_proj_fea, 

457 pos_idx_ranges, 

458 expand_length=4, 

459 ) 

460 gen_timer.end("audio_encoder") 

461 

462 gen_timer.start("wan") 

463 if not self.pipeline_wan: 

464 raise ValueError("Pipeline Wan not initialized") 

465 self.pipeline_wan.to(self.device) 

466 video_frames = await asyncio.to_thread( 

467 self.pipeline_wan, 

468 parent=self, 

469 prompt=prompt, 

470 negative_prompt=neg_prompt, 

471 input_image=img, 

472 input_video=video, 

473 width=width, 

474 height=height, 

475 num_frames=num_frames, 

476 latents_num_frames=lat_num_frames, 

477 num_inference_steps=sampling_steps, 

478 seed=self.base_seed, 

479 tiled=True, 

480 audio_scale=audio_scale, 

481 cfg_scale=cfg_scale, 

482 audio_cfg_scale=audio_cfg_scale, 

483 audio_proj=audio_proj_split, 

484 audio_context_lens=audio_context_lens, 

485 gen_timer=gen_timer, 

486 end_percent=end_percent, # Use 0.9 to control how much de-noising steps consider audio inputs 

487 ) 

488 gen_timer.end("wan") 

489 

490 if self.rank != 0: 

491 return None # Only rank 0 returns something 

492 

493 # Because of the 1+4n alignment, we might have generated more frames than audio, trim the video to match 

494 if audio_num_frames < len(video_frames): 

495 logging.info( 

496 f"[{self.rank}] Trimming number of video frames from {len(video_frames)} to " 

497 f"{audio_num_frames} to match audio frames.") 

498 video_frames = video_frames[0:audio_num_frames] 

499 

500 return await self._output_video( 

501 job_id, 

502 gen_timer, 

503 audio_path, 

504 video_frames, 

505 output_type) 

506 finally: 

507 self.running = False 

508 torch.cuda.empty_cache() 

509 gen_timer.end("total") 

510 

511 async def _output_video( 

512 self, 

513 job_id: Optional[str], 

514 gen_timer: GenTimer, 

515 audio_path: str, 

516 video_frames: List[Image.Image], 

517 output_type: str = "pil", # "pil", "video_binary", "video_path" 

518 ) -> Union[List[Image.Image], str, bytes, None]: 

519 gen_timer.start("output_video") 

520 try: 

521 if output_type == "pil": 

522 return video_frames 

523 

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

525 if not job_id: 

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

527 else: 

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

529 video_path = await save_video_audio( 

530 video_content=video_frames, 

531 audio_path=audio_path, 

532 out_video_path=video_path, 

533 fps=self.FPS) 

534 if output_type == "video_path": 

535 return video_path 

536 

537 # video_binary 

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

539 video_binary = await file.read() 

540 return video_binary 

541 

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

543 return None 

544 finally: 

545 gen_timer.end("output_video") 

546 

547 async def get_rest_args( 

548 self, 

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

550 ) -> Dict[str, Any]: 

551 if data_json is None: 

552 raise ValueError("Missing JSON body") 

553 

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

555 

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

557 img = None 

558 if img_base64: 

559 if not isinstance(img_base64, str): 

560 raise ValueError("Invalid 'img' parameter") 

561 img = base64_to_img(img_base64) 

562 

563 video_base64 = data_json.get("video", None) 

564 video = None 

565 if video_base64: 

566 if not isinstance(video_base64, str): 

567 raise ValueError("Invalid 'video' parameter") 

568 video = base64_to_video_frames(video_base64) 

569 

570 audio_base64 = data_json.get("audio", None) 

571 if not audio_base64: 

572 raise ValueError("Missing 'audio' parameter") 

573 if not isinstance(audio_base64, str): 

574 raise ValueError("Invalid 'audio' parameter") 

575 audio_path = None 

576 if not job_id: 

577 audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name 

578 else: 

579 audio_path = f"/tmp/{job_id}.wav" 

580 audio_path = await base64_to_audio_file( 

581 audio_base64, 

582 audio_path=audio_path) 

583 

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

585 if prompt is None: 

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

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

588 

589 gen_args = { 

590 "task": self.model_name, 

591 "args": { 

592 "job_id": job_id, 

593 "img": img, 

594 "video": video, 

595 "prompt": prompt, 

596 "neg_prompt": neg_prompt, 

597 "audio_path": audio_path, 

598 "width": int(data_json.get("width", 640)), 

599 "height": int(data_json.get("height", 480)), 

600 "sampling_steps": int(data_json.get("sampling_steps", 10)), 

601 "audio_scale": float(data_json.get("audio_scale", 1.0)), 

602 "cfg_scale": float(data_json.get("cfg_scale", 5.0)), 

603 "audio_cfg_scale": float(data_json.get("audio_cfg_scale", 5.0)), 

604 "end_percent": float(data_json.get("end_percent", 0.9)), 

605 "output_type": data_json.get("output_type", "pil"), 

606 } 

607 } 

608 return gen_args 

609 

610 

611class CustomWanVideoPipeline(WanVideoPipeline): 

612 """ 

613 Customize WanVideoPipeline to support end_percent (0.0-1.0). 

614 This is used to control how much de-noising steps are considering audio inputs. 

615 E.g., a 0.9 end_percent means that the last 10% of de-noising steps will not consider audio inputs. 

616 https://github.com/Fantasy-AMAP/fantasy-talking/issues/52 

617 We also add a timer. 

618 """ 

619 @staticmethod 

620 def from_model_manager( 

621 model_manager: ModelManager, 

622 torch_dtype: Optional[torch.dtype] = None, 

623 device: Optional[str] = None, 

624 ) -> "CustomWanVideoPipeline": 

625 if device is None: 

626 device = model_manager.device 

627 if torch_dtype is None: 

628 torch_dtype = model_manager.torch_dtype 

629 pipe = CustomWanVideoPipeline( 

630 device=device, 

631 torch_dtype=torch_dtype) 

632 pipe.fetch_models(model_manager) 

633 return pipe 

634 

635 @torch.no_grad() 

636 def __call__( 

637 self, 

638 parent: ModelGeneration, 

639 prompt: str, 

640 negative_prompt: str = "", 

641 input_image: Optional[Image.Image] = None, 

642 input_video: Optional[List[Image.Image]] = None, 

643 denoising_strength: float = 1.0, 

644 seed: Optional[int] = None, 

645 rand_device: str = "cpu", 

646 height: int = 480, 

647 width: int = 832, 

648 num_frames: int = 1 + 80, 

649 cfg_scale: float = 5.0, 

650 audio_cfg_scale: Optional[float] = None, 

651 num_inference_steps: int = 50, 

652 sigma_shift: float = 5.0, 

653 tiled: bool = True, 

654 tile_size: Tuple[int, int] = (30, 52), 

655 tile_stride: Tuple[int, int] = (15, 26), 

656 gen_timer: Optional[GenTimer] = None, 

657 end_percent: float = 0.9, 

658 **kwargs: Any, 

659 ) -> List[Image.Image]: 

660 # Parameter check 

661 vae_stride = (4, 8, 8) # time, height, width 

662 if gen_timer is None: 

663 raise ValueError("gen_timer is required for timing") 

664 height, width = self.check_resize_height_width(height, width) 

665 if num_frames % vae_stride[0] != 1: 

666 num_frames = (num_frames + 2) // vae_stride[0] * vae_stride[0] + 1 

667 logging.info(f"Only 'num_frames % 4 != 1' is acceptable. We round it up to {num_frames}.") 

668 if end_percent < 0.0 or end_percent > 1.0: 

669 raise ValueError(f"'end_percent' must be in [0.0, 1.0], got {end_percent}") 

670 

671 # Tiler parameters 

672 tiler_kwargs = { 

673 "tiled": tiled, 

674 "tile_size": tile_size, 

675 "tile_stride": tile_stride, 

676 } 

677 

678 # Scheduler 

679 self.scheduler.set_timesteps( 

680 num_inference_steps, 

681 denoising_strength, 

682 shift=sigma_shift 

683 ) 

684 

685 # Initialize noise 

686 gen_timer.start("vae_encode") 

687 BATCH_SIZE = 1 

688 LAT_CHANNELS = 16 # VAE latent channels 

689 num_lat_frames = (num_frames - 1) // vae_stride[0] + 1 

690 lat_h = height // vae_stride[1] 

691 lat_w = width // vae_stride[2] 

692 noise = self.generate_noise( 

693 (BATCH_SIZE, LAT_CHANNELS, num_lat_frames, lat_h, lat_w), 

694 seed=seed, 

695 device=rand_device, 

696 dtype=torch.float32, 

697 ).to(self.device) 

698 

699 if input_video is not None: 

700 # Resize each frame to match adjusted height/width from self.check_resize_height_width 

701 input_video = [ 

702 TF.resize(frame, size=(height, width), antialias=True) 

703 for frame in input_video 

704 ] 

705 input_video = self.preprocess_images(input_video) 

706 input_video = torch.stack(input_video, dim=2) # type: ignore[arg-type, assignment] 

707 latents = self.encode_video(input_video, **tiler_kwargs).to( 

708 dtype=noise.dtype, 

709 device=noise.device) 

710 latents = self.scheduler.add_noise( 

711 latents, 

712 noise, 

713 timestep=self.scheduler.timesteps[0]) 

714 else: 

715 latents = noise 

716 gen_timer.end("vae_encode") 

717 

718 # Encode prompts 

719 gen_timer.start("encode_prompt") 

720 prompt_emb_nega = {} 

721 prompt_emb_posi = self.encode_prompt(prompt, positive=True) 

722 if cfg_scale != 1.0: 

723 prompt_emb_nega = self.encode_prompt(negative_prompt, positive=False) 

724 gen_timer.end("encode_prompt") 

725 

726 # Encode image 

727 gen_timer.start("encode_image") 

728 image_emb = {} 

729 if input_image is not None and self.image_encoder is not None: 

730 image_emb = self.encode_image(input_image, num_frames, height, width) 

731 gen_timer.end("encode_image") 

732 

733 # Extra input 

734 extra_input = self.prepare_extra_input(latents) 

735 

736 # De-noise steps 

737 with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self.device).type): 

738 total_steps = len(self.scheduler.timesteps) 

739 total_steps_considering_audio = min(int(total_steps * end_percent), total_steps) 

740 for progress_id, timestep in enumerate(self.scheduler.timesteps): 

741 logging.debug(f"Running step {progress_id + 1}/{total_steps}.") 

742 

743 parent.check_interrupted() 

744 

745 gen_timer.start(f"dit_{progress_id:03d}") 

746 if progress_id >= total_steps_considering_audio: 

747 logging.debug(f"Skipping audio at step {progress_id}.") 

748 audio_cfg_scale = 0.0 

749 

750 timestep = timestep.unsqueeze(0).to( 

751 dtype=torch.float32, 

752 device=self.device) 

753 

754 # Inference 

755 noise_pred_posi = self.dit( 

756 latents, 

757 timestep=timestep, 

758 **prompt_emb_posi, 

759 **image_emb, 

760 **extra_input, 

761 **kwargs, 

762 ) # (zt,audio,prompt) 

763 if audio_cfg_scale is not None: 

764 audio_scale = kwargs["audio_scale"] 

765 kwargs["audio_scale"] = 0.0 

766 noise_pred_noaudio = self.dit( 

767 latents, 

768 timestep=timestep, 

769 **prompt_emb_posi, 

770 **image_emb, 

771 **extra_input, 

772 **kwargs, 

773 ) # (zt,0,prompt) 

774 # kwargs['ip_scale'] = ip_scale 

775 if cfg_scale != 1.0: # prompt cfg 

776 noise_pred_no_cond = self.dit( 

777 latents, 

778 timestep=timestep, 

779 **prompt_emb_nega,