wan_audio_runner.py 24.8 KB
Newer Older
wangshankun's avatar
wangshankun committed
1
2
3
4
5
6
7
8
9
10
11
12
13
import os
import gc
import numpy as np
import torch
import torchvision.transforms.functional as TF
from PIL import Image
from lightx2v.utils.registry_factory import RUNNER_REGISTER
from lightx2v.models.runners.wan.wan_runner import WanRunner
from lightx2v.models.runners.default_runner import DefaultRunner
from lightx2v.models.schedulers.wan.scheduler import WanScheduler
from lightx2v.models.networks.wan.model import WanModel
from lightx2v.utils.profiler import ProfilingContext4Debug, ProfilingContext
from lightx2v.models.input_encoders.hf.t5.model import T5EncoderModel
wangshankun's avatar
wangshankun committed
14
from lightx2v.models.input_encoders.hf.xlm_roberta.model import CLIPModel, WanVideoIPHandler
wangshankun's avatar
wangshankun committed
15
16
17
18
19
20
from lightx2v.models.networks.wan.audio_model import WanAudioModel
from lightx2v.models.networks.wan.lora_adapter import WanLoraWrapper
from lightx2v.models.video_encoders.hf.wan.vae import WanVAE

from lightx2v.models.networks.wan.audio_adapter import AudioAdapter, AudioAdapterPipe, rank0_load_state_dict_from_path

wangshankun's avatar
wangshankun committed
21
from lightx2v.models.schedulers.wan.step_distill.scheduler import WanStepDistillScheduler
wangshankun's avatar
wangshankun committed
22
from lightx2v.models.schedulers.wan.audio.scheduler import EulerSchedulerTimestepFix
wangshankun's avatar
wangshankun committed
23

wangshankun's avatar
wangshankun committed
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
from loguru import logger
import torch.distributed as dist
from einops import rearrange
import torchaudio as ta
from transformers import AutoFeatureExtractor

from torchvision.datasets.folder import IMG_EXTENSIONS
from torchvision.transforms import InterpolationMode
from torchvision.transforms.functional import resize

import subprocess
import warnings
from typing import Optional, Tuple, Union


wangshankun's avatar
wangshankun committed
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def add_mask_to_frames(
    frames: np.ndarray,
    mask_rate: float = 0.1,
    rnd_state: np.random.RandomState = None,
) -> np.ndarray:
    if mask_rate is None:
        return frames

    if rnd_state is None:
        rnd_state = np.random.RandomState()

    h, w = frames.shape[-2:]
    mask = rnd_state.rand(h, w) > mask_rate
    frames = frames * mask
    return frames


def add_noise_to_frames(
    frames: np.ndarray,
    noise_mean: float = -3.0,
    noise_std: float = 0.5,
    rnd_state: np.random.RandomState = None,
) -> np.ndarray:
    if noise_mean is None or noise_std is None:
        return frames

    if rnd_state is None:
        rnd_state = np.random.RandomState()

    shape = frames.shape
    bs = 1 if len(shape) == 4 else shape[0]
    sigma = rnd_state.normal(loc=noise_mean, scale=noise_std, size=(bs,))
    sigma = np.exp(sigma)
    sigma = np.expand_dims(sigma, axis=tuple(range(1, len(shape))))
    noise = rnd_state.randn(*shape) * sigma
    frames = frames + noise
    return frames


wangshankun's avatar
wangshankun committed
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
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):
    """
    frames: (T, C, H, W)
    size: (H, W)
    """
    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]
    resized_frames = resize(cropped_frames, size, InterpolationMode.BICUBIC, antialias=True)
    return resized_frames


def adaptive_resize(img):
    bucket_config = {
        0.667: (np.array([[480, 832], [544, 960], [720, 1280]], dtype=np.int64), np.array([0.2, 0.5, 0.3])),
        1.0: (np.array([[480, 480], [576, 576], [704, 704], [960, 960]], dtype=np.int64), np.array([0.1, 0.1, 0.5, 0.3])),
        1.5: (np.array([[480, 832], [544, 960], [720, 1280]], dtype=np.int64)[:, ::-1], np.array([0.2, 0.5, 0.3])),
    }
    ori_height = img.shape[-2]
    ori_weight = img.shape[-1]
    ori_ratio = ori_height / ori_weight
    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]
