wan_audio_runner.py 39.8 KB
Newer Older
wangshankun's avatar
wangshankun committed
1
import gc
2
import io
3
import json
PengGao's avatar
PengGao committed
4
import os
sandy's avatar
sandy committed
5
import warnings
PengGao's avatar
PengGao committed
6
from dataclasses import dataclass
7
from typing import Dict, List, Optional, Tuple, Union
PengGao's avatar
PengGao committed
8

wangshankun's avatar
wangshankun committed
9
10
import numpy as np
import torch
11
import torch.distributed as dist
sandy's avatar
sandy committed
12
import torch.nn.functional as F
gushiqiao's avatar
gushiqiao committed
13
import torchaudio as ta
helloyongyang's avatar
helloyongyang committed
14
import torchvision.transforms.functional as TF
15
from PIL import Image, ImageCms, ImageOps
gushiqiao's avatar
gushiqiao committed
16
from einops import rearrange
PengGao's avatar
PengGao committed
17
from loguru import logger
gushiqiao's avatar
gushiqiao committed
18
19
from torchvision.transforms import InterpolationMode
from torchvision.transforms.functional import resize
20

LiangLiu's avatar
LiangLiu committed
21
22
from lightx2v.deploy.common.va_reader import VAReader
from lightx2v.deploy.common.va_recorder import VARecorder
LiangLiu's avatar
LiangLiu committed
23
from lightx2v.deploy.common.va_recorder_x264 import X264VARecorder
24
from lightx2v.models.input_encoders.hf.seko_audio.audio_adapter import AudioAdapter
helloyongyang's avatar
helloyongyang committed
25
from lightx2v.models.input_encoders.hf.seko_audio.audio_encoder import SekoAudioEncoderModel
26
from lightx2v.models.networks.wan.audio_model import WanAudioModel
PengGao's avatar
PengGao committed
27
from lightx2v.models.networks.wan.lora_adapter import WanLoraWrapper
28
from lightx2v.models.runners.wan.wan_runner import WanRunner
29
from lightx2v.models.schedulers.wan.audio.scheduler import EulerScheduler
sandy's avatar
sandy committed
30
from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE
yihuiwen's avatar
yihuiwen committed
31
from lightx2v.server.metrics import monitor_cli
32
from lightx2v.utils.envs import *
33
from lightx2v.utils.profiler import *
PengGao's avatar
PengGao committed
34
from lightx2v.utils.registry_factory import RUNNER_REGISTER
LiangLiu's avatar
LiangLiu committed
35
from lightx2v.utils.utils import find_torch_model_path, load_weights, vae_to_comfyui_image_inplace
36

sandy's avatar
sandy committed
37
38
39
warnings.filterwarnings("ignore", category=UserWarning, module="torchaudio")
warnings.filterwarnings("ignore", category=UserWarning, module="torchvision.io")

wangshankun's avatar
wangshankun committed
40

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def get_optimal_patched_size_with_sp(patched_h, patched_w, sp_size):
    assert sp_size > 0 and (sp_size & (sp_size - 1)) == 0, "sp_size must be a power of 2"

    h_ratio, w_ratio = 1, 1
    while sp_size != 1:
        sp_size //= 2
        if patched_h % 2 == 0:
            patched_h //= 2
            h_ratio *= 2
        elif patched_w % 2 == 0:
            patched_w //= 2
            w_ratio *= 2
        else:
            if patched_h > patched_w:
                patched_h //= 2
56
57
                h_ratio *= 2
            else:
58
                patched_w //= 2
59
                w_ratio *= 2
60
    return patched_h * h_ratio, patched_w * w_ratio
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81


def get_crop_bbox(ori_h, ori_w, tgt_h, tgt_w):
    tgt_ar = tgt_h / tgt_w
    ori_ar = ori_h / ori_w
    if abs(ori_ar - tgt_ar) < 0.01:
        return 0, ori_h, 0, ori_w
    if ori_ar > tgt_ar:
        crop_h = int(tgt_ar * ori_w)
        y0 = (ori_h - crop_h) // 2
        y1 = y0 + crop_h
        return y0, y1, 0, ori_w
    else:
        crop_w = int(ori_h / tgt_ar)
        x0 = (ori_w - crop_w) // 2
        x1 = x0 + crop_w
        return 0, ori_h, x0, x1


def isotropic_crop_resize(frames: torch.Tensor, size: tuple):
    """
82
    frames: (C, H, W) or (T, C, H, W) or (N, C, H, W)
83
84
    size: (H, W)
    """
85
86
87
88
89
90
91
    original_shape = frames.shape

    if len(frames.shape) == 3:
        frames = frames.unsqueeze(0)
    elif len(frames.shape) == 4 and frames.shape[0] > 1:
        pass

92
93
94
95
    ori_h, ori_w = frames.shape[2:]
    h, w = size
    y0, y1, x0, x1 = get_crop_bbox(ori_h, ori_w, h, w)
    cropped_frames = frames[:, :, y0:y1, x0:x1]
96
    resized_frames = resize(cropped_frames, [h, w], InterpolationMode.BICUBIC, antialias=True)
97
98
99
100

    if len(original_shape) == 3:
        resized_frames = resized_frames.squeeze(0)

101
102
103
    return resized_frames


104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def fixed_shape_resize(img, target_height, target_width):
    orig_height, orig_width = img.shape[-2:]

    target_ratio = target_height / target_width
    orig_ratio = orig_height / orig_width

    if orig_ratio > target_ratio:
        crop_width = orig_width
        crop_height = int(crop_width * target_ratio)
    else:
        crop_height = orig_height
        crop_width = int(crop_height / target_ratio)

    cropped_img = TF.center_crop(img, [crop_height, crop_width])

    resized_img = TF.resize(cropped_img, [target_height, target_width], antialias=True)

    h, w = resized_img.shape[-2:]
    return resized_img, h, w


125
def resize_image(img, resize_mode="adaptive", bucket_shape=None, fixed_area=None, fixed_shape=None):
126
    assert resize_mode in ["adaptive", "keep_ratio_fixed_area", "fixed_min_area", "fixed_max_area", "fixed_shape", "fixed_min_side"]
127
128
129
130
131

    if resize_mode == "fixed_shape":
        assert fixed_shape is not None
        logger.info(f"[wan_audio] fixed_shape_resize fixed_height: {fixed_shape[0]}, fixed_width: {fixed_shape[1]}")
        return fixed_shape_resize(img, fixed_shape[0], fixed_shape[1])
