app.py

import torch
# torch.jit.script = lambda f: f
# General
import os
from os.path import join as opj
import argparse
import datetime
from pathlib import Path
# import spaces
import gradio as gr
import tempfile
import yaml
from t2v_enhanced.model.video_ldm import VideoLDM

# Utilities
from t2v_enhanced.inference_utils import *
from t2v_enhanced.model_init import *
from t2v_enhanced.model_func import *


on_huggingspace = os.environ.get("SPACE_AUTHOR_NAME") == "PAIR"
parser = argparse.ArgumentParser()
parser.add_argument('--public_access', action='store_true', default=True)
parser.add_argument('--where_to_log', type=str, default="gradio_output")
parser.add_argument('--device', type=str, default="cuda")
args = parser.parse_args()
default_prompt = "A man with yellow ballon head is riding a bike on the street of New York City"

Path(args.where_to_log).mkdir(parents=True, exist_ok=True)
result_fol = Path(args.where_to_log).absolute()
device = args.device
n_devices = int(os.environ.get('NDEVICES', 4))
if n_devices == 4:
    devices = [f"cuda:{idx}" for idx in range(4)]
else:
    devices = ["cuda"] * 4
# --------------------------
# ----- Configurations -----
# --------------------------
cfg_v2v = {'downscale': 1, 'upscale_size': (1280, 720), 'model_id': 'damo/Video-to-Video', 'pad': True}


# --------------------------
# ----- Initialization -----
# --------------------------
ms_model = init_modelscope(devices[1])
# # zs_model = init_zeroscope(device)
ad_model = init_animatediff(devices[1])
svd_model = init_svd(devices[2])
sdxl_model = init_sdxl(devices[2])

ckpt_file_streaming_t2v = Path("t2v_enhanced/checkpoints/streaming_t2v.ckpt").absolute()
stream_cli, stream_model = init_streamingt2v_model(ckpt_file_streaming_t2v, result_fol)
msxl_model = init_v2v_model(cfg_v2v, devices[3])




# -------------------------
# ----- Functionality -----
# -------------------------
# @spaces.GPU(duration=120)
def generate(prompt, num_frames, image, model_name_stage1, model_name_stage2, seed, t, image_guidance, where_to_log=result_fol):
    now = datetime.datetime.now()
    name = prompt[:100].replace(" ", "_") + "_" + str(now.time()).replace(":", "_").replace(".", "_")

    if num_frames == [] or num_frames is None:
        num_frames = 24
    else:
        num_frames = int(num_frames.split(" ")[0])
        if num_frames > 56:
            num_frames = 56
    
    if prompt == "" or prompt is None:
        prompt = default_prompt

    n_autoreg_gen = (num_frames-8)//8

    if model_name_stage1 == "ModelScopeT2V (text to video)":
        inference_generator = torch.Generator(device=ms_model.device).manual_seed(seed)
        short_video = ms_short_gen(prompt, ms_model, inference_generator, t, device)
    elif model_name_stage1 == "AnimateDiff (text to video)":
        inference_generator = torch.Generator(device=ad_model.device).manual_seed(seed)
        short_video = ad_short_gen(prompt, ad_model, inference_generator, t, device)
    elif model_name_stage1 == "SVD (image to video)":
        # For cached examples
        if isinstance(image, dict):
            image = image["path"]
        inference_generator = torch.Generator(device=svd_model.device).manual_seed(seed)
        short_video = svd_short_gen(image, prompt, svd_model, sdxl_model, inference_generator, t, device)

    stream_long_gen(prompt, short_video, n_autoreg_gen, seed, t, image_guidance, name, stream_cli, stream_model)
    video_path = opj(where_to_log, name+".mp4")
    return video_path

# @spaces.GPU(duration=400)
def enhance(prompt, input_to_enhance, num_frames=None, image=None, model_name_stage1=None, model_name_stage2=None, seed=33, t=50, image_guidance=9.5, result_fol=result_fol):
    if prompt == "" or prompt is None:
        prompt = default_prompt
    
    if input_to_enhance is None:
        input_to_enhance = generate(prompt, num_frames, image, model_name_stage1, model_name_stage2, seed, t, image_guidance)
    encoded_video = video2video(prompt, input_to_enhance, result_fol, cfg_v2v, msxl_model)
    # for idx in range(4):
    #     print(f">>> cuda:{idx}", torch.cuda.max_memory_allocated(f"cuda:{idx}"))
    return encoded_video

def change_visibility(value):
    if value == "SVD (image to video)":
        return gr.Image(label='Image Prompt (if not attached then SDXL will be used to generate the starting image)', show_label=True, scale=1, show_download_button=False, interactive=True, value=None)
    else:
        return gr.Image(label='Image Prompt (first select Image-to-Video model from advanced options to enable image upload)', show_label=True, scale=1, show_download_button=False, interactive=False, value=None)