wangshankun's avatar
wangshankun committed
120
121
122
123
124
125
    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
wangshankun's avatar
wangshankun committed
126
127
128
129
130
131
132
133
134
135
    for resolution in bucket_config[closet_ratio][0]:
        if ori_height * ori_weight >= resolution[0] * resolution[1]:
            target_h, target_w = resolution
    cropped_img = isotropic_crop_resize(img, (target_h, target_w))
    return cropped_img, target_h, target_w


def array_to_video(
    image_array: np.ndarray,
    output_path: str,
wangshankun's avatar
wangshankun committed
136
137
    fps: int | float = 30,
    resolution: tuple[int, int] | tuple[float, float] | None = None,
wangshankun's avatar
wangshankun committed
138
139
    disable_log: bool = False,
    lossless: bool = True,
140
    output_pix_fmt: str = "yuv420p",
wangshankun's avatar
wangshankun committed
141
) -> None:
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
    """Convert an array to a video directly, gif not supported.

    Args:
        image_array (np.ndarray): shape should be (f * h * w * 3).
        output_path (str): output video file path.
        fps (Union[int, float, optional): fps. Defaults to 30.
        resolution (Optional[Union[Tuple[int, int], Tuple[float, float]]],
            optional): (height, width) of the output video.
            Defaults to None.
        disable_log (bool, optional): whether close the ffmepg command info.
            Defaults to False.
        output_pix_fmt (str): output pix_fmt in ffmpeg command.
    Raises:
        FileNotFoundError: check output path.
        TypeError: check input array.

    Returns:
        None.
    """
wangshankun's avatar
wangshankun committed
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
    if not isinstance(image_array, np.ndarray):
        raise TypeError("Input should be np.ndarray.")
    assert image_array.ndim == 4
    assert image_array.shape[-1] == 3
    if resolution:
        height, width = resolution
        width += width % 2
        height += height % 2
    else:
        image_array = pad_for_libx264(image_array)
        height, width = image_array.shape[1], image_array.shape[2]
    if lossless:
        command = [
            "/usr/bin/ffmpeg",
            "-y",  # (optional) overwrite output file if it exists
            "-f",
            "rawvideo",
            "-s",
            f"{int(width)}x{int(height)}",  # size of one frame
            "-pix_fmt",
            "bgr24",
            "-r",
            f"{fps}",  # frames per second
            "-loglevel",
            "error",
            "-threads",
            "4",
            "-i",
            "-",  # The input comes from a pipe
            "-vcodec",
            "libx264rgb",
            "-crf",
            "0",
            "-an",  # Tells FFMPEG not to expect any audio
            output_path,
        ]
    else:
198
        output_pix_fmt = output_pix_fmt or "yuv420p"
wangshankun's avatar
wangshankun committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
        command = [
            "/usr/bin/ffmpeg",
            "-y",  # (optional) overwrite output file if it exists
            "-f",
            "rawvideo",
            "-s",
            f"{int(width)}x{int(height)}",  # size of one frame
            "-pix_fmt",
            "bgr24",
            "-r",
            f"{fps}",  # frames per second
            "-loglevel",
            "error",
            "-threads",
            "4",
            "-i",
            "-",  # The input comes from a pipe
            "-vcodec",
            "libx264",
218
219
            "-pix_fmt",
            f"{output_pix_fmt}",
wangshankun's avatar
wangshankun committed
220
221
222
223
            "-an",  # Tells FFMPEG not to expect any audio
            output_path,
        ]

wangshankun's avatar
wangshankun committed
224
225
226
    if output_pix_fmt is not None:
        command += ["-pix_fmt", output_pix_fmt]

