wan_audio_runner.py 37.7 KB
Newer Older
wangshankun's avatar
wangshankun committed
1
import gc
PengGao's avatar
PengGao committed
2
import os
sandy's avatar
sandy committed
3
import warnings
PengGao's avatar
PengGao committed
4
from dataclasses import dataclass
sandy's avatar
sandy committed
5
from pathlib import Path
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
22
from lightx2v.models.input_encoders.hf.seko_audio.audio_adapter import AudioAdapter
helloyongyang's avatar
helloyongyang committed
23
from lightx2v.models.input_encoders.hf.seko_audio.audio_encoder import SekoAudioEncoderModel
24
from lightx2v.models.networks.wan.audio_model import WanAudioModel
PengGao's avatar
PengGao committed
25
from lightx2v.models.networks.wan.lora_adapter import WanLoraWrapper
26
from lightx2v.models.runners.wan.wan_runner import WanRunner
27
from lightx2v.models.schedulers.wan.audio.scheduler import EulerScheduler
sandy's avatar
sandy committed
28
from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE
29
from lightx2v.utils.envs import *
30
from lightx2v.utils.profiler import *
PengGao's avatar
PengGao committed
31
from lightx2v.utils.registry_factory import RUNNER_REGISTER
LiangLiu's avatar
LiangLiu committed
32
from lightx2v.utils.utils import find_torch_model_path, load_weights, vae_to_comfyui_image_inplace
33

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

wangshankun's avatar
wangshankun committed
37

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


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

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

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

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

98
99
100
    return resized_frames


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


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

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

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

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

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

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


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

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


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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

328

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

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

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

        # Get audio files from person objects or legacy format
        audio_files = self._get_audio_files_from_config()
        if not audio_files:
LiangLiu's avatar
LiangLiu committed
349
            return [], 0
helloyongyang's avatar
helloyongyang committed
350

sandy's avatar
sandy committed
351
352
353
354
355
356
357
358
        # 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)

        self.config.audio_num = audio_array.size(0)
helloyongyang's avatar
helloyongyang committed
359

sandy's avatar
sandy committed
360
361
        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

sandy's avatar
sandy committed
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
        return audio_array.size(0), audio_segments, expected_frames

    def _get_audio_files_from_config(self):
        talk_objects = self.config.get("talk_objects")
        if talk_objects:
            audio_files = []
            for idx, person in enumerate(talk_objects):
                audio_path = person.get("audio")
                if audio_path and Path(audio_path).is_file():
                    audio_files.append(str(audio_path))
                else:
                    logger.warning(f"Person {idx} audio file {audio_path} does not exist or not specified")
            if audio_files:
                logger.info(f"Loaded {len(audio_files)} audio files from talk_objects")
            return audio_files

        audio_path = self.config.get("audio_path")
        if audio_path:
            return [audio_path]

        logger.error("config audio_path or talk_objects is not specified")
        return []

    def read_person_mask(self):
        mask_files = self._get_mask_files_from_config()
        if not mask_files:
            return None

        mask_latents = []
        for mask_file in mask_files:
            mask_latent = self._process_single_mask(mask_file)
            mask_latents.append(mask_latent)

        mask_latents = torch.cat(mask_latents, dim=0)
        return mask_latents

    def _get_mask_files_from_config(self):
        talk_objects = self.config.get("talk_objects")
        if talk_objects:
            mask_files = []
            for idx, person in enumerate(talk_objects):
                mask_path = person.get("mask")
                if mask_path and Path(mask_path).is_file():
                    mask_files.append(str(mask_path))
                elif mask_path:
                    logger.warning(f"Person {idx} mask file {mask_path} does not exist")
            if mask_files:
                logger.info(f"Loaded {len(mask_files)} mask files from talk_objects")
            return mask_files

        logger.info("config talk_objects is not specified")
        return None

    def _process_single_mask(self, mask_file):
        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
443
444

    def read_image_input(self, img_path):
LiangLiu's avatar
LiangLiu committed
445
446
447
448
        if isinstance(img_path, Image.Image):
            ref_img = img_path
        else:
            ref_img = Image.open(img_path).convert("RGB")
helloyongyang's avatar
helloyongyang committed
449
450
        ref_img = TF.to_tensor(ref_img).sub_(0.5).div_(0.5).unsqueeze(0).cuda()

451
452
453
454
455
456
457
        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),
        )
458
        logger.info(f"[wan_audio] resize_image target_h: {h}, target_w: {w}")