# [prompt_stage1, video_stage2, num_frames, image_stage1, model_name_stage1, seed, t, image_guidance]
examples_1 = [
        ["Experience the dance of jellyfish: float through mesmerizing swarms of jellyfish, pulsating with otherworldly grace and beauty.",
            "__assets__/examples/t2v/1.mp4", "56 - frames", None, "ModelScopeT2V (text to video)", 33, 50, 9.0],
        ["People dancing in room filled with fog and colorful lights.",
            "__assets__/examples/t2v/2.mp4", "56 - frames", None, "ModelScopeT2V (text to video)", 33, 50, 9.0],
        ["Discover the secret language of bees: delve into the complex communication system that allows bees to coordinate their actions and navigate the world.",
            "__assets__/examples/t2v/3.mp4", "56 - frames", None, "AnimateDiff (text to video)", 33, 50, 9.0],
        ["sunset, orange sky, warm lighting, fishing boats, ocean waves seagulls, rippling water, wharf, silhouette, serene atmosphere, dusk, evening glow, coastal landscape, seaside scenery.",
            "__assets__/examples/t2v/4.mp4", "56 - frames", None, "AnimateDiff (text to video)", 33, 50, 9.0],
        ["Dive into the depths of the ocean: explore vibrant coral reefs, mysterious underwater caves, and the mesmerizing creatures that call the sea home.",
            "__assets__/examples/t2v/5.mp4", "56 - frames", None, "SVD (image to video)", 33, 50, 9.0],
        ["Ants, beetles and centipede nest.",
            "__assets__/examples/t2v/6.mp4", "56 - frames", None, "SVD (image to video)", 33, 50, 9.0],
        ]

examples_2 = [
        ["Fishes swimming in ocean camera moving, cinematic.",
            "__assets__/examples/i2v/1.mp4", "56 - frames", "__assets__/fish.jpg", "SVD (image to video)", 33, 50, 9.0],
        ["A squirrel on a table full of big nuts.",
            "__assets__/examples/i2v/2.mp4", "56 - frames", "__assets__/squirrel.jpg", "SVD (image to video)", 33, 50, 9.0],
        ]

