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

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

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

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

wangshankun's avatar
wangshankun committed
38

39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
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
54
55
                h_ratio *= 2
            else:
56
                patched_w //= 2
57
                w_ratio *= 2
58
    return patched_h * h_ratio, patched_w * w_ratio
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79


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):
    """
80
    frames: (C, H, W) or (T, C, H, W) or (N, C, H, W)
81
82
    size: (H, W)
    """
83
84
85
86
87
88
89
    original_shape = frames.shape

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

90
91
92
93
    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]
94
    resized_frames = resize(cropped_frames, [h, w], InterpolationMode.BICUBIC, antialias=True)
95
96
97
98

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

99
100
101
    return resized_frames


102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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


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

    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])
130

131
132
133
134
135
136
137
138
139
140
141
    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)
142
        # logger.info(f"[wan_audio] use custom bucket_shape: {bucket_config}")
143
144
145
146
147
148
    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),
        }
149
        # logger.info(f"[wan_audio] use default bucket_shape: {bucket_config}")
150

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

    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
165
        for resolution in bucket_config[closet_ratio]:
166
167
168
169
170
171
172
173
174
175
176
            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]
177
        target_h, target_w = bucket_config[closet_ratio][0]
178
179
180
181
182
183
184
185
186
187
    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)
188
189
190
191
    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]
192
        target_h, target_w = bucket_config[closet_ratio][-1]
193

194
195
196
197
    cropped_img = isotropic_crop_resize(img, (target_h, target_w))
    return cropped_img, target_h, target_w


198
199
200
201
@dataclass
class AudioSegment:
    """Data class for audio segment information"""

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


207
class FramePreprocessorTorchVersion:
208
209
210
211
212
213
214
    """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

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

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

        # 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
231
232
        return frames + noise

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

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

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

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


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
256
        self.audio_frame_rate = audio_sr // target_fps
257

sandy's avatar
sandy committed
258
    def load_audio(self, audio_path: str):
259
260
        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
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
        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
280
281
282

    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
283
        return round(start_frame * self.audio_frame_rate), round(end_frame * self.audio_frame_rate)
284

sandy's avatar
sandy committed
285
286
287
288
289
    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
        """
290
        segments = []
sandy's avatar
sandy committed
291
        segments_idx = self.init_segments_idx(expected_frames, max_num_frames, prev_frame_length)
292

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

sandy's avatar
sandy committed
296
297
        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
298
            audio_array = audio_array_ori[:, audio_start:audio_end]
299

sandy's avatar
sandy committed
300
301
            if idx < len(segments_idx) - 1:
                end_idx = segments_idx[idx + 1][0]
sandy's avatar
sandy committed
302
303
304
305
306
            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
307
                    end_idx = end_idx - padding_len // self.audio_frame_rate
308

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

sandy's avatar
sandy committed
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
    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

329

Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
330
@RUNNER_REGISTER("seko_talk")
helloyongyang's avatar
helloyongyang committed
331
332
333
class WanAudioRunner(WanRunner):  # type:ignore
    def __init__(self, config):
        super().__init__(config)
334
        self.prev_frame_length = self.config.get("prev_frame_length", 5)
335
        self.frame_preprocessor = FramePreprocessorTorchVersion()
helloyongyang's avatar
helloyongyang committed
336
337
338

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

341
    def read_audio_input(self, audio_path):
sandy's avatar
sandy committed
342
        """Read audio input - handles both single and multi-person scenarios"""
helloyongyang's avatar
helloyongyang committed
343
344
345
        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
346

LiangLiu's avatar
LiangLiu committed
347
348
349
        if not isinstance(audio_path, str):
            return [], 0, None, 0

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

sandy's avatar
sandy committed
353
354
355
356
357
358
359
360
361
        # 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)
helloyongyang's avatar
helloyongyang committed
362
363
364
        expected_frames = min(max(1, int(video_duration * target_fps)), audio_len)

        # Segment audio
365
        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
366

367
368
369
370
371
372
        # 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
373

374
        return audio_segments, expected_frames, mask_latents, len(audio_files)
sandy's avatar
sandy committed
375

376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
    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
392

393
        return audio_files, mask_files
sandy's avatar
sandy committed
394

395
    def process_single_mask(self, mask_file):
sandy's avatar
sandy committed
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
        mask_img = Image.open(mask_file).convert("RGB")
        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
418
419

    def read_image_input(self, img_path):
LiangLiu's avatar
LiangLiu committed
420
421
422
423
        if isinstance(img_path, Image.Image):
            ref_img = img_path
        else:
            ref_img = Image.open(img_path).convert("RGB")