helloyongyang's avatar
helloyongyang committed
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
        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]

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

        self.config.lat_h = patched_h * self.config.patch_size[1]
        self.config.lat_w = patched_w * self.config.patch_size[2]

        self.config.tgt_h = self.config.lat_h * self.config.vae_stride[1]
        self.config.tgt_w = self.config.lat_w * self.config.vae_stride[2]

        logger.info(f"[wan_audio] tgt_h: {self.config.tgt_h}, tgt_w: {self.config.tgt_w}, lat_h: {self.config.lat_h}, lat_w: {self.config.lat_w}")

        ref_img = torch.nn.functional.interpolate(ref_img, size=(self.config.tgt_h, self.config.tgt_w), mode="bicubic")
        return ref_img

    def run_image_encoder(self, first_frame, last_frame=None):
476
477
        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
478
        clip_encoder_out = self.image_encoder.visual([first_frame]).squeeze(0).to(GET_DTYPE()) if self.config.get("use_image_encoder", True) else None
479
480
481
482
        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
483
484
485
        return clip_encoder_out

    def run_vae_encoder(self, img):
486
487
488
        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
489
        img = rearrange(img, "1 C H W -> 1 C 1 H W")
490
        vae_encoder_out = self.vae_encoder.encode(img.to(GET_DTYPE()))
sandy's avatar
sandy committed
491

492
493
494
495
        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
496
497
        return vae_encoder_out

498
    @ProfilingContext4DebugL2("Run Encoders")
helloyongyang's avatar
helloyongyang committed
499
500
501
502
503
    def _run_input_encoder_local_r2v_audio(self):
        prompt = self.config["prompt_enhanced"] if self.config["use_prompt_enhancer"] else self.config["prompt"]
        img = self.read_image_input(self.config["image_path"])
        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
504
505
506
507
508
509
510
511

        audio_num, audio_segments, expected_frames = self.read_audio_input()
        person_mask_latens = self.read_person_mask()
        self.config.person_num = 0
        if person_mask_latens is not None:
            assert audio_num == person_mask_latens.size(0), "audio_num and person_mask_latens.size(0) must be the same"
            self.config.person_num = person_mask_latens.size(0)

helloyongyang's avatar
helloyongyang committed
512
513
514
515
516
517
518
519
520
521
522
        text_encoder_output = self.run_text_encoder(prompt, None)
        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
523
            "person_mask_latens": person_mask_latens,
helloyongyang's avatar
helloyongyang committed
524
        }
525
526
527

    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
528
        device = torch.device("cuda")
529
        dtype = GET_DTYPE()
530
531
532
533

        tgt_h, tgt_w = self.config.tgt_h, self.config.tgt_w
        prev_frames = torch.zeros((1, 3, self.config.target_video_length, tgt_h, tgt_w), device=device)

534
535
536
        if prev_video is not None:
            # Extract and process last frames
            last_frames = prev_video[:, :, -prev_frame_length:].clone().to(device)
sandy's avatar
sandy committed
537
538
            if self.config.model_cls != "wan2.2_audio":
                last_frames = self.frame_preprocessor.process_prev_frames(last_frames)
539
            prev_frames[:, :, :prev_frame_length] = last_frames
sandy's avatar
sandy committed
540
541
542
            prev_len = (prev_frame_length - 1) // 4 + 1
        else:
            prev_len = 0
543

544
545
546
        if self.config.get("lazy_load", False) or self.config.get("unload_modules", False):
            self.vae_encoder = self.load_vae_encoder()

547
        _, nframe, height, width = self.model.scheduler.latents.shape
548
        with ProfilingContext4DebugL1("vae_encoder in init run segment"):
549
550
551
552
553
554
            if self.config.model_cls == "wan2.2_audio":
                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
555
            else:
556
                prev_latents = self.vae_encoder.encode(prev_frames.to(dtype))
557

558
559
            frames_n = (nframe - 1) * 4 + 1
            prev_mask = torch.ones((1, frames_n, height, width), device=device, dtype=dtype)
560
561
            prev_frame_len = max((prev_len - 1) * 4 + 1, 0)
            prev_mask[:, prev_frame_len:] = 0
562
            prev_mask = self._wan_mask_rearrange(prev_mask)
helloyongyang's avatar
fix ci  
helloyongyang committed
563

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

569
570
571
572
573
        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
574
        return {"prev_latents": prev_latents, "prev_mask": prev_mask, "prev_len": prev_len}
575
576
577
578
579
580
581
582
583
584
585

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

helloyongyang's avatar
helloyongyang committed
588
589
    def get_video_segment_num(self):
        self.video_segment_num = len(self.inputs["audio_segments"])
wangshankun's avatar
wangshankun committed
590

helloyongyang's avatar
helloyongyang committed
591
592
    def init_run(self):
        super().init_run()
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
593
        self.scheduler.set_audio_adapter(self.audio_adapter)
helloyongyang's avatar
helloyongyang committed
594
        self.prev_video = None
LiangLiu's avatar
LiangLiu committed
595
596
        if self.config.get("return_video", False):
            self.gen_video_final = torch.zeros((self.inputs["expected_frames"], self.config.tgt_h, self.config.tgt_w, 3), dtype=torch.float32, device="cpu")