wangshankun's avatar
wangshankun committed
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
    if not disable_log:
        print(f'Running "{" ".join(command)}"')
    process = subprocess.Popen(
        command,
        stdin=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    if process.stdin is None or process.stderr is None:
        raise BrokenPipeError("No buffer received.")
    index = 0
    while True:
        if index >= image_array.shape[0]:
            break
        process.stdin.write(image_array[index].tobytes())
        index += 1
    process.stdin.close()
    process.stderr.close()
    process.wait()


def pad_for_libx264(image_array):
    if image_array.ndim == 2 or (image_array.ndim == 3 and image_array.shape[2] == 3):
        hei_index = 0
        wid_index = 1
    elif image_array.ndim == 4 or (image_array.ndim == 3 and image_array.shape[2] != 3):
        hei_index = 1
        wid_index = 2
    else:
        return image_array
    hei_pad = image_array.shape[hei_index] % 2
    wid_pad = image_array.shape[wid_index] % 2
    if hei_pad + wid_pad > 0:
        pad_width = []
        for dim_index in range(image_array.ndim):
            if dim_index == hei_index:
                pad_width.append((0, hei_pad))
            elif dim_index == wid_index:
                pad_width.append((0, wid_pad))
            else:
                pad_width.append((0, 0))
        values = 0
        image_array = np.pad(image_array, pad_width, mode="constant", constant_values=values)
    return image_array


def generate_unique_path(path):
    if not os.path.exists(path):
        return path
    root, ext = os.path.splitext(path)
    index = 1
    new_path = f"{root}-{index}{ext}"
    while os.path.exists(new_path):
        index += 1
        new_path = f"{root}-{index}{ext}"
    return new_path


def save_to_video(gen_lvideo, out_path, target_fps):
    gen_lvideo = rearrange(gen_lvideo, "B C T H W -> B T H W C")
    gen_lvideo = (gen_lvideo[0].cpu().numpy() * 127.5 + 127.5).astype(np.uint8)
    gen_lvideo = gen_lvideo[..., ::-1].copy()
    generate_unique_path(out_path)
wangshankun's avatar
wangshankun committed
289
    array_to_video(gen_lvideo, output_path=out_path, fps=target_fps, lossless=False, output_pix_fmt="yuv444p")
wangshankun's avatar
wangshankun committed
290
291
292


def save_audio(
wangshankun's avatar
wangshankun committed
293
    audio_array,
wangshankun's avatar
wangshankun committed
294
    audio_name: str,
PengGao's avatar
PengGao committed
295
    video_name: str,
wangshankun's avatar
wangshankun committed
296
    sr: int = 16000,
PengGao's avatar
PengGao committed
297
    output_path: Optional[str] = None,
wangshankun's avatar
wangshankun committed
298
299
):
    logger.info(f"Saving audio to {audio_name} type: {type(audio_array)}")
wangshankun's avatar
wangshankun committed
300
301
302
303
304
305

    ta.save(
        audio_name,
        torch.tensor(audio_array[None]),
        sample_rate=sr,
    )
wangshankun's avatar
wangshankun committed
306

PengGao's avatar
PengGao committed
307
308
309
310
311
    if output_path is None:
        out_video = f"{video_name[:-4]}_with_audio.mp4"
    else:
        out_video = output_path

wangshankun's avatar
wangshankun committed
312
313
314
315
316
317
318
    parent_dir = os.path.dirname(out_video)
    if parent_dir and not os.path.exists(parent_dir):
        os.makedirs(parent_dir, exist_ok=True)

    if os.path.exists(out_video):
        os.remove(out_video)

PengGao's avatar
PengGao committed
319
320
321
    subprocess.call(["/usr/bin/ffmpeg", "-y", "-i", video_name, "-i", audio_name, out_video])

    return out_video
wangshankun's avatar
wangshankun committed
322
323
324
325
326
327
328


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

wangshankun's avatar
wangshankun committed
329
330
331
332
    def init_scheduler(self):
        scheduler = EulerSchedulerTimestepFix(self.config)
        self.model.set_scheduler(scheduler)

wangshankun's avatar
wangshankun committed
333
    def load_audio_models(self):
wangshankun's avatar
wangshankun committed
334
335
        ##音频特征提取器
        self.audio_preprocess = AutoFeatureExtractor.from_pretrained(self.config["model_path"], subfolder="audio_encoder")
wangshankun's avatar
wangshankun committed
336
337
338

        ##音频驱动视频生成adapter
        audio_adapter_path = self.config["model_path"] + "/audio_adapter.safetensors"
wangshankun's avatar
wangshankun committed
339
340
341
342
343
344
345
        audio_adaper = AudioAdapter.from_transformer(
            self.model,
            audio_feature_dim=1024,
            interval=1,
            time_freq_dim=256,
            projection_transformer_layers=4,
        )
wangshankun's avatar
wangshankun committed
346
        audio_adapter = rank0_load_state_dict_from_path(audio_adaper, audio_adapter_path, strict=False)
wangshankun's avatar
wangshankun committed
347

wangshankun's avatar
wangshankun committed
348
        ##音频特征编码器
wangshankun's avatar
wangshankun committed
349
        device = self.model.device
wangshankun's avatar
wangshankun committed
350
        audio_encoder_repo = self.config["model_path"] + "/audio_encoder"
wangshankun's avatar
wangshankun committed
351
352
353
354
355
356
        audio_adapter_pipe = AudioAdapterPipe(audio_adapter, audio_encoder_repo=audio_encoder_repo, dtype=torch.bfloat16, device=device, generator=torch.Generator(device), weight=1.0)

        return audio_adapter_pipe

    def load_transformer(self):
        base_model = WanAudioModel(self.config.model_path, self.config, self.init_device)
wangshankun's avatar
wangshankun committed
357

358
        if self.config.get("lora_configs") and self.config.lora_configs:
wangshankun's avatar
wangshankun committed
359
360
            assert not self.config.get("dit_quantized", False) or self.config.mm_config.get("weight_auto_quant", False)
            lora_wrapper = WanLoraWrapper(base_model)
361
362
363
364
365
366
            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
367

wangshankun's avatar
wangshankun committed
368
369
        return base_model

wangshankun's avatar
wangshankun committed
370
    def load_image_encoder(self):
wangshankun's avatar
wangshankun committed
371
372
        clip_model_dir = self.config["model_path"] + "/image_encoder"
        image_encoder = WanVideoIPHandler("CLIPModel", repo_or_path=clip_model_dir, require_grad=False, mode="eval", device=self.init_device, dtype=torch.float16)
wangshankun's avatar
wangshankun committed
373
374
375

        return image_encoder

wangshankun's avatar
wangshankun committed
376
377
378
379
380
381
382
383
384
385
386
    def run_image_encoder(self, config, vae_model):
        ref_img = Image.open(config.image_path)
        ref_img = (np.array(ref_img).astype(np.float32) - 127.5) / 127.5
        ref_img = torch.from_numpy(ref_img).to(vae_model.device)
        ref_img = rearrange(ref_img, "H W C -> 1 C H W")
        ref_img = ref_img[:, :3]

        # resize and crop image
        cond_frms, tgt_h, tgt_w = adaptive_resize(ref_img)
        config.tgt_h = tgt_h
        config.tgt_w = tgt_w
wangshankun's avatar
wangshankun committed
387
        clip_encoder_out = self.image_encoder.encode(cond_frms).squeeze(0).to(torch.bfloat16)
wangshankun's avatar
wangshankun committed
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409

        cond_frms = rearrange(cond_frms, "1 C H W -> 1 C 1 H W")
        lat_h, lat_w = tgt_h // 8, tgt_w // 8
        config.lat_h = lat_h
        config.lat_w = lat_w
        vae_encode_out = vae_model.encode(cond_frms.to(torch.float), config)
        if isinstance(vae_encode_out, list):  #
            # list转tensor
            vae_encode_out = torch.stack(vae_encode_out, dim=0).to(torch.bfloat16)

        return vae_encode_out, clip_encoder_out

    def run_input_encoder_internal(self):
        image_encoder_output = None
        if os.path.isfile(self.config.image_path):
            with ProfilingContext("Run Img Encoder"):
                vae_encode_out, clip_encoder_out = self.run_image_encoder(self.config, self.vae_encoder)
                image_encoder_output = {
                    "clip_encoder_out": clip_encoder_out,
                    "vae_encode_out": vae_encode_out,
                }
                logger.info(f"clip_encoder_out:{clip_encoder_out.shape} vae_encode_out:{vae_encode_out.shape}")
PengGao's avatar
PengGao committed
410

wangshankun's avatar
wangshankun committed
411
        with ProfilingContext("Run Text Encoder"):
PengGao's avatar
PengGao committed
412
413
414
            logger.info(f"Prompt: {self.config['prompt']}")
            img = Image.open(self.config["image_path"]).convert("RGB")
            text_encoder_output = self.run_text_encoder(self.config["prompt"], img)
wangshankun's avatar
wangshankun committed
415
416
417
418

        self.set_target_shape()
        self.inputs = {"text_encoder_output": text_encoder_output, "image_encoder_output": image_encoder_output}

PengGao's avatar
PengGao committed
419
        # del self.image_encoder  # 删除ref的clip模型,只使用一次
wangshankun's avatar
wangshankun committed
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
        gc.collect()
        torch.cuda.empty_cache()

    def set_target_shape(self):
        ret = {}
        num_channels_latents = 16
        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"
            assert 1 == 0, error_msg

        ret["target_shape"] = self.config.target_shape
        return ret

    def run(self):
        def load_audio(in_path: str, sr: float = 16000):
            audio_array, ori_sr = ta.load(in_path)
            audio_array = ta.functional.resample(audio_array.mean(0), orig_freq=ori_sr, new_freq=sr)
            return audio_array.numpy()

        def get_audio_range(start_frame: int, end_frame: int, fps: float, audio_sr: float = 16000):
            audio_frame_rate = audio_sr / fps
            return round(start_frame * audio_frame_rate), round((end_frame + 1) * audio_frame_rate)

wangshankun's avatar
wangshankun committed
452
453
454
455
456
457
458
459
460
461
462
463
        def wan_mask_rearrange(mask: torch.Tensor):
            # mask: 1, T, H, W, where 1 means the input mask is one-channel
            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)
            return mask.transpose(0, 1)  # 4, T // 4, H, W