# --------------------------
# ----- Gradio-Demo UI -----
# --------------------------
with gr.Blocks() as demo:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
        <h1 style="font-weight: 900; font-size: 3rem; margin: 0rem">
            <a href="https://github.com/Picsart-AI-Research/StreamingT2V" style="color:blue;">StreamingT2V</a> 
        </h1>
        <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
        Roberto Henschel<sup>1*</sup>, Levon Khachatryan<sup>1*</sup>, Daniil Hayrapetyan<sup>1*</sup>, Hayk Poghosyan<sup>1</sup>, Vahram Tadevosyan<sup>1</sup>, Zhangyang Wang<sup>1,2</sup>, Shant Navasardyan<sup>1</sup>, Humphrey Shi<sup>1,3</sup>
        </h2>
        <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
        <sup>1</sup>Picsart AI Resarch (PAIR), <sup>2</sup>UT Austin, <sup>3</sup>SHI Labs @ Georgia Tech, Oregon & UIUC
        </h2>
        <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
        *Equal Contribution
        </h2>
        <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
        [<a href="https://arxiv.org/abs/2403.14773" style="color:blue;">arXiv</a>] 
        [<a href="https://github.com/Picsart-AI-Research/StreamingT2V" style="color:blue;">GitHub</a>]
        [<a href="https://streamingt2v.github.io/" style="color:blue;">Project page</a>]
        </h2>
        <h2 style="font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
        <b>StreamingT2V</b> is an advanced autoregressive technique that enables the creation of long videos featuring rich motion dynamics without any stagnation. 
        It ensures temporal consistency throughout the video, aligns closely with the descriptive text, and maintains high frame-level image quality. 
        Our demonstrations include successful examples of videos up to <b>1200 frames, spanning 2 minutes</b>, and can be extended for even longer durations. 
        Importantly, the effectiveness of StreamingT2V is not limited by the specific Text2Video model used, indicating that improvements in base models could yield even higher-quality videos.
        </h2>
        </div>
        """)

    if on_huggingspace:
        gr.HTML("""
        <p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
        <br/>
        <a href="https://huggingface.co/spaces/PAIR/StreamingT2V?duplicate=true">
        <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
        </p>""")

    with gr.Row():
        with gr.Column():
            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        num_frames = gr.Dropdown(["24 - frames", "32 - frames", "40 - frames", "48 - frames", "56 - frames", "80 - recommended to run on local GPUs", "240 - recommended to run on local GPUs", "600 - recommended to run on local GPUs", "1200 - recommended to run on local GPUs", "10000 - recommended to run on local GPUs"], label="Number of Video Frames", info="For >56 frames use local workstation!", value="24 - frames")
                    with gr.Row():
                        prompt_stage1 = gr.Textbox(label='Textual Prompt', placeholder=f"Ex: {default_prompt}")
                    with gr.Row():
                        image_stage1 = gr.Image(label='Image Prompt (first select Image-to-Video model from advanced options to enable image upload)', show_label=True, scale=1, show_download_button=False, interactive=False)
                with gr.Column():
                    video_stage1 = gr.Video(label='Long Video Preview', show_label=True, interactive=False, scale=2, show_download_button=True)
            with gr.Row():
                with gr.Row():
                    run_button_stage1 = gr.Button("Long Video Generation (faster preview)")
                with gr.Row():
                    run_button_stage2 = gr.Button("Long Video Generation")

            with gr.Row():
                with gr.Column():
                    with gr.Accordion('Advanced options', open=False):
                        model_name_stage1 = gr.Dropdown(
                            choices=["ModelScopeT2V (text to video)", "AnimateDiff (text to video)", "SVD (image to video)"],
                            label="Base Model",
                            value="ModelScopeT2V (text to video)"
                        )
                        model_name_stage2 = gr.Dropdown(
                            choices=["MS-Vid2Vid-XL"],
                            label="Enhancement Model",
                            value="MS-Vid2Vid-XL"
                        )
                        seed = gr.Slider(label='Seed', minimum=0, maximum=65536, value=33,step=1,)

                        t = gr.Slider(label="Timesteps", minimum=0, maximum=100, value=50, step=1,)
                        image_guidance = gr.Slider(label='Image guidance scale', minimum=1, maximum=10, value=9.0, step=1.0)

        with gr.Column():
            with gr.Row():
                video_stage2 = gr.Video(label='Long Video', show_label=True, interactive=False, height=588, show_download_button=True)

    model_name_stage1.change(fn=change_visibility, inputs=[model_name_stage1], outputs=image_stage1)

    inputs_t2v = [prompt_stage1, num_frames, image_stage1, model_name_stage1, model_name_stage2, seed, t, image_guidance]
    run_button_stage1.click(fn=generate, inputs=inputs_t2v, outputs=video_stage1,)

    inputs_v2v = [prompt_stage1, video_stage1, num_frames, image_stage1, model_name_stage1, model_name_stage2, seed, t, image_guidance]

    inputs_examples = [prompt_stage1, video_stage2, num_frames, image_stage1, model_name_stage1, seed, t, image_guidance]

    gr.HTML("""
        <h2>
            You can check the inference time for different number of frames
            <p style=" display: inline">
                <a href="https://github.com/Picsart-AI-Research/StreamingT2V/blob/main/README.md#inference-time" style="color:blue;" target="_blank">[here].</a> 
            </p>
        </h2>
        """)

    gr.Examples(examples=examples_1,
                inputs=inputs_examples,
                # outputs=[video_stage2],
                # fn=enhance,
                # run_on_click=False,
                # cache_examples=False,
                )

    gr.Examples(examples=examples_2,
                inputs=inputs_examples,
                # outputs=[video_stage2],
                # fn=enhance,
                # run_on_click=False,
                # cache_examples=False,
                # # preprocess=False,
                # # postprocess=True,
                )
    
    run_button_stage2.click(fn=enhance, inputs=inputs_v2v, outputs=video_stage2,)

    '''
    '''
    gr.HTML(
        """
        <div style="text-align: justify; max-width: 1200px; margin: 20px auto;">
        <h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
        <b>Version: v1.0</b>
        </h3>
        <h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
        <b>Caution</b>: 
        We would like the raise the awareness of users of this demo of its potential issues and concerns.
        Like previous large foundation models, StreamingT2V could be problematic in some cases, partially we use pretrained ModelScope, therefore StreamingT2V can Inherit Its Imperfections.
        So far, we keep all features available for research testing both to show the great potential of the StreamingT2V framework and to collect important feedback to improve the model in the future.
        We welcome researchers and users to report issues with the HuggingFace community discussion feature or email the authors.
        </h3>
        <h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
        <b>Biases and content acknowledgement</b>:
        Beware that StreamingT2V may output content that reinforces or exacerbates societal biases, as well as realistic faces, pornography, and violence. 
        StreamingT2V in this demo is meant only for research purposes.
        </h3>
        </div>
        """)


if on_huggingspace:
    demo.queue(max_size=10)
    demo.launch(debug=True)
else:
    demo.queue(api_open=False).launch(share=args.public_access)