132

133
134
135
136
137
138
139
140
141
142
143
    if bucket_shape is not None:
        """
        "adaptive_shape": {
            "0.667": [[480, 832], [544, 960], [720, 1280]],
            "1.500": [[832, 480], [960, 544], [1280, 720]],
            "1.000": [[480, 480], [576, 576], [704, 704], [960, 960]]
        }
        """
        bucket_config = {}
        for ratio, resolutions in bucket_shape.items():
            bucket_config[float(ratio)] = np.array(resolutions, dtype=np.int64)
144
        # logger.info(f"[wan_audio] use custom bucket_shape: {bucket_config}")
145
146
147
148
149
150
    else:
        bucket_config = {
            0.667: np.array([[480, 832], [544, 960], [720, 1280]], dtype=np.int64),
            1.500: np.array([[832, 480], [960, 544], [1280, 720]], dtype=np.int64),
            1.000: np.array([[480, 480], [576, 576], [704, 704], [960, 960]], dtype=np.int64),
        }
151
        # logger.info(f"[wan_audio] use default bucket_shape: {bucket_config}")
152

153
154
155
    ori_height = img.shape[-2]
    ori_weight = img.shape[-1]
    ori_ratio = ori_height / ori_weight
156
157
158
159
160
161
162
163
164
165
166

    if resize_mode == "adaptive":
        aspect_ratios = np.array(np.array(list(bucket_config.keys())))
        closet_aspect_idx = np.argmin(np.abs(aspect_ratios - ori_ratio))
        closet_ratio = aspect_ratios[closet_aspect_idx]
        if ori_ratio < 1.0:
            target_h, target_w = 480, 832
        elif ori_ratio == 1.0:
            target_h, target_w = 480, 480
        else:
            target_h, target_w = 832, 480
167
        for resolution in bucket_config[closet_ratio]:
168
169
170
171
172
173
174
175
176
177
178
            if ori_height * ori_weight >= resolution[0] * resolution[1]:
                target_h, target_w = resolution
    elif resize_mode == "keep_ratio_fixed_area":
        assert fixed_area in ["480p", "720p"], f"fixed_area must be in ['480p', '720p'], but got {fixed_area}, please set fixed_area in config."
        fixed_area = 480 * 832 if fixed_area == "480p" else 720 * 1280
        target_h = round(np.sqrt(fixed_area * ori_ratio))
        target_w = round(np.sqrt(fixed_area / ori_ratio))
    elif resize_mode == "fixed_min_area":
        aspect_ratios = np.array(np.array(list(bucket_config.keys())))
        closet_aspect_idx = np.argmin(np.abs(aspect_ratios - ori_ratio))
        closet_ratio = aspect_ratios[closet_aspect_idx]
179
        target_h, target_w = bucket_config[closet_ratio][0]
180
181
182
183
184
185
186
187
188
189
    elif resize_mode == "fixed_min_side":
        assert fixed_area in ["480p", "720p"], f"fixed_min_side mode requires fixed_area to be '480p' or '720p', got {fixed_area}"

        min_side = 720 if fixed_area == "720p" else 480
        if ori_ratio < 1.0:
            target_h = min_side
            target_w = round(target_h / ori_ratio)
        else:
            target_w = min_side
            target_h = round(target_w * ori_ratio)
190
191
192
193
    elif resize_mode == "fixed_max_area":
        aspect_ratios = np.array(np.array(list(bucket_config.keys())))
        closet_aspect_idx = np.argmin(np.abs(aspect_ratios - ori_ratio))
        closet_ratio = aspect_ratios[closet_aspect_idx]
194
        target_h, target_w = bucket_config[closet_ratio][-1]
195

196
197
198
199
    cropped_img = isotropic_crop_resize(img, (target_h, target_w))
    return cropped_img, target_h, target_w


200
201
202
203
@dataclass
class AudioSegment:
    """Data class for audio segment information"""

sandy's avatar
sandy committed
204
    audio_array: torch.Tensor
205
206
207
208
    start_frame: int
    end_frame: int


209
class FramePreprocessorTorchVersion:
210
211
212
213
214
215
216
    """Handles frame preprocessing including noise and masking"""

    def __init__(self, noise_mean: float = -3.0, noise_std: float = 0.5, mask_rate: float = 0.1):
        self.noise_mean = noise_mean
        self.noise_std = noise_std
        self.mask_rate = mask_rate

217
    def add_noise(self, frames: torch.Tensor, generator: Optional[torch.Generator] = None) -> torch.Tensor:
218
219
        """Add noise to frames"""

220
        device = frames.device
221
222
        shape = frames.shape
        bs = 1 if len(shape) == 4 else shape[0]
223
224
225
226
227
228
229
230
231
232

        # Generate sigma values on the same device
        sigma = torch.normal(mean=self.noise_mean, std=self.noise_std, size=(bs,), device=device, generator=generator)
        sigma = torch.exp(sigma)

        for _ in range(1, len(shape)):
            sigma = sigma.unsqueeze(-1)

        # Generate noise on the same device
        noise = torch.randn(*shape, device=device, generator=generator) * sigma
233
234
        return frames + noise

235
    def add_mask(self, frames: torch.Tensor, generator: Optional[torch.Generator] = None) -> torch.Tensor:
236
237
        """Add mask to frames"""

238
        device = frames.device
239
        h, w = frames.shape[-2:]
240
241
242

        # Generate mask on the same device
        mask = torch.rand(h, w, device=device, generator=generator) > self.mask_rate
243
244
245
246
        return frames * mask

    def process_prev_frames(self, frames: torch.Tensor) -> torch.Tensor:
        """Process previous frames with noise and masking"""
247
248
249
        frames = self.add_noise(frames, torch.Generator(device=frames.device))
        frames = self.add_mask(frames, torch.Generator(device=frames.device))
        return frames
250
251
252
253
254
255
256
257


class AudioProcessor:
    """Handles audio loading and segmentation"""

    def __init__(self, audio_sr: int = 16000, target_fps: int = 16):
        self.audio_sr = audio_sr
        self.target_fps = target_fps
sandy's avatar
sandy committed
258
        self.audio_frame_rate = audio_sr // target_fps
259

sandy's avatar
sandy committed
260
    def load_audio(self, audio_path: str):