wangshankun's avatar
wangshankun committed
464
465
466
        self.inputs["audio_adapter_pipe"] = self.load_audio_models()

        # process audio
wangshankun's avatar
wangshankun committed
467
468
        audio_sr = self.config.get("audio_sr", 16000)
        max_num_frames = self.config.get("target_video_length", 81)  # wan2.1一段最多81帧,5秒,16fps
wangshankun's avatar
wangshankun committed
469
        target_fps = self.config.get("target_fps", 16)  # 音视频同步帧率
wangshankun's avatar
wangshankun committed
470
        video_duration = self.config.get("video_duration", 5)  # 期望视频输出时长
wangshankun's avatar
wangshankun committed
471
472
473
474
        audio_array = load_audio(self.config["audio_path"], sr=audio_sr)
        audio_len = int(audio_array.shape[0] / audio_sr * target_fps)
        prev_frame_length = 5
        prev_token_length = (prev_frame_length - 1) // 4 + 1
wangshankun's avatar
wangshankun committed
475
        max_num_audio_length = int((max_num_frames + 1) / target_fps * audio_sr)
wangshankun's avatar
wangshankun committed
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502

        interval_num = 1
        # expected_frames
        expected_frames = min(max(1, int(float(video_duration) * target_fps)), audio_len)
        res_frame_num = 0
        if expected_frames <= max_num_frames:
            interval_num = 1
        else:
            interval_num = max(int((expected_frames - max_num_frames) / (max_num_frames - prev_frame_length)) + 1, 1)
            res_frame_num = expected_frames - interval_num * (max_num_frames - prev_frame_length)
            if res_frame_num > 5:
                interval_num += 1

        audio_start, audio_end = get_audio_range(0, expected_frames, fps=target_fps, audio_sr=audio_sr)
        audio_array_ori = audio_array[audio_start:audio_end]

        gen_video_list = []
        cut_audio_list = []
        # reference latents

        tgt_h = self.config.tgt_h
        tgt_w = self.config.tgt_w
        device = self.model.scheduler.latents.device
        dtype = torch.bfloat16
        vae_dtype = torch.float

        for idx in range(interval_num):
