src/videogen_hub/pipelines/streamingt2v/model_func.py

# General
import os
from os.path import join as opj
import datetime
import torch
from einops import rearrange, repeat

# Utilities
from videogen_hub.pipelines.streamingt2v.inference_utils import *

from modelscope.outputs import OutputKeys
import imageio
from PIL import Image
import numpy as np

import torch.nn.functional as F
import torchvision.transforms as transforms
from diffusers.utils import load_image

transform = transforms.Compose([transforms.PILToTensor()])


def ms_short_gen(prompt, ms_model, inference_generator, t=50, device="cuda"):
    frames = ms_model(
        prompt,
        num_inference_steps=t,
        generator=inference_generator,
        eta=1.0,
        height=256,
        width=256,
        latents=None,
    ).frames
    frames = torch.stack([torch.from_numpy(frame) for frame in frames])
    frames = frames.to(device).to(torch.float32)
    return rearrange(frames[0], "F W H C -> F C W H")


def ad_short_gen(prompt, ad_model, inference_generator, t=25, device="cuda"):
    frames = ad_model(
        prompt,
        negative_prompt="bad quality, worse quality",
        num_frames=16,
        num_inference_steps=t,
        generator=inference_generator,
        guidance_scale=7.5,
    ).frames[0]
    frames = torch.stack([transform(frame) for frame in frames])
    frames = frames.to(device).to(torch.float32)
    frames = F.interpolate(frames, size=256)
    frames = frames / 255.0
    return frames


def sdxl_image_gen(prompt, sdxl_model):
    image = sdxl_model(prompt=prompt).images[0]
    return image


def svd_short_gen(
    image, prompt, svd_model, sdxl_model, inference_generator, t=25, device="cuda"
):
    if image is None or image == "":
        image = sdxl_image_gen(prompt, sdxl_model)
        image = image.resize((576, 576))
        image = add_margin(image, 0, 224, 0, 224, (0, 0, 0))
    elif type(image) is str:
        image = load_image(image)
        image = resize_and_keep(image)
        image = center_crop(image)
        image = add_margin(image, 0, 224, 0, 224, (0, 0, 0))
    else:
        image = Image.fromarray(np.uint8(image))
        image = resize_and_keep(image)
        image = center_crop(image)
        image = add_margin(image, 0, 224, 0, 224, (0, 0, 0))

    frames = svd_model(
        image, decode_chunk_size=8, generator=inference_generator
    ).frames[0]
    frames = torch.stack([transform(frame) for frame in frames])
    frames = frames.to(device).to(torch.float32)
    frames = frames[:16, :, :, 224:-224]
    frames = F.interpolate(frames, size=256)
    frames = frames / 255.0
    return frames


def stream_long_gen(
    prompt,
    short_video,
    n_autoreg_gen,
    negative_prompt,
    seed,
    t,
    image_guidance,
    result_file_stem,
    stream_cli,
    stream_model,
):
    trainer = stream_cli.trainer
    trainer.limit_predict_batches = 1

    trainer.predict_cfg = {
        "predict_dir": stream_cli.config["result_fol"].as_posix(),
        "result_file_stem": result_file_stem,
        "prompt": prompt,
        "video": short_video,
        "seed": seed,
        "num_inference_steps": t,
        "guidance_scale": image_guidance,
        "n_autoregressive_generations": n_autoreg_gen,
    }
    stream_model.inference_params.negative_prompt = negative_prompt
    trainer.predict(model=stream_model, datamodule=stream_cli.datamodule)


def video2video(prompt, video, where_to_log, cfg_v2v, model_v2v, square=True):
    downscale = cfg_v2v["downscale"]
    upscale_size = cfg_v2v["upscale_size"]
    pad = cfg_v2v["pad"]

    now = datetime.datetime.now()
    now = str(now.time()).replace(":", "_").replace(".", "_")
    name = prompt[:100].replace(" ", "_") + "_" + now
    enhanced_video_mp4 = opj(where_to_log, name + "_enhanced.mp4")

    video_frames = imageio.mimread(video)
    h, w, _ = video_frames[0].shape

    # Downscale video, then resize to fit the upscale size
    video = [
        Image.fromarray(frame).resize((w // downscale, h // downscale))
        for frame in video_frames
    ]
    video = [resize_to_fit(frame, upscale_size) for frame in video]

    if pad:
        video = [pad_to_fit(frame, upscale_size) for frame in video]
    # video = [np.array(frame) for frame in video]

    imageio.mimsave(opj(where_to_log, "temp_" + now + ".mp4"), video, fps=8)

    p_input = {
        "video_path": opj(where_to_log, "temp_" + now + ".mp4"),
        "text": prompt,
        "positive_prompt": prompt,
        "total_noise_levels": 600,
    }
    model_v2v(p_input, output_video=enhanced_video_mp4)[OutputKeys.OUTPUT_VIDEO]

    # Remove padding
    video_frames = imageio.mimread(enhanced_video_mp4)
    video_frames_square = []
    for frame in video_frames:
        frame = frame[:, 280:-280, :]
        video_frames_square.append(frame)
    imageio.mimsave(enhanced_video_mp4, video_frames_square)

    return enhanced_video_mp4


# The main functionality for video to video
def video2video_randomized(
    prompt,
    video,
    where_to_log,
    cfg_v2v,
    model_v2v,
    square=True,
    chunk_size=24,
    overlap_size=8,
    negative_prompt="",
):
    downscale = cfg_v2v["downscale"]
    upscale_size = cfg_v2v["upscale_size"]
    pad = cfg_v2v["pad"]

    now = datetime.datetime.now()
    name = (
        prompt[:100].replace(" ", "_")
        + "_"
        + str(now.time()).replace(":", "_").replace(".", "_")
    )
    enhanced_video_mp4 = opj(where_to_log, name + "_enhanced.mp4")

    video_frames = imageio.mimread(video)
    h, w, _ = video_frames[0].shape

    n_chunks = (len(video_frames) - overlap_size) // (chunk_size - overlap_size)
    trim_length = n_chunks * (chunk_size - overlap_size) + overlap_size
    if trim_length < chunk_size:
        raise ValueError(
            f"Chunk size [{chunk_size}] cannot be larger than the number of frames in the video [{len(video_frames)}], please provide smaller chunk size"
        )
    if trim_length < len(video_frames):
        print(
            "Video cannot be processed with chunk size {chunk_size} and overlap size {overlap_size}, "
            "trimming it to length {trim_length} to be able to process it"
        )
        video_frames = video_frames[:trim_length]

    model_v2v.chunk_size = chunk_size
    model_v2v.overlap_size = overlap_size

    # Downscale video, then resize to fit the upscale size
    video = [
        Image.fromarray(frame).resize((w // downscale, h // downscale))
        for frame in video_frames
    ]
    video = [resize_to_fit(frame, upscale_size) for frame in video]

    if pad:
        video = [pad_to_fit(frame, upscale_size) for frame in video]

    video = list(map(np.array, video))

    imageio.mimsave(opj(where_to_log, "temp.mp4"), video, fps=8)

    p_input = {
        "video_path": opj(where_to_log, "temp.mp4"),
        "text": prompt,
        "positive_prompt": "",
        "negative_prompt": negative_prompt,
        "total_noise_levels": 600,
    }

    output_video_path = model_v2v(p_input, output_video=enhanced_video_mp4)[
        OutputKeys.OUTPUT_VIDEO
    ]

    return enhanced_video_mp4