helloyongyang's avatar
helloyongyang committed
424
425
        ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(0).cuda()

426
427
428
429
430
431
432
        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),
        )
433
        logger.info(f"[wan_audio] resize_image target_h: {h}, target_w: {w}")
434
435
        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
436
437
438

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

439
440
        latent_h = patched_h * self.config["patch_size"][1]
        latent_w = patched_w * self.config["patch_size"][2]
helloyongyang's avatar
helloyongyang committed
441

442
443
        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
444

445
        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
446

447
448
        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
449
450

    def run_image_encoder(self, first_frame, last_frame=None):
451
452
        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
453
        clip_encoder_out = self.image_encoder.visual([first_frame]).squeeze(0).to(GET_DTYPE()) if self.config.get("use_image_encoder", True) else None
454
455
456
457
        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
458
459
460
        return clip_encoder_out

    def run_vae_encoder(self, img):
461
462
463
        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
464
        img = rearrange(img, "1 C H W -> 1 C 1 H W")
465
        vae_encoder_out = self.vae_encoder.encode(img.to(GET_DTYPE()))
sandy's avatar
sandy committed
466

467
468
469
470
        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
471
472
        return vae_encoder_out

473
    @ProfilingContext4DebugL2("Run Encoders")
474
475
476
477
    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
478
479
        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
480

481
482
483
484
        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
485
486
487
488
489
490
491
492
493
494
        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
495
            "person_mask_latens": person_mask_latens,
helloyongyang's avatar
helloyongyang committed
496
        }
497
498
499

    def prepare_prev_latents(self, prev_video: Optional[torch.Tensor], prev_frame_length: int) -> Optional[Dict[str, torch.Tensor]]:
        """Prepare previous latents for conditioning"""
wangshankun's avatar
wangshankun committed
500
        device = torch.device("cuda")
501
        dtype = GET_DTYPE()
502

503
504
        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)
505

506
507
508
        if prev_video is not None:
            # Extract and process last frames
            last_frames = prev_video[:, :, -prev_frame_length:].clone().to(device)
509
            if self.config["model_cls"] != "wan2.2_audio":
sandy's avatar
sandy committed
510
                last_frames = self.frame_preprocessor.process_prev_frames(last_frames)
511
            prev_frames[:, :, :prev_frame_length] = last_frames
sandy's avatar
sandy committed
512
513
514
            prev_len = (prev_frame_length - 1) // 4 + 1
        else:
            prev_len = 0
515

516
517
518
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            self.vae_encoder = self.load_vae_encoder()

519
        _, nframe, height, width = self.model.scheduler.latents.shape
520
        with ProfilingContext4DebugL1("vae_encoder in init run segment"):
521
            if self.config["model_cls"] == "wan2.2_audio":
522
523
524
525
526
                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
527
            else:
528
                prev_latents = self.vae_encoder.encode(prev_frames.to(dtype))
529

530
531
            frames_n = (nframe - 1) * 4 + 1
            prev_mask = torch.ones((1, frames_n, height, width), device=device, dtype=dtype)
532
533
            prev_frame_len = max((prev_len - 1) * 4 + 1, 0)
            prev_mask[:, prev_frame_len:] = 0
534
            prev_mask = self._wan_mask_rearrange(prev_mask)
helloyongyang's avatar
fix ci  
helloyongyang committed
535

sandy's avatar
sandy committed
536
537
        if prev_latents is not None:
            if prev_latents.shape[-2:] != (height, width):
538
                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
539
                prev_latents = torch.nn.functional.interpolate(prev_latents, size=(height, width), mode="bilinear", align_corners=False)
540

541
542
543
544
545
        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
546
        return {"prev_latents": prev_latents, "prev_mask": prev_mask, "prev_len": prev_len}
547
548
549
550
551
552
553
554
555
556
557

    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
558
        return mask.transpose(0, 1).contiguous()
559

helloyongyang's avatar
helloyongyang committed
560
561
    def get_video_segment_num(self):
        self.video_segment_num = len(self.inputs["audio_segments"])
wangshankun's avatar
wangshankun committed
562

helloyongyang's avatar
helloyongyang committed
563
564
    def init_run(self):
        super().init_run()
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
565
        self.scheduler.set_audio_adapter(self.audio_adapter)
helloyongyang's avatar
helloyongyang committed
566
        self.prev_video = None
567
568
        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
569
            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
570
571
        else:
            self.gen_video_final = None