wangshankun's avatar
wangshankun committed
503
504
505
            self.config.seed = self.config.seed + idx
            torch.manual_seed(self.config.seed)
            logger.info(f"###  manual_seed: {self.config.seed} ####")
wangshankun's avatar
wangshankun committed
506
507
508
509
510
511
512
513
514
515
            useful_length = -1
            if idx == 0:  # 第一段 Condition padding0
                prev_frames = torch.zeros((1, 3, max_num_frames, tgt_h, tgt_w), device=device)
                prev_latents = self.vae_encoder.encode(prev_frames.to(vae_dtype), self.config)[0].to(dtype)
                prev_len = 0
                audio_start, audio_end = get_audio_range(0, max_num_frames, fps=target_fps, audio_sr=audio_sr)
                audio_array = audio_array_ori[audio_start:audio_end]
                if expected_frames < max_num_frames:
                    useful_length = audio_array.shape[0]
                    audio_array = np.concatenate((audio_array, np.zeros(max_num_audio_length)[: max_num_audio_length - useful_length]), axis=0)
wangshankun's avatar
wangshankun committed
516
                audio_input_feat = self.audio_preprocess(audio_array, sampling_rate=audio_sr, return_tensors="pt").input_values.squeeze(0)
wangshankun's avatar
wangshankun committed
517
518
519

            elif res_frame_num > 5 and idx == interval_num - 1:  # 最后一段可能不够81帧
                prev_frames = torch.zeros((1, 3, max_num_frames, tgt_h, tgt_w), device=device)
