run_demo_video_continuation.py

import os
import argparse

import cv2
import datetime
import PIL.Image
import numpy as np

import torch
import torch.distributed as dist

from transformers import AutoTokenizer, UMT5EncoderModel
from torchvision.io import write_video
from diffusers.utils import load_video

from longcat_video.pipeline_longcat_video import LongCatVideoPipeline
from longcat_video.modules.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
from longcat_video.modules.autoencoder_kl_wan import AutoencoderKLWan
from longcat_video.modules.longcat_video_dit import LongCatVideoTransformer3DModel
from longcat_video.context_parallel import context_parallel_util
from longcat_video.context_parallel.context_parallel_util import init_context_parallel


def torch_gc():
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()


def get_fps(video_path):
    cap = cv2.VideoCapture(video_path)
    original_fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()
    
    return original_fps


def generate(args):
    # case setup
    video_path = "assets/motorcycle.mp4"
    video = load_video(video_path)
    prompt = "A person rides a motorcycle along a long, straight road that stretches between a body of water and a forested hillside. The rider steadily accelerates, keeping the motorcycle centered between the guardrails, while the scenery passes by on both sides. The video captures the journey from the rider’s perspective, emphasizing the sense of motion and adventure."
    negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
    num_cond_frames = 13
    spatial_refine_only = False

    # load parsed args
    checkpoint_dir = args.checkpoint_dir
    context_parallel_size = args.context_parallel_size
    enable_compile = args.enable_compile

    # prepare distributed environment
    rank = int(os.environ['RANK'])
    num_gpus = torch.cuda.device_count()
    local_rank = rank % num_gpus
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl", timeout=datetime.timedelta(seconds=3600*24))
    global_rank    = dist.get_rank()
    num_processes  = dist.get_world_size()

    # initialize context parallel before loading models
    init_context_parallel(context_parallel_size=context_parallel_size, global_rank=global_rank, world_size=num_processes)
    cp_size = context_parallel_util.get_cp_size()
    cp_split_hw = context_parallel_util.get_optimal_split(cp_size)

    tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir, subfolder="tokenizer", torch_dtype=torch.bfloat16)
    text_encoder = UMT5EncoderModel.from_pretrained(checkpoint_dir, subfolder="text_encoder", torch_dtype=torch.bfloat16)
    vae = AutoencoderKLWan.from_pretrained(checkpoint_dir, subfolder="vae", torch_dtype=torch.bfloat16)
    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(checkpoint_dir, subfolder="scheduler", torch_dtype=torch.bfloat16)
    dit = LongCatVideoTransformer3DModel.from_pretrained(checkpoint_dir, subfolder="dit", cp_split_hw=cp_split_hw, torch_dtype=torch.bfloat16)

    if enable_compile:
        dit = torch.compile(dit)

    pipe = LongCatVideoPipeline(
        tokenizer = tokenizer,
        text_encoder = text_encoder,
        vae = vae,
        scheduler = scheduler,
        dit = dit,
    )
    pipe.to(local_rank)

    global_seed = 42
    seed = global_seed + global_rank

    generator = torch.Generator(device=local_rank)
    generator.manual_seed(seed)

    target_fps = 15
    target_size = video[0].size  # (width, height)
    current_fps = get_fps(video_path)
    stride = max(1, round(current_fps / target_fps))

    ### vc (480p)
    output = pipe.generate_vc(
        video=video[::stride],
        prompt=prompt,
        negative_prompt=negative_prompt,
        resolution='480p', # 480p / 720p
        num_frames=93,
        num_cond_frames=num_cond_frames,
        num_inference_steps=50,
        guidance_scale=4.0,
        generator=generator,
        use_kv_cache=True,
        offload_kv_cache=False,
    )[0]

    if local_rank == 0:
        output = [(output[i] * 255).astype(np.uint8) for i in range(output.shape[0])]
        output = [PIL.Image.fromarray(img) for img in output]
        output = [frame.resize(target_size, PIL.Image.BICUBIC) for frame in output]

        output = video[::stride] + output[num_cond_frames:]
        output_tensor = torch.from_numpy(np.array(output))
        write_video("output_vc.mp4", output_tensor, fps=15, video_codec="libx264", options={"crf": f"{18}"})
    del output
    torch_gc()

    ### vc distill (480p)
    cfg_step_lora_path = os.path.join(checkpoint_dir, 'lora/cfg_step_lora.safetensors')
    pipe.dit.load_lora(cfg_step_lora_path, 'cfg_step_lora')
    pipe.dit.enable_loras(['cfg_step_lora'])

    if enable_compile:
        dit = torch.compile(dit)

    output_distill = pipe.generate_vc(
        video=video[::stride],
        prompt=prompt,
        resolution='480p', # 480p / 720p
        num_frames=93,
        num_cond_frames=num_cond_frames,
        num_inference_steps=16,
        use_distill=True,
        guidance_scale=1.0,
        generator=generator,
        use_kv_cache=True,
        offload_kv_cache=False,
        enhance_hf=False,
    )[0]
    pipe.dit.disable_all_loras()

    if local_rank == 0:
        output_processed = [(output_distill[i] * 255).astype(np.uint8) for i in range(output_distill.shape[0])]
        output_processed = [PIL.Image.fromarray(img) for img in output_processed]
        output_processed = [frame.resize(target_size, PIL.Image.BICUBIC) for frame in output_processed]

        output = video[::stride] + output_processed[num_cond_frames:]
        output_tensor = torch.from_numpy(np.array(output))
        write_video("output_vc_distill.mp4", output_tensor, fps=15, video_codec="libx264", options={"crf": f"{18}"})

    ### vc refinement (720p)
    refinement_lora_path = os.path.join(checkpoint_dir, 'lora/refinement_lora.safetensors')
    pipe.dit.load_lora(refinement_lora_path, 'refinement_lora')
    pipe.dit.enable_loras(['refinement_lora'])
    pipe.dit.enable_bsa()

    if enable_compile:
        dit = torch.compile(dit)
    
    stage1_video = [(output_distill[i] * 255).astype(np.uint8) for i in range(output_distill.shape[0])]
    stage1_video = [PIL.Image.fromarray(img) for img in stage1_video]
    del output_distill 
    torch_gc()

    target_fps = 30
    stride = max(1, round(current_fps / target_fps))
    cur_num_cond_frames = num_cond_frames if spatial_refine_only else num_cond_frames * 2

    output_refine = pipe.generate_refine(
        video=video[::stride],
        prompt=prompt,
        stage1_video=stage1_video,
        num_cond_frames=cur_num_cond_frames,
        num_inference_steps=50,
        generator=generator,
        spatial_refine_only=spatial_refine_only
    )[0]

    pipe.dit.disable_all_loras()
    pipe.dit.disable_bsa()

    if local_rank == 0:
        output_refine = [(output_refine[i] * 255).astype(np.uint8) for i in range(output_refine.shape[0])]
        output_refine = [PIL.Image.fromarray(img) for img in output_refine]
        output_refine = [frame.resize(target_size, PIL.Image.BICUBIC) for frame in output_refine]
        output_refine = video[::stride] + output_refine[cur_num_cond_frames:]

        output_tensor = torch.from_numpy(np.array(output_refine))
        fps = 15 if spatial_refine_only else 30
        write_video("output_vc_refine.mp4", output_tensor, fps=fps, video_codec="libx264", options={"crf": f"{10}"})


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--context_parallel_size",
        type=int,
        default=1,
    )
    parser.add_argument(
        "--checkpoint_dir",
        type=str,
        default=None,
    )
    parser.add_argument(
        '--enable_compile',
        action='store_true',
    )

    args = parser.parse_args()

    return args


if __name__ == "__main__":
    args = _parse_args()
    generate(args)