sandy's avatar
sandy committed
572
            self.cut_audio_final = None
wangshankun's avatar
wangshankun committed
573

574
    @ProfilingContext4DebugL1("Init run segment")
LiangLiu's avatar
LiangLiu committed
575
    def init_run_segment(self, segment_idx, audio_array=None):
helloyongyang's avatar
helloyongyang committed
576
        self.segment_idx = segment_idx
LiangLiu's avatar
LiangLiu committed
577
        if audio_array is not None:
LiangLiu's avatar
LiangLiu committed
578
579
580
            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
581
582
        else:
            self.segment = self.inputs["audio_segments"][segment_idx]
wangshankun's avatar
wangshankun committed
583

584
585
        self.input_info.seed = self.input_info.seed + segment_idx
        torch.manual_seed(self.input_info.seed)
586
        # logger.info(f"Processing segment {segment_idx + 1}/{self.video_segment_num}, seed: {self.config.seed}")
wangshankun's avatar
wangshankun committed
587

588
589
590
        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
591
592
593
594
595
596
        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
597

helloyongyang's avatar
helloyongyang committed
598
        self.inputs["audio_encoder_output"] = audio_features
599
        self.inputs["previmg_encoder_output"] = self.prepare_prev_latents(self.prev_video, prev_frame_length=self.prev_frame_length)
wangshankun's avatar
wangshankun committed
600

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

605
    @ProfilingContext4DebugL1("End run segment")
606
    def end_run_segment(self, segment_idx):
helloyongyang's avatar
helloyongyang committed
607
        self.gen_video = torch.clamp(self.gen_video, -1, 1).to(torch.float)
sandy's avatar
sandy committed
608
        useful_length = self.segment.end_frame - self.segment.start_frame
LiangLiu's avatar
LiangLiu committed
609
        video_seg = self.gen_video[:, :, :useful_length].cpu()
sandy's avatar
sandy committed
610
611
        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
612
613
614
615
616
617
618
619
620
621
622
        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
623

LiangLiu's avatar
LiangLiu committed
624
625
        if self.va_recorder:
            self.va_recorder.pub_livestream(video_seg, audio_seg)
626
        elif self.input_info.return_result_tensor:
LiangLiu's avatar
LiangLiu committed
627
            self.gen_video_final[self.segment.start_frame : self.segment.end_frame].copy_(video_seg)
sandy's avatar
sandy committed
628
            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
629

helloyongyang's avatar
helloyongyang committed
630
631
632
        # Update prev_video for next iteration
        self.prev_video = self.gen_video

LiangLiu's avatar
LiangLiu committed
633
        del video_seg, audio_seg
helloyongyang's avatar
helloyongyang committed
634
635
        torch.cuda.empty_cache()

LiangLiu's avatar
LiangLiu committed
636
637
638
639
640
641
642
643
644
    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):
645
        output_video_path = self.input_info.save_result_path
LiangLiu's avatar
LiangLiu committed
646
647
        self.va_recorder = None
        if isinstance(output_video_path, dict):
LiangLiu's avatar
LiangLiu committed
648
649
650
651
652
653
654
655
            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
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670

            whip_shared_path = os.getenv("WHIP_SHARED_LIB", None)
            if whip_shared_path and output_video_path.startswith("http"):
                self.va_recorder = VAX64Recorder(
                    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
671
672

    def init_va_reader(self):
LiangLiu's avatar
LiangLiu committed
673
        audio_path = self.input_info.audio_path
LiangLiu's avatar
LiangLiu committed
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
        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,
            )

    def run_main(self, total_steps=None):
        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:
                return super().run_main(total_steps)

LiangLiu's avatar
LiangLiu committed
701
            self.va_reader.start()
LiangLiu's avatar
LiangLiu committed
702
            rank, world_size = self.get_rank_and_world_size()
LiangLiu's avatar
LiangLiu committed
703
704
            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
705
706
707
                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
708
709

            self.init_run()
LiangLiu's avatar
LiangLiu committed
710
            if self.config.get("compile", False):
711
                self.model.select_graph_for_compile(self.input_info)
LiangLiu's avatar
LiangLiu committed
712
713
714
715
716
717
718
719
            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:
720
                with ProfilingContext4DebugL1(f"stream segment get audio segment {segment_idx}"):
LiangLiu's avatar
LiangLiu committed
721
722
723
724
725
726
727
728
729
                    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

730
                with ProfilingContext4DebugL1(f"stream segment end2end {segment_idx}"):