wangshankun's avatar
wangshankun committed
520
521
522
523
524
525
526
527
                last_frames = gen_video_list[-1][:, :, -prev_frame_length:].clone().to(device)

                last_frames = last_frames.cpu().detach().numpy()
                last_frames = add_noise_to_frames(last_frames)
                last_frames = add_mask_to_frames(last_frames, mask_rate=0.1)  # mask 0.10
                last_frames = torch.from_numpy(last_frames).to(dtype=dtype, device=device)

                prev_frames[:, :, :prev_frame_length] = last_frames
wangshankun's avatar
wangshankun committed
528
529
530
531
532
533
                prev_latents = self.vae_encoder.encode(prev_frames.to(vae_dtype), self.config)[0].to(dtype)
                prev_len = prev_token_length
                audio_start, audio_end = get_audio_range(idx * max_num_frames - idx * prev_frame_length, expected_frames, fps=target_fps, audio_sr=audio_sr)
                audio_array = audio_array_ori[audio_start:audio_end]
                useful_length = audio_array.shape[0]
                audio_array = np.concatenate((audio_array, np.zeros(max_num_audio_length)[: max_num_audio_length - useful_length]), axis=0)
wangshankun's avatar
wangshankun committed
534
                audio_input_feat = self.audio_preprocess(audio_array, sampling_rate=audio_sr, return_tensors="pt").input_values.squeeze(0)
wangshankun's avatar
wangshankun committed
535
536
537

            else:  # 中间段满81帧带pre_latens
                prev_frames = torch.zeros((1, 3, max_num_frames, tgt_h, tgt_w), device=device)
wangshankun's avatar
wangshankun committed
538
539
540
541
542
543
544
545
                last_frames = gen_video_list[-1][:, :, -prev_frame_length:].clone().to(device)

                last_frames = last_frames.cpu().detach().numpy()
                last_frames = add_noise_to_frames(last_frames)
                last_frames = add_mask_to_frames(last_frames, mask_rate=0.1)  # mask 0.10
                last_frames = torch.from_numpy(last_frames).to(dtype=dtype, device=device)

                prev_frames[:, :, :prev_frame_length] = last_frames
wangshankun's avatar
wangshankun committed
546
547
548
549
                prev_latents = self.vae_encoder.encode(prev_frames.to(vae_dtype), self.config)[0].to(dtype)
                prev_len = prev_token_length
                audio_start, audio_end = get_audio_range(idx * max_num_frames - idx * prev_frame_length, (idx + 1) * max_num_frames - idx * prev_frame_length, fps=target_fps, audio_sr=audio_sr)
                audio_array = audio_array_ori[audio_start:audio_end]
wangshankun's avatar
wangshankun committed
550
                audio_input_feat = self.audio_preprocess(audio_array, sampling_rate=audio_sr, return_tensors="pt").input_values.squeeze(0)