sandy's avatar
sandy committed
597
            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
598
599
        else:
            self.gen_video_final = None
sandy's avatar
sandy committed
600
            self.cut_audio_final = None
wangshankun's avatar
wangshankun committed
601

602
    @ProfilingContext4DebugL1("Init run segment")
LiangLiu's avatar
LiangLiu committed
603
    def init_run_segment(self, segment_idx, audio_array=None):
helloyongyang's avatar
helloyongyang committed
604
        self.segment_idx = segment_idx
LiangLiu's avatar
LiangLiu committed
605
        if audio_array is not None:
sandy's avatar
sandy committed
606
            end_idx = audio_array.shape[1] // self._audio_processor.audio_frame_rate - self.prev_frame_length
LiangLiu's avatar
LiangLiu committed
607
            self.segment = AudioSegment(audio_array, 0, end_idx)
LiangLiu's avatar
LiangLiu committed
608
609
        else:
            self.segment = self.inputs["audio_segments"][segment_idx]
wangshankun's avatar
wangshankun committed
610

helloyongyang's avatar
helloyongyang committed
611
612
        self.config.seed = self.config.seed + segment_idx
        torch.manual_seed(self.config.seed)
613
        # logger.info(f"Processing segment {segment_idx + 1}/{self.video_segment_num}, seed: {self.config.seed}")
wangshankun's avatar
wangshankun committed
614

615
616
617
        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
618
619
620
621
622
623
        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
624

helloyongyang's avatar
helloyongyang committed
625
        self.inputs["audio_encoder_output"] = audio_features
626
        self.inputs["previmg_encoder_output"] = self.prepare_prev_latents(self.prev_video, prev_frame_length=self.prev_frame_length)
wangshankun's avatar
wangshankun committed
627

helloyongyang's avatar
helloyongyang committed
628
629
        # Reset scheduler for non-first segments
        if segment_idx > 0:
sandy's avatar
sandy committed
630
            self.model.scheduler.reset(self.inputs["previmg_encoder_output"])
wangshankun's avatar
wangshankun committed
631

632
    @ProfilingContext4DebugL1("End run segment")
helloyongyang's avatar
helloyongyang committed
633
634
    def end_run_segment(self):
        self.gen_video = torch.clamp(self.gen_video, -1, 1).to(torch.float)
sandy's avatar
sandy committed
635
        useful_length = self.segment.end_frame - self.segment.start_frame
LiangLiu's avatar
LiangLiu committed
636
        video_seg = self.gen_video[:, :, :useful_length].cpu()
sandy's avatar
sandy committed
637
638
        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
639
640
641
642
643
644
645
646
647
648
649
        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
650

LiangLiu's avatar
LiangLiu committed
651
652
653
654
        if self.va_recorder:
            self.va_recorder.pub_livestream(video_seg, audio_seg)
        elif self.config.get("return_video", False):
            self.gen_video_final[self.segment.start_frame : self.segment.end_frame].copy_(video_seg)
sandy's avatar
sandy committed
655
            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
656

helloyongyang's avatar
helloyongyang committed
657
658
659
        # Update prev_video for next iteration
        self.prev_video = self.gen_video

LiangLiu's avatar
LiangLiu committed
660
        del video_seg, audio_seg
helloyongyang's avatar
helloyongyang committed
661
662
        torch.cuda.empty_cache()

LiangLiu's avatar
LiangLiu committed
663
664
665
666
667
668
669
670
671
672
673
674
    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):
        output_video_path = self.config.get("save_video_path", None)
        self.va_recorder = None
        if isinstance(output_video_path, dict):
LiangLiu's avatar
LiangLiu committed
675
676
677
678
679
680
681
682
683
684
685
686
687
            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"]
            self.va_recorder = VARecorder(
                livestream_url=output_video_path,
                fps=record_fps,
                sample_rate=audio_sr,
            )
LiangLiu's avatar
LiangLiu committed
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718

    def init_va_reader(self):
        audio_path = self.config.get("audio_path", None)
        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)

            rank, world_size = self.get_rank_and_world_size()
LiangLiu's avatar
LiangLiu committed
719
720
            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
721
722
723
            self.va_reader.start()

            self.init_run()
LiangLiu's avatar
LiangLiu committed
724
725
            if self.config.get("compile", False):
                self.model.select_graph_for_compile()
LiangLiu's avatar
LiangLiu committed
726
727
728
729
730
731
732
733
            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:
734
                with ProfilingContext4DebugL1(f"stream segment get audio segment {segment_idx}"):
LiangLiu's avatar
LiangLiu committed
735
736
737
738
739
740
741
742
743
                    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

744
                with ProfilingContext4DebugL1(f"stream segment end2end {segment_idx}"):