261
262
        audio_array, ori_sr = ta.load(audio_path)
        audio_array = ta.functional.resample(audio_array.mean(0), orig_freq=ori_sr, new_freq=self.audio_sr)
sandy's avatar
sandy committed
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
        return audio_array

    def load_multi_person_audio(self, audio_paths: List[str]):
        audio_arrays = []
        max_len = 0

        for audio_path in audio_paths:
            audio_array = self.load_audio(audio_path)
            audio_arrays.append(audio_array)
            max_len = max(max_len, audio_array.numel())

        num_files = len(audio_arrays)
        padded = torch.zeros(num_files, max_len, dtype=torch.float32)

        for i, arr in enumerate(audio_arrays):
            length = arr.numel()
            padded[i, :length] = arr

        return padded
282
283
284

    def get_audio_range(self, start_frame: int, end_frame: int) -> Tuple[int, int]:
        """Calculate audio range for given frame range"""
sandy's avatar
sandy committed
285
        return round(start_frame * self.audio_frame_rate), round(end_frame * self.audio_frame_rate)
286

sandy's avatar
sandy committed
287
288
289
290
291
    def segment_audio(self, audio_array: torch.Tensor, expected_frames: int, max_num_frames: int, prev_frame_length: int = 5) -> List[AudioSegment]:
        """
        Segment audio based on frame requirements
        audio_array is (N, T) tensor
        """
292
        segments = []
sandy's avatar
sandy committed
293
        segments_idx = self.init_segments_idx(expected_frames, max_num_frames, prev_frame_length)
294

sandy's avatar
sandy committed
295
        audio_start, audio_end = self.get_audio_range(0, expected_frames)
sandy's avatar
sandy committed
296
        audio_array_ori = audio_array[:, audio_start:audio_end]
297

sandy's avatar
sandy committed
298
299
        for idx, (start_idx, end_idx) in enumerate(segments_idx):
            audio_start, audio_end = self.get_audio_range(start_idx, end_idx)
sandy's avatar
sandy committed
300
            audio_array = audio_array_ori[:, audio_start:audio_end]
301

sandy's avatar
sandy committed
302
303
            if idx < len(segments_idx) - 1:
                end_idx = segments_idx[idx + 1][0]
sandy's avatar
sandy committed
304
305
306
307
308
            else:  # for last segments
                if audio_array.shape[1] < audio_end - audio_start:
                    padding_len = audio_end - audio_start - audio_array.shape[1]
                    audio_array = F.pad(audio_array, (0, padding_len))
                    # Adjust end_idx to account for the frames added by padding
sandy's avatar
sandy committed
309
                    end_idx = end_idx - padding_len // self.audio_frame_rate
310

sandy's avatar
sandy committed
311
312
            segments.append(AudioSegment(audio_array, start_idx, end_idx))
        del audio_array, audio_array_ori
313
314
        return segments