wangshankun's avatar
wangshankun committed
551
552
553
554
555
556
557
558

            self.inputs["audio_encoder_output"] = audio_input_feat.to(device)

            if idx != 0:
                self.model.scheduler.reset()

            if prev_latents is not None:
                ltnt_channel, nframe, height, width = self.model.scheduler.latents.shape
wangshankun's avatar
wangshankun committed
559
560
                # bs = 1
                frames_n = (nframe - 1) * 4 + 1
wangshankun's avatar
wangshankun committed
561
562
563
                prev_frame_len = max((prev_len - 1) * 4 + 1, 0)
                prev_mask = torch.ones((1, frames_n, height, width), device=device, dtype=dtype)
                prev_mask[:, prev_frame_len:] = 0
wangshankun's avatar
wangshankun committed
564
                prev_mask = wan_mask_rearrange(prev_mask).unsqueeze(0)
wangshankun's avatar
wangshankun committed
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
                previmg_encoder_output = {
                    "prev_latents": prev_latents,
                    "prev_mask": prev_mask,
                }
                self.inputs["previmg_encoder_output"] = previmg_encoder_output

            for step_index in range(self.model.scheduler.infer_steps):
                logger.info(f"==> step_index: {step_index} / {self.model.scheduler.infer_steps}")

                with ProfilingContext4Debug("step_pre"):
                    self.model.scheduler.step_pre(step_index=step_index)

                with ProfilingContext4Debug("infer"):
                    self.model.infer(self.inputs)

                with ProfilingContext4Debug("step_post"):
                    self.model.scheduler.step_post()

            latents = self.model.scheduler.latents
            generator = self.model.scheduler.generator
            gen_video = self.vae_decoder.decode(latents, generator=generator, config=self.config)
            gen_video = torch.clamp(gen_video, -1, 1)
            start_frame = 0 if idx == 0 else prev_frame_length
            start_audio_frame = 0 if idx == 0 else int((prev_frame_length + 1) * audio_sr / target_fps)
wangshankun's avatar
wangshankun committed
589

wangshankun's avatar
wangshankun committed
590
            if res_frame_num > 5 and idx == interval_num - 1:
wangshankun's avatar
wangshankun committed
591
                gen_video_list.append(gen_video[:, :, start_frame:res_frame_num].cpu())
wangshankun's avatar
wangshankun committed
592
593
                cut_audio_list.append(audio_array[start_audio_frame:useful_length])
            elif expected_frames < max_num_frames and useful_length != -1:
wangshankun's avatar
wangshankun committed
594
                gen_video_list.append(gen_video[:, :, start_frame:expected_frames].cpu())
wangshankun's avatar
wangshankun committed
595
596
                cut_audio_list.append(audio_array[start_audio_frame:useful_length])
            else:
wangshankun's avatar
wangshankun committed
597
                gen_video_list.append(gen_video[:, :, start_frame:].cpu())
wangshankun's avatar
wangshankun committed
598
599
600
601
                cut_audio_list.append(audio_array[start_audio_frame:])

        gen_lvideo = torch.cat(gen_video_list, dim=2).float()
        merge_audio = np.concatenate(cut_audio_list, axis=0).astype(np.float32)
PengGao's avatar
PengGao committed
602
        out_path = os.path.join("./", "video_merge.mp4")
wangshankun's avatar
wangshankun committed
603
604
        audio_file = os.path.join("./", "audio_merge.wav")
        save_to_video(gen_lvideo, out_path, target_fps)
PengGao's avatar
PengGao committed
605
        save_audio(merge_audio, audio_file, out_path, output_path=self.config.get("save_video_path", None))
wangshankun's avatar
wangshankun committed
606
607
608
        os.remove(out_path)
        os.remove(audio_file)

609
    def run_pipeline(self):
wangshankun's avatar
wangshankun committed
610
611
612
613
614
615
616
617
618
619
620
621
        if self.config["use_prompt_enhancer"]:
            self.config["prompt_enhanced"] = self.post_prompt_enhancer()

        self.run_input_encoder_internal()
        self.set_target_shape()

        self.init_scheduler()
        self.model.scheduler.prepare(self.inputs["image_encoder_output"])
        self.run()
        self.end_run()

        gc.collect()
wangshankun's avatar
wangshankun committed
622
        torch.cuda.empty_cache()