Coverage for wrapper/wan22/wrapper_wan22.py: 75%
131 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Handle video generation using the Wan 2.2 S2V (Sound-to-Video) model.
4Reference: https://github.com/Wan-Video/Wan2.2
5"""
7import asyncio
8import logging
9import os
10import tempfile
12from typing import Dict
13from typing import Union
14from typing import List
15from typing import Optional
16from typing import Any
18from PIL import Image
20import torch
21import torch.distributed as dist
22from torch import inference_mode
24from wrapper_wan import WanVideoGeneration
25from media_utils import base64_to_audio_file as _base64_to_audio_file
26from media_utils import empty_audio_file as _empty_audio_file
28import wan
29from wan.configs import WAN_CONFIGS
30from wan.utils.utils import save_video
32from xfuser.config import EngineConfig
35class Wan22VideoGeneration(WanVideoGeneration):
36 """Handle video generation using the Wan 2.2 S2V (Sound-to-Video) model.
38 Uses the official wan.WanS2V pipeline to generate video clips driven by
39 speech/audio and conditioned on a reference image and text prompt.
41 Supports:
42 - Direct audio input (WAV/MP3 file path)
43 - TTS audio synthesis via CosyVoice (enable_tts=True)
44 - Automatic video length based on audio duration (num_clip=None)
45 """
47 def __init__(
48 self,
49 model_name: str = "wan22",
50 ckpt_dir: str = "./Wan2.2-S2V-14B",
51 engine_config: EngineConfig = None,
52 param_dtype: torch.dtype = torch.bfloat16,
53 offload_model: bool = True,
54 ) -> None:
55 super().__init__(
56 model_name=model_name,
57 ckpt_dir=ckpt_dir,
58 engine_config=engine_config,
59 param_dtype=param_dtype,
60 )
61 self.offload_model = offload_model
63 # WanS2V pipeline instance (set in load_model)
64 self.wan_s2v: Optional[wan.WanS2V] = None
66 # S2V defaults matching wan_s2v_14B config
67 self.shift = 3.0
68 self.guide_scale = 4.5
70 def __del__(self) -> None:
71 if self.wan_s2v is not None:
72 del self.wan_s2v
73 super().__del__()
75 def load_model(self) -> None:
76 """Load the Wan 2.2 S2V model using the official WanS2V pipeline."""
77 assert torch.cuda.is_available()
79 cfg = WAN_CONFIGS['s2v-14B']
81 use_fsdp = self.world_size > 1
83 prev_memory = torch.cuda.memory_allocated()
84 self.load_timer.start("wan_s2v")
85 self.wan_s2v = wan.WanS2V(
86 config=cfg,
87 checkpoint_dir=self.ckpt_dir,
88 device_id=self.device_id,
89 rank=self.rank,
90 t5_fsdp=use_fsdp,
91 dit_fsdp=use_fsdp,
92 use_sp=False,
93 t5_cpu=True,
94 convert_model_dtype=True,
95 )
96 self.load_timer.end("wan_s2v")
98 diff_memory = torch.cuda.memory_allocated() - prev_memory
99 logging.info(f"[{self.rank}] WanS2V memory allocated: {diff_memory / 1024 / 1024 ** 2:.2f} GB.")
101 # Expose pipeline sub-components for _assert_model_init compatibility
102 self.text_encoder = self.wan_s2v.text_encoder
103 self.vae = self.wan_s2v.vae
105 def init_model_parallelism(self) -> None:
106 """Model parallelism is configured inside WanS2V.__init__; no additional setup needed."""
108 def model_compile(self) -> None:
109 """Compile the DiT model with torch.compile()."""
110 if not self.torch_compile:
111 return
112 if self.wan_s2v is None:
113 return
114 logging.info(f"[{self.rank}] Compiling WanS2V DiT with torch.compile().")
115 self.load_timer.start("dit_compile")
116 self.wan_s2v.noise_model = torch.compile(
117 self.wan_s2v.noise_model,
118 mode="max-autotune-no-cudagraphs",
119 )
120 self.load_timer.end("dit_compile")
122 def _assert_model_init(self) -> None:
123 super()._assert_model_init()
124 if self.wan_s2v is None:
125 raise ValueError("WanS2V model not initialized")
127 @inference_mode()
128 async def warmup(self) -> None:
129 logging.info(f"[{self.rank}] Warmup for Wan 2.2 S2V generation.")
130 audio_path: Optional[str] = None
131 try:
132 audio_path = _empty_audio_file(duration_seconds=1.0)
133 await self.generate(
134 img=Image.new("RGB", (704, 480), (128, 128, 128)),
135 prompt="Warmup prompt",
136 neg_prompt="",
137 max_area=480 * 704,
138 sampling_steps=2,
139 audio_path=audio_path,
140 infer_frames=4,
141 num_clip=1,
142 )
143 finally:
144 if audio_path and os.path.exists(audio_path):
145 os.unlink(audio_path)
147 @inference_mode()
148 async def generate(
149 self,
150 img: Image.Image,
151 prompt: str,
152 neg_prompt: str = "",
153 max_area: int = 1024 * 704,
154 sampling_steps: int = 40,
155 audio_path: Optional[str] = None,
156 enable_tts: bool = False,
157 tts_prompt_audio: Optional[str] = None,
158 tts_prompt_text: Optional[str] = None,
159 tts_text: Optional[str] = None,
160 num_clip: Optional[int] = None,
161 infer_frames: int = 80,
162 job_id: Optional[str] = None,
163 output_type: str = "tensor",
164 ) -> Union[List[Image.Image], str, bytes, torch.Tensor, None]:
165 """Generate a video clip driven by audio and conditioned on an image and prompt.
167 Args:
168 img: Reference image used as the visual anchor for generation.
169 prompt: Text description guiding video content.
170 neg_prompt: Negative prompt to suppress unwanted content.
171 max_area: Maximum pixel area for the output video (width * height).
172 The actual resolution is derived from the input image aspect ratio.
173 sampling_steps: Number of diffusion denoising steps.
174 audio_path: Path to the driving audio file (WAV/MP3).
175 Required when enable_tts is False.
176 enable_tts: If True, synthesise audio from text via CosyVoice instead
177 of using a pre-recorded audio file.
178 tts_prompt_audio: Path to reference speaker audio for zero-shot TTS.
179 Used only when enable_tts is True.
180 tts_prompt_text: Transcript matching tts_prompt_audio.
181 Used only when enable_tts is True.
182 tts_text: Text to synthesise into speech.
183 Used only when enable_tts is True.
184 num_clip: Number of video clips to generate. When None the pipeline
185 infers the count automatically from the audio length.
186 infer_frames: Frames generated per clip (must be a multiple of 4).
187 job_id: Identifier for this generation job (used for temp file naming).
188 output_type: One of "tensor", "pil", "video_binary", or "video_path".
190 Returns:
191 Generated video in the requested format, or None on non-primary ranks.
192 """
193 gen_timer = self._new_gen_timer(job_id)
194 self._assert_model_init()
195 assert self.wan_s2v is not None # guaranteed by _assert_model_init
196 self.running = True
198 img_path: Optional[str] = None
199 try:
200 # Save PIL image to a temporary file; WanS2V expects a file path
201 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
202 img_path = tmp.name
203 img.save(img_path)
205 gen_timer.start("s2v_generate")
206 video_tensor = await asyncio.to_thread(
207 self.wan_s2v.generate,
208 input_prompt=prompt,
209 ref_image_path=img_path,
210 audio_path=audio_path,
211 enable_tts=enable_tts,
212 tts_prompt_audio=tts_prompt_audio,
213 tts_prompt_text=tts_prompt_text,
214 tts_text=tts_text,
215 num_repeat=num_clip,
216 max_area=max_area,
217 infer_frames=infer_frames,
218 shift=self.shift,
219 sampling_steps=sampling_steps,
220 guide_scale=self.guide_scale,
221 n_prompt=neg_prompt,
222 seed=self.base_seed,
223 offload_model=self.offload_model,
224 )
225 gen_timer.end("s2v_generate")
227 if dist.is_initialized():
228 dist.barrier()
230 if self.rank != 0:
231 return None
233 if video_tensor is None:
234 raise ValueError("No video generated")
236 return await self._output_video(job_id, gen_timer, video_tensor, output_type)
237 finally:
238 self.running = False
239 gen_timer.end("total")
240 if img_path and os.path.exists(img_path):
241 os.unlink(img_path)
243 def _save_video(
244 self,
245 video_tensor: torch.Tensor, # C, T, H, W
246 video_path: str,
247 ) -> str:
248 assert video_tensor is not None
249 assert isinstance(video_tensor, torch.Tensor)
250 assert video_tensor.dim() == 4
251 assert video_tensor.shape[0] == 3 # RGB channels
252 return save_video(
253 tensor=video_tensor[None], # C, T, H, W -> B, C, T, H, W
254 save_file=video_path,
255 fps=self.FPS,
256 nrow=1,
257 normalize=True,
258 value_range=(-1, 1),
259 )
261 async def get_rest_args(
262 self,
263 data_json: Dict[str, Any],
264 ) -> Dict[str, Any]:
265 """Parse REST API request for Wan 2.2 S2V generation.
267 Expected JSON fields:
268 img (str): Base64-encoded reference image.
269 prompt (str): Text prompt describing the video.
270 neg_prompt (str, optional): Negative prompt.
271 max_area (int, optional): Maximum output pixel area (default 1024*704).
272 sampling_steps (int, optional): Diffusion steps (default 40).
273 audio (str, optional): Base64-encoded audio file. Required when
274 enable_tts is false/absent.
275 enable_tts (bool, optional): Use CosyVoice TTS to generate audio.
276 tts_prompt_audio (str, optional): Base64-encoded TTS reference audio.
277 tts_prompt_text (str, optional): Transcript for TTS reference audio.
278 tts_text (str, optional): Text to synthesise when enable_tts is True.
279 num_clip (int, optional): Number of video clips (auto from audio if absent).
280 infer_frames (int, optional): Frames per clip (default 80).
281 output_type (str, optional): "tensor", "pil", "video_binary", or "video_path".
282 """
283 if data_json is None:
284 raise ValueError("Missing JSON body")
286 from image_utils import base64_to_img
288 job_id = data_json.get("job_id", None)
290 img_base64 = data_json.get("img", None)
291 if not img_base64:
292 raise ValueError("Missing 'img' parameter")
293 if not isinstance(img_base64, str):
294 raise ValueError("'img' parameter must be a base64-encoded string")
295 img = base64_to_img(img_base64)
297 prompt = data_json.get("prompt", None)
298 if prompt is None:
299 raise ValueError("Missing 'prompt' parameter")
300 neg_prompt = data_json.get("neg_prompt", "")
302 enable_tts = bool(data_json.get("enable_tts", False))
304 audio_path: Optional[str] = None
305 tts_prompt_audio_path: Optional[str] = None
307 if enable_tts:
308 tts_prompt_audio_b64 = data_json.get("tts_prompt_audio", None)
309 if tts_prompt_audio_b64:
310 if not isinstance(tts_prompt_audio_b64, str):
311 raise ValueError("'tts_prompt_audio' must be a base64-encoded string")
312 tts_prompt_audio_path = await _base64_to_audio_file(tts_prompt_audio_b64)
313 else:
314 audio_b64 = data_json.get("audio", None)
315 if not audio_b64:
316 raise ValueError("Missing 'audio' parameter (or set enable_tts=true)")
317 if not isinstance(audio_b64, str):
318 raise ValueError("'audio' parameter must be a base64-encoded string")
319 audio_path_dest: Optional[str] = None
320 if job_id:
321 audio_path_dest = f"/tmp/{job_id}.wav"
322 audio_path = await _base64_to_audio_file(audio_b64, audio_path=audio_path_dest)
324 return {
325 "task": self.model_name,
326 "args": {
327 "job_id": job_id,
328 "img": img,
329 "prompt": prompt,
330 "neg_prompt": neg_prompt,
331 "max_area": int(data_json.get("max_area", 1024 * 704)),
332 "sampling_steps": int(data_json.get("sampling_steps", 40)),
333 "audio_path": audio_path,
334 "enable_tts": enable_tts,
335 "tts_prompt_audio": tts_prompt_audio_path,
336 "tts_prompt_text": data_json.get("tts_prompt_text", None),
337 "tts_text": data_json.get("tts_text", None),
338 "num_clip": int(data_json["num_clip"]) if data_json.get("num_clip") is not None else None,
339 "infer_frames": int(data_json.get("infer_frames", 80)),
340 "output_type": data_json.get("output_type", "tensor"),
341 }
342 }