sandy's avatar
sandy committed
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
    def init_segments_idx(self, total_frame: int, clip_frame: int = 81, overlap_frame: int = 5) -> list[tuple[int, int, int]]:
        """Initialize segment indices with overlap"""
        start_end_list = []
        min_frame = clip_frame
        for start in range(0, total_frame, clip_frame - overlap_frame):
            is_last = start + clip_frame >= total_frame
            end = min(start + clip_frame, total_frame)
            if end - start < min_frame:
                end = start + min_frame
            if ((end - start) - 1) % 4 != 0:
                end = start + (((end - start) - 1) // 4) * 4 + 1
            start_end_list.append((start, end))
            if is_last:
                break
        return start_end_list

331

332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
def load_image(image: Union[str, Image.Image], to_rgb: bool = True) -> Image.Image:
    _image = image
    if isinstance(image, str):
        if os.path.isfile(image):
            _image = Image.open(image)
        else:
            raise ValueError(f"Incorrect path. {image} is not a valid path.")
    # orientation transpose
    _image = ImageOps.exif_transpose(_image)
    # convert color space to sRGB
    icc_profile = _image.info.get("icc_profile")
    if icc_profile:
        srgb_profile = ImageCms.createProfile("sRGB")
        input_profile = ImageCms.ImageCmsProfile(io.BytesIO(icc_profile))
        _image = ImageCms.profileToProfile(_image, input_profile, srgb_profile)
    # convert to "RGB"
    if to_rgb:
        _image = _image.convert("RGB")

    return _image


Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
354
@RUNNER_REGISTER("seko_talk")
helloyongyang's avatar
helloyongyang committed
355
356
357
class WanAudioRunner(WanRunner):  # type:ignore
    def __init__(self, config):
        super().__init__(config)
358
        self.prev_frame_length = self.config.get("prev_frame_length", 5)
359
        self.frame_preprocessor = FramePreprocessorTorchVersion()
helloyongyang's avatar
helloyongyang committed
360
361
362

    def init_scheduler(self):
        """Initialize consistency model scheduler"""
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
363
        self.scheduler = EulerScheduler(self.config)
helloyongyang's avatar
helloyongyang committed
364

365
    def read_audio_input(self, audio_path):
sandy's avatar
sandy committed
366
        """Read audio input - handles both single and multi-person scenarios"""
helloyongyang's avatar
helloyongyang committed
367
368
369
        audio_sr = self.config.get("audio_sr", 16000)
        target_fps = self.config.get("target_fps", 16)
        self._audio_processor = AudioProcessor(audio_sr, target_fps)
sandy's avatar
sandy committed
370

LiangLiu's avatar
LiangLiu committed
371
372
373
        if not isinstance(audio_path, str):
            return [], 0, None, 0

sandy's avatar
sandy committed
374
        # Get audio files from person objects or legacy format
375
        audio_files, mask_files = self.get_audio_files_from_audio_path(audio_path)
helloyongyang's avatar
helloyongyang committed
376

sandy's avatar
sandy committed
377
378
379
380
381
382
383
384
385
        # Load audio based on single or multi-person mode
        if len(audio_files) == 1:
            audio_array = self._audio_processor.load_audio(audio_files[0])
            audio_array = audio_array.unsqueeze(0)  # Add batch dimension for consistency
        else:
            audio_array = self._audio_processor.load_multi_person_audio(audio_files)

        video_duration = self.config.get("video_duration", 5)
        audio_len = int(audio_array.shape[1] / audio_sr * target_fps)
yihuiwen's avatar
yihuiwen committed
386
387
388
        if GET_RECORDER_MODE():
            monitor_cli.lightx2v_input_audio_len.observe(audio_len)

helloyongyang's avatar
helloyongyang committed
389
        expected_frames = min(max(1, int(video_duration * target_fps)), audio_len)
gushiqiao's avatar
gushiqiao committed
390
391
        if expected_frames < int(video_duration * target_fps):
            logger.warning(f"Input video duration is greater than actual audio duration, using audio duration instead: audio_duration={audio_len / target_fps}, video_duration={video_duration}")
helloyongyang's avatar
helloyongyang committed
392
393

        # Segment audio
394
        audio_segments = self._audio_processor.segment_audio(audio_array, expected_frames, self.config.get("target_video_length", 81), self.prev_frame_length)
helloyongyang's avatar
helloyongyang committed
395

396
397
398
399
400
401
        # Mask latent for multi-person s2v
        if mask_files is not None:
            mask_latents = [self.process_single_mask(mask_file) for mask_file in mask_files]
            mask_latents = torch.cat(mask_latents, dim=0)
        else:
            mask_latents = None
sandy's avatar
sandy committed
402

403
        return audio_segments, expected_frames, mask_latents, len(audio_files)
sandy's avatar
sandy committed
404

405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
    def get_audio_files_from_audio_path(self, audio_path):
        if os.path.isdir(audio_path):
            audio_files = []
            mask_files = []
            logger.info(f"audio_path is a directory, loading config.json from {audio_path}")
            audio_config_path = os.path.join(audio_path, "config.json")
            assert os.path.exists(audio_config_path), "config.json not found in audio_path"
            with open(audio_config_path, "r") as f:
                audio_config = json.load(f)
            for talk_object in audio_config["talk_objects"]:
                audio_files.append(os.path.join(audio_path, talk_object["audio"]))
                mask_files.append(os.path.join(audio_path, talk_object["mask"]))
        else:
            logger.info(f"audio_path is a file without mask: {audio_path}")
            audio_files = [audio_path]
            mask_files = None
sandy's avatar
sandy committed
421

422
        return audio_files, mask_files
sandy's avatar
sandy committed
423

424
    def process_single_mask(self, mask_file):
425
        mask_img = load_image(mask_file)
sandy's avatar
sandy committed
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
        mask_img = TF.to_tensor(mask_img).sub_(0.5).div_(0.5).unsqueeze(0).cuda()

        if mask_img.shape[1] == 3:  # If it is an RGB three-channel image
            mask_img = mask_img[:, :1]  # Only take the first channel

        mask_img, h, w = resize_image(
            mask_img,
            resize_mode=self.config.get("resize_mode", "adaptive"),
            bucket_shape=self.config.get("bucket_shape", None),
            fixed_area=self.config.get("fixed_area", None),
            fixed_shape=self.config.get("fixed_shape", None),
        )

        mask_latent = torch.nn.functional.interpolate(
            mask_img,  # (1, 1, H, W)
            size=(h // 16, w // 16),
            mode="bicubic",
        )

        mask_latent = (mask_latent > 0).to(torch.int8)
        return mask_latent
helloyongyang's avatar
helloyongyang committed
447
448

    def read_image_input(self, img_path):
LiangLiu's avatar
LiangLiu committed
449
450
451
        if isinstance(img_path, Image.Image):
            ref_img = img_path
        else:
452
            ref_img = load_image(img_path)
Kane's avatar
Kane committed
453
        ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(0).to(self.init_device)
helloyongyang's avatar
helloyongyang committed
454

455
456
457
458
459
460
461
        ref_img, h, w = resize_image(
            ref_img,
            resize_mode=self.config.get("resize_mode", "adaptive"),
            bucket_shape=self.config.get("bucket_shape", None),
            fixed_area=self.config.get("fixed_area", None),
            fixed_shape=self.config.get("fixed_shape", None),
        )
462
        logger.info(f"[wan_audio] resize_image target_h: {h}, target_w: {w}")
463
464
        patched_h = h // self.config["vae_stride"][1] // self.config["patch_size"][1]
        patched_w = w // self.config["vae_stride"][2] // self.config["patch_size"][2]
helloyongyang's avatar
helloyongyang committed
465
466
467

        patched_h, patched_w = get_optimal_patched_size_with_sp(patched_h, patched_w, 1)

468
469
        latent_h = patched_h * self.config["patch_size"][1]
        latent_w = patched_w * self.config["patch_size"][2]
helloyongyang's avatar
helloyongyang committed
470

471
472
        latent_shape = self.get_latent_shape_with_lat_hw(latent_h, latent_w)
        target_shape = [latent_h * self.config["vae_stride"][1], latent_w * self.config["vae_stride"][2]]
helloyongyang's avatar
helloyongyang committed
473

474
        logger.info(f"[wan_audio] target_h: {target_shape[0]}, target_w: {target_shape[1]}, latent_h: {latent_h}, latent_w: {latent_w}")
helloyongyang's avatar
helloyongyang committed
475

476
477
        ref_img = torch.nn.functional.interpolate(ref_img, size=(target_shape[0], target_shape[1]), mode="bicubic")
        return ref_img, latent_shape, target_shape
helloyongyang's avatar
helloyongyang committed
478

yihuiwen's avatar
yihuiwen committed
479
480
481
482
483
484
    @ProfilingContext4DebugL1(
        "Run Image Encoder",
        recorder_mode=GET_RECORDER_MODE(),
        metrics_func=monitor_cli.lightx2v_run_img_encode_duration,
        metrics_labels=["WanAudioRunner"],
    )
helloyongyang's avatar
helloyongyang committed
485
    def run_image_encoder(self, first_frame, last_frame=None):
486
487
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            self.image_encoder = self.load_image_encoder()
helloyongyang's avatar
helloyongyang committed
488
        clip_encoder_out = self.image_encoder.visual([first_frame]).squeeze(0).to(GET_DTYPE()) if self.config.get("use_image_encoder", True) else None
489
490
491
492
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            del self.image_encoder
            torch.cuda.empty_cache()
            gc.collect()
helloyongyang's avatar
helloyongyang committed
493
494
        return clip_encoder_out

yihuiwen's avatar
yihuiwen committed
495
496
497
    @ProfilingContext4DebugL1(
        "Run VAE Encoder",
        recorder_mode=GET_RECORDER_MODE(),
498
        metrics_func=monitor_cli.lightx2v_run_vae_encoder_image_duration,
yihuiwen's avatar
yihuiwen committed
499
500
        metrics_labels=["WanAudioRunner"],
    )
helloyongyang's avatar
helloyongyang committed
501
    def run_vae_encoder(self, img):
502
503
504
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            self.vae_encoder = self.load_vae_encoder()

helloyongyang's avatar
helloyongyang committed
505
        img = rearrange(img, "1 C H W -> 1 C 1 H W")
506
        vae_encoder_out = self.vae_encoder.encode(img.to(GET_DTYPE()))
sandy's avatar
sandy committed
507

508
509
510
511
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            del self.vae_encoder
            torch.cuda.empty_cache()
            gc.collect()
helloyongyang's avatar
helloyongyang committed
512
513
        return vae_encoder_out

514
    @ProfilingContext4DebugL2("Run Encoders")
515
516
517
518
    def _run_input_encoder_local_s2v(self):
        img, latent_shape, target_shape = self.read_image_input(self.input_info.image_path)
        self.input_info.latent_shape = latent_shape  # Important: set latent_shape in input_info
        self.input_info.target_shape = target_shape  # Important: set target_shape in input_info
helloyongyang's avatar
helloyongyang committed
519
520
        clip_encoder_out = self.run_image_encoder(img) if self.config.get("use_image_encoder", True) else None
        vae_encode_out = self.run_vae_encoder(img)
sandy's avatar
sandy committed
521

522
523
524
525
        audio_segments, expected_frames, person_mask_latens, audio_num = self.read_audio_input(self.input_info.audio_path)
        self.input_info.audio_num = audio_num
        self.input_info.with_mask = person_mask_latens is not None
        text_encoder_output = self.run_text_encoder(self.input_info)
helloyongyang's avatar
helloyongyang committed
526
527
528
529
530
531
532
533
534
535
        torch.cuda.empty_cache()
        gc.collect()
        return {
            "text_encoder_output": text_encoder_output,
            "image_encoder_output": {
                "clip_encoder_out": clip_encoder_out,
                "vae_encoder_out": vae_encode_out,
            },
            "audio_segments": audio_segments,
            "expected_frames": expected_frames,
sandy's avatar
sandy committed
536
            "person_mask_latens": person_mask_latens,
helloyongyang's avatar
helloyongyang committed
537
        }
538
539
540

    def prepare_prev_latents(self, prev_video: Optional[torch.Tensor], prev_frame_length: int) -> Optional[Dict[str, torch.Tensor]]:
        """Prepare previous latents for conditioning"""
Kane's avatar
Kane committed
541
        device = self.init_device
542
        dtype = GET_DTYPE()
543

544
545
        tgt_h, tgt_w = self.input_info.target_shape[0], self.input_info.target_shape[1]
        prev_frames = torch.zeros((1, 3, self.config["target_video_length"], tgt_h, tgt_w), device=device)
546

547
548
549
        if prev_video is not None:
            # Extract and process last frames
            last_frames = prev_video[:, :, -prev_frame_length:].clone().to(device)
550
            if self.config["model_cls"] != "wan2.2_audio":
sandy's avatar
sandy committed
551
                last_frames = self.frame_preprocessor.process_prev_frames(last_frames)
552
            prev_frames[:, :, :prev_frame_length] = last_frames
sandy's avatar
sandy committed
553
554
555
            prev_len = (prev_frame_length - 1) // 4 + 1
        else:
            prev_len = 0
556

557
558
559
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            self.vae_encoder = self.load_vae_encoder()

560
        _, nframe, height, width = self.model.scheduler.latents.shape
561
562
563
564
565
566
        with ProfilingContext4DebugL1(
            "vae_encoder in init run segment",
            recorder_mode=GET_RECORDER_MODE(),
            metrics_func=monitor_cli.lightx2v_run_vae_encoder_pre_latent_duration,
            metrics_labels=["WanAudioRunner"],
        ):
567
            if self.config["model_cls"] == "wan2.2_audio":
568
569
570
571
572
                if prev_video is not None:
                    prev_latents = self.vae_encoder.encode(prev_frames.to(dtype))
                else:
                    prev_latents = None
                prev_mask = self.model.scheduler.mask
573
            else:
574
                prev_latents = self.vae_encoder.encode(prev_frames.to(dtype))
575

576
577
            frames_n = (nframe - 1) * 4 + 1
            prev_mask = torch.ones((1, frames_n, height, width), device=device, dtype=dtype)
578
579
            prev_frame_len = max((prev_len - 1) * 4 + 1, 0)
            prev_mask[:, prev_frame_len:] = 0
580
            prev_mask = self._wan_mask_rearrange(prev_mask)
helloyongyang's avatar
fix ci  
helloyongyang committed
581

sandy's avatar
sandy committed
582
583
        if prev_latents is not None:
            if prev_latents.shape[-2:] != (height, width):
584
                logger.warning(f"Size mismatch: prev_latents {prev_latents.shape} vs scheduler latents (H={height}, W={width}). Config tgt_h={tgt_h}, tgt_w={tgt_w}")
sandy's avatar
sandy committed
585
                prev_latents = torch.nn.functional.interpolate(prev_latents, size=(height, width), mode="bilinear", align_corners=False)
586

587
588
589
590
591
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            del self.vae_encoder
            torch.cuda.empty_cache()
            gc.collect()

sandy's avatar
sandy committed
592
        return {"prev_latents": prev_latents, "prev_mask": prev_mask, "prev_len": prev_len}
593
594
595
596
597
598
599
600
601
602
603

    def _wan_mask_rearrange(self, mask: torch.Tensor) -> torch.Tensor:
        """Rearrange mask for WAN model"""
        if mask.ndim == 3:
            mask = mask[None]
        assert mask.ndim == 4
        _, t, h, w = mask.shape
        assert t == ((t - 1) // 4 * 4 + 1)
        mask_first_frame = torch.repeat_interleave(mask[:, 0:1], repeats=4, dim=1)
        mask = torch.concat([mask_first_frame, mask[:, 1:]], dim=1)
        mask = mask.view(mask.shape[1] // 4, 4, h, w)
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
604
        return mask.transpose(0, 1).contiguous()
605

helloyongyang's avatar
helloyongyang committed
606
607
    def get_video_segment_num(self):
        self.video_segment_num = len(self.inputs["audio_segments"])
wangshankun's avatar
wangshankun committed
608

helloyongyang's avatar
helloyongyang committed
609
610
    def init_run(self):
        super().init_run()
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
611
        self.scheduler.set_audio_adapter(self.audio_adapter)
helloyongyang's avatar
helloyongyang committed
612
        self.prev_video = None
613
614
        if self.input_info.return_result_tensor:
            self.gen_video_final = torch.zeros((self.inputs["expected_frames"], self.input_info.target_shape[0], self.input_info.target_shape[1], 3), dtype=torch.float32, device="cpu")
sandy's avatar
sandy committed
615
            self.cut_audio_final = torch.zeros((self.inputs["expected_frames"] * self._audio_processor.audio_frame_rate), dtype=torch.float32, device="cpu")
LiangLiu's avatar
LiangLiu committed
616
617
        else:
            self.gen_video_final = None
sandy's avatar
sandy committed
618
            self.cut_audio_final = None
wangshankun's avatar
wangshankun committed
619

620
621
622
623
624
625
    @ProfilingContext4DebugL1(
        "Init run segment",
        recorder_mode=GET_RECORDER_MODE(),
        metrics_func=monitor_cli.lightx2v_run_init_run_segment_duration,
        metrics_labels=["WanAudioRunner"],
    )
LiangLiu's avatar
LiangLiu committed
626
    def init_run_segment(self, segment_idx, audio_array=None):
helloyongyang's avatar
helloyongyang committed
627
        self.segment_idx = segment_idx
LiangLiu's avatar
LiangLiu committed
628
        if audio_array is not None:
LiangLiu's avatar
LiangLiu committed
629
630
631
            end_idx = audio_array.shape[0] // self._audio_processor.audio_frame_rate - self.prev_frame_length
            audio_tensor = torch.Tensor(audio_array).float().unsqueeze(0)
            self.segment = AudioSegment(audio_tensor, 0, end_idx)
LiangLiu's avatar
LiangLiu committed
632
633
        else:
            self.segment = self.inputs["audio_segments"][segment_idx]
wangshankun's avatar
wangshankun committed
634

635
636
        self.input_info.seed = self.input_info.seed + segment_idx
        torch.manual_seed(self.input_info.seed)
637
        # logger.info(f"Processing segment {segment_idx + 1}/{self.video_segment_num}, seed: {self.config.seed}")
wangshankun's avatar
wangshankun committed
638

639
640
641
        if (self.config.get("lazy_load", False) or self.config.get("unload_modules", False)) and not hasattr(self, "audio_encoder"):
            self.audio_encoder = self.load_audio_encoder()

sandy's avatar
sandy committed
642
643
644
645
646
647
        features_list = []
        for i in range(self.segment.audio_array.shape[0]):
            feat = self.audio_encoder.infer(self.segment.audio_array[i])
            feat = self.audio_adapter.forward_audio_proj(feat, self.model.scheduler.latents.shape[1])
            features_list.append(feat.squeeze(0))
        audio_features = torch.stack(features_list, dim=0)
PengGao's avatar
PengGao committed
648

helloyongyang's avatar
helloyongyang committed
649
        self.inputs["audio_encoder_output"] = audio_features
650
        self.inputs["previmg_encoder_output"] = self.prepare_prev_latents(self.prev_video, prev_frame_length=self.prev_frame_length)
wangshankun's avatar
wangshankun committed
651

helloyongyang's avatar
helloyongyang committed
652
653
        # Reset scheduler for non-first segments
        if segment_idx > 0:
654
            self.model.scheduler.reset(self.input_info.seed, self.input_info.latent_shape, self.inputs["previmg_encoder_output"])
wangshankun's avatar
wangshankun committed
655

656
657
658
659
660
661
    @ProfilingContext4DebugL1(
        "End run segment",
        recorder_mode=GET_RECORDER_MODE(),
        metrics_func=monitor_cli.lightx2v_run_end_run_segment_duration,
        metrics_labels=["WanAudioRunner"],
    )
662
    def end_run_segment(self, segment_idx):
helloyongyang's avatar
helloyongyang committed
663
        self.gen_video = torch.clamp(self.gen_video, -1, 1).to(torch.float)
sandy's avatar
sandy committed
664
        useful_length = self.segment.end_frame - self.segment.start_frame
LiangLiu's avatar
LiangLiu committed
665
        video_seg = self.gen_video[:, :, :useful_length].cpu()
sandy's avatar
sandy committed
666
667
        audio_seg = self.segment.audio_array[:, : useful_length * self._audio_processor.audio_frame_rate]
        audio_seg = audio_seg.sum(dim=0)  # Multiple audio tracks, mixed into one track
LiangLiu's avatar
LiangLiu committed
668
669
670
671
672
673
674
675
676
677
678
        video_seg = vae_to_comfyui_image_inplace(video_seg)

        # [Warning] Need check whether video segment interpolation works...
        if "video_frame_interpolation" in self.config and self.vfi_model is not None:
            target_fps = self.config["video_frame_interpolation"]["target_fps"]
            logger.info(f"Interpolating frames from {self.config.get('fps', 16)} to {target_fps}")
            video_seg = self.vfi_model.interpolate_frames(
                video_seg,
                source_fps=self.config.get("fps", 16),
                target_fps=target_fps,
            )
LiangLiu's avatar
LiangLiu committed
679

680
681
682
683
684
685
686
687
        if "video_super_resolution" in self.config and self.vsr_model is not None:
            logger.info(f"Applying video super resolution with scale {self.config['video_super_resolution']['scale']}")
            video_seg = self.vsr_model.super_resolve_frames(
                video_seg,
                seed=self.config["video_super_resolution"]["seed"],
                scale=self.config["video_super_resolution"]["scale"],
            )

LiangLiu's avatar
LiangLiu committed
688
689
        if self.va_recorder:
            self.va_recorder.pub_livestream(video_seg, audio_seg)
690
        elif self.input_info.return_result_tensor:
LiangLiu's avatar
LiangLiu committed
691
            self.gen_video_final[self.segment.start_frame : self.segment.end_frame].copy_(video_seg)
sandy's avatar
sandy committed
692
            self.cut_audio_final[self.segment.start_frame * self._audio_processor.audio_frame_rate : self.segment.end_frame * self._audio_processor.audio_frame_rate].copy_(audio_seg)
LiangLiu's avatar
LiangLiu committed
693

helloyongyang's avatar
helloyongyang committed
694
695
696
        # Update prev_video for next iteration
        self.prev_video = self.gen_video

LiangLiu's avatar
LiangLiu committed
697
        del video_seg, audio_seg
helloyongyang's avatar
helloyongyang committed
698
699
        torch.cuda.empty_cache()

LiangLiu's avatar
LiangLiu committed
700
701
702
703
704
705
706
707
708
    def get_rank_and_world_size(self):
        rank = 0
        world_size = 1
        if dist.is_initialized():
            rank = dist.get_rank()
            world_size = dist.get_world_size()
        return rank, world_size

    def init_va_recorder(self):
709
        output_video_path = self.input_info.save_result_path
LiangLiu's avatar
LiangLiu committed
710
711
        self.va_recorder = None
        if isinstance(output_video_path, dict):
LiangLiu's avatar
LiangLiu committed
712
713
714
715
716
717
718
719
            output_video_path = output_video_path["data"]
        logger.info(f"init va_recorder with output_video_path: {output_video_path}")
        rank, world_size = self.get_rank_and_world_size()
        if output_video_path and rank == world_size - 1:
            record_fps = self.config.get("target_fps", 16)
            audio_sr = self.config.get("audio_sr", 16000)
            if "video_frame_interpolation" in self.config and self.vfi_model is not None:
                record_fps = self.config["video_frame_interpolation"]["target_fps"]
LiangLiu's avatar
LiangLiu committed
720
721
722

            whip_shared_path = os.getenv("WHIP_SHARED_LIB", None)
            if whip_shared_path and output_video_path.startswith("http"):
LiangLiu's avatar
LiangLiu committed
723
                self.va_recorder = X264VARecorder(
LiangLiu's avatar
LiangLiu committed
724
725
726
727
728
729
730
731
732
733
734
                    whip_shared_path=whip_shared_path,
                    livestream_url=output_video_path,
                    fps=record_fps,
                    sample_rate=audio_sr,
                )
            else:
                self.va_recorder = VARecorder(
                    livestream_url=output_video_path,
                    fps=record_fps,
                    sample_rate=audio_sr,
                )
LiangLiu's avatar
LiangLiu committed
735
736

    def init_va_reader(self):
LiangLiu's avatar
LiangLiu committed
737
        audio_path = self.input_info.audio_path
LiangLiu's avatar
LiangLiu committed
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
        self.va_reader = None
        if isinstance(audio_path, dict):
            assert audio_path["type"] == "stream", f"unexcept audio_path: {audio_path}"
            rank, world_size = self.get_rank_and_world_size()
            target_fps = self.config.get("target_fps", 16)
            max_num_frames = self.config.get("target_video_length", 81)
            audio_sr = self.config.get("audio_sr", 16000)
            prev_frames = self.config.get("prev_frame_length", 5)
            self.va_reader = VAReader(
                rank=rank,
                world_size=world_size,
                stream_url=audio_path["data"],
                sample_rate=audio_sr,
                segment_duration=max_num_frames / target_fps,
                prev_duration=prev_frames / target_fps,
                target_rank=1,
            )

PengGao's avatar
PengGao committed
756
    def run_main(self):
LiangLiu's avatar
LiangLiu committed
757
758
759
760
761
762
        try:
            self.init_va_recorder()
            self.init_va_reader()
            logger.info(f"init va_recorder: {self.va_recorder} and va_reader: {self.va_reader}")

            if self.va_reader is None:
PengGao's avatar
PengGao committed
763
                return super().run_main()
LiangLiu's avatar
LiangLiu committed
764

LiangLiu's avatar
LiangLiu committed
765
            self.va_reader.start()
LiangLiu's avatar
LiangLiu committed
766
            rank, world_size = self.get_rank_and_world_size()
LiangLiu's avatar
LiangLiu committed
767
768
            if rank == world_size - 1:
                assert self.va_recorder is not None, "va_recorder is required for stream audio input for rank 2"
LiangLiu's avatar
LiangLiu committed
769
770
771
                self.va_recorder.start(self.input_info.target_shape[1], self.input_info.target_shape[0])
            if world_size > 1:
                dist.barrier()
LiangLiu's avatar
LiangLiu committed
772
773

            self.init_run()
LiangLiu's avatar
LiangLiu committed
774
            if self.config.get("compile", False):
775
                self.model.select_graph_for_compile(self.input_info)
LiangLiu's avatar
LiangLiu committed
776
777
778
779
780
781
782
783
            self.video_segment_num = "unlimited"

            fetch_timeout = self.va_reader.segment_duration + 1
            segment_idx = 0
            fail_count = 0
            max_fail_count = 10

            while True:
784
                with ProfilingContext4DebugL1(f"stream segment get audio segment {segment_idx}"):
LiangLiu's avatar
LiangLiu committed
785
786
787
788
789
790
791
792
793
                    self.check_stop()
                    audio_array = self.va_reader.get_audio_segment(timeout=fetch_timeout)
                    if audio_array is None:
                        fail_count += 1
                        logger.warning(f"Failed to get audio chunk {fail_count} times")
                        if fail_count > max_fail_count:
                            raise Exception(f"Failed to get audio chunk {fail_count} times, stop reader")
                        continue

794
                with ProfilingContext4DebugL1(f"stream segment end2end {segment_idx}"):
LiangLiu's avatar
LiangLiu committed
795
796
                    fail_count = 0
                    self.init_run_segment(segment_idx, audio_array)
PengGao's avatar
PengGao committed
797
                    latents = self.run_segment(segment_idx)
LiangLiu's avatar
LiangLiu committed
798
                    self.gen_video = self.run_vae_decoder(latents)
LiangLiu's avatar
LiangLiu committed
799
                    self.end_run_segment(segment_idx)
LiangLiu's avatar
LiangLiu committed
800
801
802
                    segment_idx += 1

        finally:
LiangLiu's avatar
LiangLiu committed
803
            if hasattr(self.model, "inputs"):
LiangLiu's avatar
LiangLiu committed
804
805
806
807
808
                self.end_run()
            if self.va_reader:
                self.va_reader.stop()
                self.va_reader = None
            if self.va_recorder:
LiangLiu's avatar
LiangLiu committed
809
                self.va_recorder.stop()
LiangLiu's avatar
LiangLiu committed
810
811
                self.va_recorder = None

812
    @ProfilingContext4DebugL1("Process after vae decoder")
813
814
    def process_images_after_vae_decoder(self):
        if self.input_info.return_result_tensor:
sandy's avatar
sandy committed
815
            audio_waveform = self.cut_audio_final.unsqueeze(0).unsqueeze(0)
LiangLiu's avatar
LiangLiu committed
816
817
818
            comfyui_audio = {"waveform": audio_waveform, "sample_rate": self._audio_processor.audio_sr}
            return {"video": self.gen_video_final, "audio": comfyui_audio}
        return {"video": None, "audio": None}
819

wangshankun's avatar
wangshankun committed
820
    def load_transformer(self):
821
        """Load transformer with LoRA support"""
822
823
        base_model = WanAudioModel(self.config["model_path"], self.config, self.init_device)
        if self.config.get("lora_configs") and self.config["lora_configs"]:
824
            assert not self.config.get("dit_quantized", False)
wangshankun's avatar
wangshankun committed
825
            lora_wrapper = WanLoraWrapper(base_model)
826
            for lora_config in self.config["lora_configs"]:
827
828
829
830
831
                lora_path = lora_config["path"]
                strength = lora_config.get("strength", 1.0)
                lora_name = lora_wrapper.load_lora(lora_path)
                lora_wrapper.apply_lora(lora_name, strength)
                logger.info(f"Loaded LoRA: {lora_name} with strength: {strength}")
wangshankun's avatar
wangshankun committed
832

wangshankun's avatar
wangshankun committed
833
834
        return base_model

helloyongyang's avatar
helloyongyang committed
835
    def load_audio_encoder(self):
gushiqiao's avatar
gushiqiao committed
836
        audio_encoder_path = self.config.get("audio_encoder_path", os.path.join(self.config["model_path"], "TencentGameMate-chinese-hubert-large"))
837
        audio_encoder_offload = self.config.get("audio_encoder_cpu_offload", self.config.get("cpu_offload", False))
Kane's avatar
Kane committed
838
        model = SekoAudioEncoderModel(audio_encoder_path, self.config["audio_sr"], audio_encoder_offload, device=self.config.get("run_device", "cuda"))
helloyongyang's avatar
helloyongyang committed
839
        return model
840

helloyongyang's avatar
helloyongyang committed
841
    def load_audio_adapter(self):
842
843
844
845
        audio_adapter_offload = self.config.get("audio_adapter_cpu_offload", self.config.get("cpu_offload", False))
        if audio_adapter_offload:
            device = torch.device("cpu")
        else:
Kane's avatar
Kane committed
846
            device = torch.device(self.config.get("run_device", "cuda"))
helloyongyang's avatar
helloyongyang committed
847
        audio_adapter = AudioAdapter(
sandy's avatar
sandy committed
848
            attention_head_dim=self.config["dim"] // self.config["num_heads"],
helloyongyang's avatar
helloyongyang committed
849
850
851
852
853
854
855
856
857
            num_attention_heads=self.config["num_heads"],
            base_num_layers=self.config["num_layers"],
            interval=1,
            audio_feature_dim=1024,
            time_freq_dim=256,
            projection_transformer_layers=4,
            mlp_dims=(1024, 1024, 32 * 1024),
            quantized=self.config.get("adapter_quantized", False),
            quant_scheme=self.config.get("adapter_quant_scheme", None),
858
            cpu_offload=audio_adapter_offload,
Kane's avatar
Kane committed
859
            device=self.config.get("run_device", "cuda"),
helloyongyang's avatar
helloyongyang committed
860
        )
861

862
        audio_adapter.to(device)
863
        load_from_rank0 = self.config.get("load_from_rank0", False)
864
        weights_dict = load_weights(self.config["adapter_model_path"], cpu_offload=audio_adapter_offload, remove_key="ca", load_from_rank0=load_from_rank0)
865
        audio_adapter.load_state_dict(weights_dict, strict=False)
helloyongyang's avatar
helloyongyang committed
866
        return audio_adapter.to(dtype=GET_DTYPE())
wangshankun's avatar
wangshankun committed
867

helloyongyang's avatar
helloyongyang committed
868
869
    def load_model(self):
        super().load_model()
870
871
872
        with ProfilingContext4DebugL2("Load audio encoder and adapter"):
            self.audio_encoder = self.load_audio_encoder()
            self.audio_adapter = self.load_audio_adapter()
wangshankun's avatar
wangshankun committed
873

874
875
876
877
878
879
880
881
    def get_latent_shape_with_lat_hw(self, latent_h, latent_w):
        latent_shape = [
            self.config.get("num_channels_latents", 16),
            (self.config["target_video_length"] - 1) // self.config["vae_stride"][0] + 1,
            latent_h,
            latent_w,
        ]
        return latent_shape
sandy's avatar
sandy committed
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896


@RUNNER_REGISTER("wan2.2_audio")
class Wan22AudioRunner(WanAudioRunner):
    def __init__(self, config):
        super().__init__(config)

    def load_vae_decoder(self):
        # offload config
        vae_offload = self.config.get("vae_cpu_offload", self.config.get("cpu_offload"))
        if vae_offload:
            vae_device = torch.device("cpu")
        else:
            vae_device = torch.device("cuda")
        vae_config = {
gushiqiao's avatar
gushiqiao committed
897
            "vae_path": find_torch_model_path(self.config, "vae_path", "Wan2.2_VAE.pth"),
sandy's avatar
sandy committed
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
            "device": vae_device,
            "cpu_offload": vae_offload,
            "offload_cache": self.config.get("vae_offload_cache", False),
        }
        vae_decoder = Wan2_2_VAE(**vae_config)
        return vae_decoder

    def load_vae_encoder(self):
        # offload config
        vae_offload = self.config.get("vae_cpu_offload", self.config.get("cpu_offload"))
        if vae_offload:
            vae_device = torch.device("cpu")
        else:
            vae_device = torch.device("cuda")
        vae_config = {
gushiqiao's avatar
gushiqiao committed
913
            "vae_path": find_torch_model_path(self.config, "vae_path", "Wan2.2_VAE.pth"),
sandy's avatar
sandy committed
914
915
916
917
            "device": vae_device,
            "cpu_offload": vae_offload,
            "offload_cache": self.config.get("vae_offload_cache", False),
        }
918
        if self.config.task not in ["i2v", "s2v"]:
sandy's avatar
sandy committed
919
920
921
922
923
924
925
926
            return None
        else:
            return Wan2_2_VAE(**vae_config)

    def load_vae(self):
        vae_encoder = self.load_vae_encoder()
        vae_decoder = self.load_vae_decoder()
        return vae_encoder, vae_decoder