scripts/inference.py

import os
import math
import time
from pprint import pformat
import warnings
warnings.filterwarnings('ignore')

import colossalai
import torch
import torch.distributed as dist
from peft import PeftModel
from colossalai.cluster import DistCoordinator
from mmengine.runner import set_random_seed
from tqdm import tqdm
import numpy as np
import pandas as pd

from javisdit.acceleration.parallel_states import set_sequence_parallel_group
from javisdit.datasets import save_sample
from javisdit.datasets.aspect import get_image_size, get_num_frames
from javisdit.models.text_encoder.t5 import text_preprocessing
from javisdit.registry import MODELS, SCHEDULERS, build_module
from javisdit.utils.config_utils import parse_configs
from javisdit.utils.inference_utils import (
    add_watermark,
    append_generated,
    append_score_to_prompts,
    apply_va_mask_strategy,
    collect_va_references_batch,
    dframe_to_frame,
    extract_json_from_prompts,
    extract_prompts_loop,
    get_save_path_name,
    load_prompts,
    merge_prompt,
    prepare_multi_resolution_info,
    refine_prompts_by_openai,
    split_prompt
)
from javisdit.utils.misc import all_exists, create_logger, is_distributed, is_main_process, to_torch_dtype


def main():
    torch.set_grad_enabled(False)
    # ======================================================
    # configs & runtime variables
    # ======================================================
    # == parse configs ==
    cfg = parse_configs(training=False)
    audio_only = cfg.get('audio_only', False)

    # == device and dtype ==
    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg_dtype = cfg.get("dtype", "fp32")
    assert cfg_dtype in ["fp16", "bf16", "fp32"], f"Unknown mixed precision {cfg_dtype}"
    dtype = to_torch_dtype(cfg.get("dtype", "bf16"))
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    # == init distributed env ==
    if is_distributed():
        colossalai.launch_from_torch({})
        coordinator = DistCoordinator()
        enable_sequence_parallelism = cfg.enable_sequence_parallelism and coordinator.world_size > 1
        if enable_sequence_parallelism:
            set_sequence_parallel_group(dist.group.WORLD)
    else:
        coordinator = None
        enable_sequence_parallelism = False
    set_random_seed(seed=cfg.get("seed", 1024))

    # == init logger ==
    logger = create_logger()
    logger.info("Inference configuration:\n %s", pformat(cfg.to_dict()))
    verbose = cfg.get("verbose", 1)
    progress_wrap = tqdm if verbose == 1 else (lambda x: x)

    torch.set_num_threads(1) # NOTE: without it, loading audioldm2 is really really slow

    # ======================================================
    # build model & load weights
    # ======================================================
    logger.info("Building models...")
    # == build text-encoder and vae ==
    text_encoder = build_module(cfg.text_encoder, MODELS, device=device, dtype=dtype)
    prior_encoder = build_module(cfg.get('prior_encoder', None), MODELS)
    if prior_encoder is not None:
        prior_encoder = prior_encoder.to(device, dtype).eval()        
    vae = build_module(cfg.vae, MODELS).to(device, dtype).eval()
    audio_vae = build_module(cfg.audio_vae, MODELS, device=device, dtype=dtype)

    # == prepare video size ==
    image_size = cfg.get("image_size", None)
    if image_size is None:
        resolution = cfg.get("resolution", None)
        aspect_ratio = cfg.get("aspect_ratio", None)
        assert (
            resolution is not None and aspect_ratio is not None
        ), "resolution and aspect_ratio must be provided if image_size is not provided"
        image_size = get_image_size(resolution, aspect_ratio)
    num_frames = get_num_frames(cfg.num_frames)

    # == build diffusion model ==
    input_size = (num_frames, *image_size)
    v_latent_size = vae.get_latent_size(input_size)  # [t//4 for every 17 frame, h//8, w//8]
    ckpt_path = cfg.model.pop('weight_init_from', cfg.get('model_path', ''))
    model = (
        build_module(
            cfg.model,
            MODELS,
            input_size=v_latent_size,
            in_channels=vae.out_channels,
            caption_channels=text_encoder.output_dim,
            model_max_length=text_encoder.model_max_length,
            enable_sequence_parallelism=enable_sequence_parallelism,
            weight_init_from=ckpt_path,
        )
        .to(device, dtype)
        .eval()
    )
    text_encoder.y_embedder = model.y_embedder  # HACK: for classifier-free guidance
    if prior_encoder is not None:
        prior_encoder.st_prior_embedder = model.st_prior_embedder
    
    lora_ckpt_path = os.path.join(ckpt_path, cfg.get("lora_dir", "lora"))
    if os.path.exists(lora_ckpt_path):
        logger.info(f'Loading LoRA weight from {lora_ckpt_path}')
        model = PeftModel.from_pretrained(model, lora_ckpt_path, is_trainable=False)
        logger.info("Merging LoRA weights...")
        model = model.merge_and_unload()

    # == build scheduler ==
    scheduler = build_module(cfg.scheduler, SCHEDULERS)

    # ======================================================
    # inference
    # ======================================================
    # == load prompts ==
    prompts = cfg.get("prompt", None)
    start_idx = cfg.get("start_index", 0)
    if prompts is None:
        if cfg.get("prompt_path", None) is not None:
            prompts = load_prompts(cfg.prompt_path, start_idx, cfg.get("end_index", None),
                                   prompt_key=cfg.get("prompt_key", "text"))
        else:
            prompts = [cfg.get("prompt_generator", "")] * 1_000_000  # endless loop
    elif isinstance(prompts, str):
        prompts = [prompts]

    # == prepare reference ==
    reference_path = cfg.get("reference_path", [("", "")] * len(prompts))
    mask_strategy = cfg.get("mask_strategy", [""] * len(prompts))
    if isinstance(reference_path, str) and os.path.isfile(reference_path):
        reference_df = pd.read_csv(reference_path)
        reference_path = list(zip(reference_df['path'].tolist(), reference_df['audio_path'].tolist()))
    if isinstance(mask_strategy, str):
        mask_strategy = [mask_strategy] * len(prompts)
    assert len(reference_path) == len(prompts), "Length of reference must be the same as prompts"
    assert len(mask_strategy) == len(prompts), "Length of mask_strategy must be the same as prompts"

    # == prepare distributed inference ==
    if is_distributed() and not enable_sequence_parallelism:
        local_rank = dist.get_rank()
        num_per_rank = math.ceil(len(prompts) / coordinator.world_size)
        local_start_idx = local_rank * num_per_rank
        local_end_idx = local_start_idx + num_per_rank
        prompts = prompts[local_start_idx:local_end_idx]
        reference_path = reference_path[local_start_idx:local_end_idx]
        mask_strategy = mask_strategy[local_start_idx:local_end_idx]
        start_idx += local_start_idx

    # == prepare arguments ==
    fps = cfg.fps
    save_fps = cfg.get("save_fps", fps // cfg.get("frame_interval", 1))
    multi_resolution = cfg.get("multi_resolution", None)
    batch_size = cfg.get("batch_size", 1)
    num_sample = cfg.get("num_sample", 1)
    loop = cfg.get("loop", 1)
    condition_frame_length = cfg.get("condition_frame_length", 5)
    condition_frame_edit = cfg.get("condition_frame_edit", 0.0)
    align = cfg.get("align", None)
    assert loop == 1, "not implemented"
    audio_fps = cfg.get("audio_fps", 16000)
    neg_prompts = cfg.get("neg_prompt", None)
    if isinstance(neg_prompts, str):
        neg_prompts = [neg_prompts] * len(prompts)
    use_text_preprocessing = cfg.get("use_text_preprocessing", True)

    save_dir = cfg.save_dir
    os.makedirs(save_dir, exist_ok=True)
    sample_name = cfg.get("sample_name", None)
    prompt_as_path = cfg.get("prompt_as_path", False)

    # == Iter over all samples ==
    for i in progress_wrap(range(0, len(prompts), batch_size)):
        # == prepare batch prompts ==
        batch_prompts = prompts[i : i + batch_size]
        ms = mask_strategy[i : i + batch_size]
        refs = reference_path[i : i + batch_size]
        neg_prompts_batch = neg_prompts[i : i + batch_size] if neg_prompts is not None else None

        # == get json from prompts ==
        batch_prompts, refs, ms = extract_json_from_prompts(batch_prompts, refs, ms, noimpl=True)
        original_batch_prompts = batch_prompts

        # == get reference for condition ==
        refs = collect_va_references_batch(refs, vae, image_size, 
                                           audio_vae=audio_vae, audio_cfg=cfg.get("audio_cfg", {}))

        # == multi-resolution info ==
        model_args = prepare_multi_resolution_info(
            multi_resolution, len(batch_prompts), image_size, num_frames, fps, device, dtype
        )

        # == Iter over number of sampling for one prompt ==
        for k in range(num_sample):
            # == prepare save paths ==
            save_paths = [
                get_save_path_name(
                    save_dir,
                    sample_name=sample_name,
                    sample_idx=start_idx + idx,
                    prompt=original_batch_prompts[idx],
                    prompt_as_path=prompt_as_path,
                    num_sample=num_sample,
                    k=k,
                )
                for idx in range(len(batch_prompts))
            ]
            if enable_sequence_parallelism and coordinator.world_size > 1:
                dist.barrier()
            if all(os.path.exists(f'{_path}.mp4') for _path in save_paths):
                continue

            # NOTE: Skip if the sample already exists
            # This is useful for resuming sampling VBench
            if prompt_as_path and all_exists(save_paths):
                continue

            # == process prompts step by step ==
            # 0. split prompt
            # each element in the list is [prompt_segment_list, loop_idx_list]
            batched_prompt_segment_list = []
            batched_loop_idx_list = []
            for prompt in batch_prompts:
                prompt_segment_list, loop_idx_list = split_prompt(prompt)
                batched_prompt_segment_list.append(prompt_segment_list)
                batched_loop_idx_list.append(loop_idx_list)

            # 1. refine prompt by openai
            if cfg.get("llm_refine", False):
                # only call openai API when
                # 1. seq parallel is not enabled
                # 2. seq parallel is enabled and the process is rank 0
                if not enable_sequence_parallelism or (enable_sequence_parallelism and is_main_process()):
                    for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):
                        batched_prompt_segment_list[idx] = refine_prompts_by_openai(prompt_segment_list)

                # sync the prompt if using seq parallel
                if enable_sequence_parallelism:
                    coordinator.block_all()
                    prompt_segment_length = [
                        len(prompt_segment_list) for prompt_segment_list in batched_prompt_segment_list
                    ]

                    # flatten the prompt segment list
                    batched_prompt_segment_list = [
                        prompt_segment
                        for prompt_segment_list in batched_prompt_segment_list
                        for prompt_segment in prompt_segment_list
                    ]

                    # create a list of size equal to world size
                    broadcast_obj_list = [batched_prompt_segment_list] * coordinator.world_size
                    dist.broadcast_object_list(broadcast_obj_list, 0)

                    # recover the prompt list
                    batched_prompt_segment_list = []
                    segment_start_idx = 0
                    all_prompts = broadcast_obj_list[0]
                    for num_segment in prompt_segment_length:
                        batched_prompt_segment_list.append(
                            all_prompts[segment_start_idx : segment_start_idx + num_segment]
                        )
                        segment_start_idx += num_segment

            # 2. append score
            for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):
                batched_prompt_segment_list[idx] = append_score_to_prompts(
                    prompt_segment_list,
                    aes=cfg.get("aes", None),
                    flow=cfg.get("flow", None),
                    camera_motion=cfg.get("camera_motion", None),
                )

            # 3. clean prompt with T5
            for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):
                batched_prompt_segment_list[idx] = [
                    text_preprocessing(prompt, use_text_preprocessing=use_text_preprocessing) \
                        for prompt in prompt_segment_list
                ]

            # 4. merge to obtain the final prompt
            batch_prompts = []
            for prompt_segment_list, loop_idx_list in zip(batched_prompt_segment_list, batched_loop_idx_list):
                batch_prompts.append(merge_prompt(prompt_segment_list, loop_idx_list))

            # == Iter over loop generation ==
            video_clips, audio_clips = [], []
            for loop_i in range(loop):
                # == get prompt for loop i ==
                batch_prompts_loop = extract_prompts_loop(batch_prompts, loop_i)

                # == add condition frames for loop ==
                if loop_i > 0:
                    raise NotImplementedError
                    refs, ms = append_generated(
                        vae, video_clips[-1], refs, ms, loop_i, condition_frame_length, condition_frame_edit
                    )

                # == sampling ==
                if cfg.get('fix_noise_seed', None):
                    torch.manual_seed(1024)  ## TODO: fix or not
                vz = torch.randn(len(batch_prompts), vae.out_channels, *v_latent_size, device=device, dtype=dtype)
                audio_length_in_s = num_frames / fps
                az, original_waveform_length = audio_vae.prepare_latents(audio_length_in_s, len(batch_prompts), device=device, dtype=dtype)
                masks = apply_va_mask_strategy(vz, az, refs, ms, loop_i, align=align,
                                               v2a_t_scale=1/5*17/fps/(10.24/1024)/4)

                samples = scheduler.multimodal_sample(
                    model,
                    text_encoder,
                    {'video': vz, 'audio': az},
                    batch_prompts_loop,
                    device=device,
                    additional_args=model_args,
                    progress=verbose >= 2,
                    mask=masks,
                    prior_encoder=prior_encoder,
                    neg_prompts=neg_prompts_batch
                )
                video_samples, audio_samples = samples['video'], samples['audio']
                
                video_samples = vae.decode(video_samples.to(dtype), num_frames=num_frames)
                video_clips.append(video_samples)
                audio_samples = audio_vae.decode_audio(audio_samples, original_waveform_length=original_waveform_length)
                audio_clips.append(audio_samples)

            # == save samples ==
            if not enable_sequence_parallelism or is_main_process():
                for idx, batch_prompt in enumerate(batch_prompts):
                    if verbose >= 2:
                        logger.info("Prompt: %s", batch_prompt)
                    save_path = save_paths[idx]
                    video = [video_clips[i][idx] for i in range(loop)]
                    audio = [audio_clips[i][idx] for i in range(loop)]
                    for i in range(1, loop):  # TODO: audio
                        raise NotImplementedError
                        video[i] = video[i][:, dframe_to_frame(condition_frame_length) :]
                        # audio[i] = audio[i][:, audio_dframe_to_frame(condition_frame_length) :]
                    video = torch.cat(video, dim=1)
                    audio = torch.cat(audio, dim=0)
                    save_path = save_sample(
                        video,
                        fps=save_fps,
                        audio=audio,
                        audio_fps=audio_fps,
                        save_path=save_path,
                        verbose=verbose >= 2,
                        audio_only=audio_only,
                    )
                    if save_path.endswith(".mp4") and cfg.get("watermark", False):
                        time.sleep(1)  # prevent loading previous generated video
                        add_watermark(save_path)
        start_idx += len(batch_prompts)
    logger.info("Inference finished.")
    logger.info("Saved %s samples to %s", len(prompts), save_dir)


if __name__ == "__main__":
    main()