LiangLiu's avatar
LiangLiu committed
731
732
                    fail_count = 0
                    self.init_run_segment(segment_idx, audio_array)
helloyongyang's avatar
helloyongyang committed
733
                    latents = self.run_segment(total_steps=None)
LiangLiu's avatar
LiangLiu committed
734
                    self.gen_video = self.run_vae_decoder(latents)
LiangLiu's avatar
LiangLiu committed
735
                    self.end_run_segment(segment_idx)
LiangLiu's avatar
LiangLiu committed
736
737
738
                    segment_idx += 1

        finally:
LiangLiu's avatar
LiangLiu committed
739
            if hasattr(self.model, "inputs"):
LiangLiu's avatar
LiangLiu committed
740
741
742
743
744
                self.end_run()
            if self.va_reader:
                self.va_reader.stop()
                self.va_reader = None
            if self.va_recorder:
LiangLiu's avatar
LiangLiu committed
745
                self.va_recorder.stop()
LiangLiu's avatar
LiangLiu committed
746
747
                self.va_recorder = None

748
    @ProfilingContext4DebugL1("Process after vae decoder")
749
750
    def process_images_after_vae_decoder(self):
        if self.input_info.return_result_tensor:
sandy's avatar
sandy committed
751
            audio_waveform = self.cut_audio_final.unsqueeze(0).unsqueeze(0)
LiangLiu's avatar
LiangLiu committed
752
753
754
            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}
755

wangshankun's avatar
wangshankun committed
756
    def load_transformer(self):
757
        """Load transformer with LoRA support"""
758
759
760
        base_model = WanAudioModel(self.config["model_path"], self.config, self.init_device)
        if self.config.get("lora_configs") and self.config["lora_configs"]:
            assert not self.config.get("dit_quantized", False) or self.config["mm_config"].get("weight_auto_quant", False)
wangshankun's avatar
wangshankun committed
761
            lora_wrapper = WanLoraWrapper(base_model)
762
            for lora_config in self.config["lora_configs"]:
763
764
765
766
767
                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
768

wangshankun's avatar
wangshankun committed
769
770
        return base_model

helloyongyang's avatar
helloyongyang committed
771
    def load_audio_encoder(self):
gushiqiao's avatar
gushiqiao committed
772
        audio_encoder_path = self.config.get("audio_encoder_path", os.path.join(self.config["model_path"], "TencentGameMate-chinese-hubert-large"))
773
774
        audio_encoder_offload = self.config.get("audio_encoder_cpu_offload", self.config.get("cpu_offload", False))
        model = SekoAudioEncoderModel(audio_encoder_path, self.config["audio_sr"], audio_encoder_offload)
helloyongyang's avatar
helloyongyang committed
775
        return model
776

helloyongyang's avatar
helloyongyang committed
777
    def load_audio_adapter(self):
778
779
780
781
782
        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:
            device = torch.device("cuda")
helloyongyang's avatar
helloyongyang committed
783
        audio_adapter = AudioAdapter(
sandy's avatar
sandy committed
784
            attention_head_dim=self.config["dim"] // self.config["num_heads"],
helloyongyang's avatar
helloyongyang committed
785
786
787
788
789
790
791
792
793
            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),
794
            cpu_offload=audio_adapter_offload,
helloyongyang's avatar
helloyongyang committed
795
        )
796

797
        audio_adapter.to(device)
798
        load_from_rank0 = self.config.get("load_from_rank0", False)
799
        weights_dict = load_weights(self.config["adapter_model_path"], cpu_offload=audio_adapter_offload, remove_key="ca", load_from_rank0=load_from_rank0)
800
        audio_adapter.load_state_dict(weights_dict, strict=False)
helloyongyang's avatar
helloyongyang committed
801
        return audio_adapter.to(dtype=GET_DTYPE())
wangshankun's avatar
wangshankun committed
802

helloyongyang's avatar
helloyongyang committed
803
804
    def load_model(self):
        super().load_model()
805
806
807
        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
808

809
810
811
812
813
814
815
816
    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
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852


@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 = {
            "vae_pth": find_torch_model_path(self.config, "vae_pth", "Wan2.2_VAE.pth"),
            "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 = {
            "vae_pth": find_torch_model_path(self.config, "vae_pth", "Wan2.2_VAE.pth"),
            "device": vae_device,
            "cpu_offload": vae_offload,
            "offload_cache": self.config.get("vae_offload_cache", False),
        }
853
        if self.config.task not in ["i2v", "s2v"]:
sandy's avatar
sandy committed
854
855
856
857
858
859
860
861
            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