LiangLiu's avatar
LiangLiu committed
745
746
                    fail_count = 0
                    self.init_run_segment(segment_idx, audio_array)
helloyongyang's avatar
helloyongyang committed
747
                    latents = self.run_segment(total_steps=None)
LiangLiu's avatar
LiangLiu committed
748
749
750
751
752
                    self.gen_video = self.run_vae_decoder(latents)
                    self.end_run_segment()
                    segment_idx += 1

        finally:
LiangLiu's avatar
LiangLiu committed
753
            if hasattr(self.model, "inputs"):
LiangLiu's avatar
LiangLiu committed
754
755
756
757
758
759
760
761
                self.end_run()
            if self.va_reader:
                self.va_reader.stop()
                self.va_reader = None
            if self.va_recorder:
                self.va_recorder.stop(wait=False)
                self.va_recorder = None

762
    @ProfilingContext4DebugL1("Process after vae decoder")
LiangLiu's avatar
LiangLiu committed
763
764
    def process_images_after_vae_decoder(self, save_video=False):
        if self.config.get("return_video", False):
sandy's avatar
sandy committed
765
            audio_waveform = self.cut_audio_final.unsqueeze(0).unsqueeze(0)
LiangLiu's avatar
LiangLiu committed
766
767
768
            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}
769

helloyongyang's avatar
helloyongyang committed
770
771
772
    def init_modules(self):
        super().init_modules()
        self.run_input_encoder = self._run_input_encoder_local_r2v_audio
wangshankun's avatar
wangshankun committed
773
774

    def load_transformer(self):
775
        """Load transformer with LoRA support"""
wangshankun's avatar
wangshankun committed
776
        base_model = WanAudioModel(self.config.model_path, self.config, self.init_device)
777
        if self.config.get("lora_configs") and self.config.lora_configs:
wangshankun's avatar
wangshankun committed
778
779
            assert not self.config.get("dit_quantized", False) or self.config.mm_config.get("weight_auto_quant", False)
            lora_wrapper = WanLoraWrapper(base_model)
780
781
782
783
784
785
            for lora_config in self.config.lora_configs:
                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
786

wangshankun's avatar
wangshankun committed
787
788
        return base_model

helloyongyang's avatar
helloyongyang committed
789
    def load_audio_encoder(self):
gushiqiao's avatar
gushiqiao committed
790
        audio_encoder_path = self.config.get("audio_encoder_path", os.path.join(self.config["model_path"], "TencentGameMate-chinese-hubert-large"))
791
792
        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
793
        return model
794

helloyongyang's avatar
helloyongyang committed
795
    def load_audio_adapter(self):
796
797
798
799
800
        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
801
        audio_adapter = AudioAdapter(
sandy's avatar
sandy committed
802
            attention_head_dim=self.config["dim"] // self.config["num_heads"],
helloyongyang's avatar
helloyongyang committed
803
804
805
806
807
808
809
810
811
            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),
812
            cpu_offload=audio_adapter_offload,
helloyongyang's avatar
helloyongyang committed
813
        )
814

815
        audio_adapter.to(device)
816
817
        load_from_rank0 = self.config.get("load_from_rank0", False)
        weights_dict = load_weights(self.config.adapter_model_path, cpu_offload=audio_adapter_offload, remove_key="ca", load_from_rank0=load_from_rank0)
818
        audio_adapter.load_state_dict(weights_dict, strict=False)
helloyongyang's avatar
helloyongyang committed
819
        return audio_adapter.to(dtype=GET_DTYPE())
wangshankun's avatar
wangshankun committed
820

helloyongyang's avatar
helloyongyang committed
821
822
    def load_model(self):
        super().load_model()
823
824
825
        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
826
827

    def set_target_shape(self):
828
        """Set target shape for generation"""
wangshankun's avatar
wangshankun committed
829
830
        ret = {}
        num_channels_latents = 16
wangshankun's avatar
wangshankun committed
831
832
        if self.config.model_cls == "wan2.2_audio":
            num_channels_latents = self.config.num_channels_latents
833

wangshankun's avatar
wangshankun committed
834
835
836
837
838
839
840
841
842
843
844
        if self.config.task == "i2v":
            self.config.target_shape = (
                num_channels_latents,
                (self.config.target_video_length - 1) // self.config.vae_stride[0] + 1,
                self.config.lat_h,
                self.config.lat_w,
            )
            ret["lat_h"] = self.config.lat_h
            ret["lat_w"] = self.config.lat_w
        else:
            error_msg = "t2v task is not supported in WanAudioRunner"
845
            assert False, error_msg
wangshankun's avatar
wangshankun committed
846
847
848

        ret["target_shape"] = self.config.target_shape
        return ret
sandy's avatar
sandy committed
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893


@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),
        }
        if self.config.task != "i2v":
            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