base/app.py

import gradio as gr
from text_to_video import model_t2v_fun, setup_seed
from omegaconf import OmegaConf
import torch
import imageio
import os
import cv2
import pandas as pd
import torchvision
import random
from models import get_models

from pipelines.pipeline_videogen import VideoGenPipeline
from download import find_model
from diffusers.schedulers import DDIMScheduler, DDPMScheduler, PNDMScheduler, EulerDiscreteScheduler
from diffusers.models import AutoencoderKL
from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection

config_path = "./base/configs/sample.yaml"
args = OmegaConf.load("./base/configs/sample.yaml")
device = "cpu"  # Force CPU usage

css = """
h1 {
  text-align: center;
}
#component-0 {
  max-width: 730px;
  margin: auto;
}
"""

sd_path = args.pretrained_path
unet = get_models(args, sd_path).to(device, dtype=torch.float32)  # Use float32 for CPU
state_dict = find_model("./pretrained_models/lavie_base.pt")
unet.load_state_dict(state_dict)
vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae", torch_dtype=torch.float32).to(device)  # Use float32 for CPU
tokenizer_one = CLIPTokenizer.from_pretrained(sd_path, subfolder="tokenizer")
text_encoder_one = CLIPTextModel.from_pretrained(sd_path, subfolder="text_encoder", torch_dtype=torch.float32).to(device)  # Use float32 for CPU
unet.eval()
vae.eval()
text_encoder_one.eval()

def infer(prompt, seed_inp, ddim_steps, cfg, infer_type):
    if seed_inp != -1:
        setup_seed(seed_inp)
    else:
        seed_inp = random.choice(range(10000000))
        setup_seed(seed_inp)
    if infer_type == 'ddim':
        scheduler = DDIMScheduler.from_pretrained(sd_path, 
                                               subfolder="scheduler",
                                               beta_start=args.beta_start, 
                                               beta_end=args.beta_end, 
                                               beta_schedule=args.beta_schedule)
    elif infer_type == 'eulerdiscrete':
        scheduler = EulerDiscreteScheduler.from_pretrained(sd_path,
                                               subfolder="scheduler",
                                               beta_start=args.beta_start,
                                               beta_end=args.beta_end,
                                               beta_schedule=args.beta_schedule)
    elif infer_type == 'ddpm':
        scheduler = DDPMScheduler.from_pretrained(sd_path,
                                              subfolder="scheduler",
                                              beta_start=args.beta_start,
                                              beta_end=args.beta_end,
                                              beta_schedule=args.beta_schedule)
    model = VideoGenPipeline(vae=vae, text_encoder=text_encoder_one, tokenizer=tokenizer_one, scheduler=scheduler, unet=unet)
    model.to(device)
    # Disable xformers for CPU
    # if device == "cuda":
    #     model.enable_xformers_memory_efficient_attention()
    videos = model(prompt, video_length=8, height=160, width=256, num_inference_steps=ddim_steps, guidance_scale=cfg).video  # Reduced resolution and length
    if not os.path.exists(args.output_folder):
        os.mkdir(args.output_folder)
    torchvision.io.write_video(args.output_folder + prompt[0:30].replace(' ', '_') + '-' + str(seed_inp) + '-' + str(ddim_steps) + '-' + str(cfg) + '-.mp4', videos[0], fps=4)  # Reduced FPS

    return args.output_folder + prompt[0:30].replace(' ', '_') + '-' + str(seed_inp) + '-' + str(ddim_steps) + '-' + str(cfg) + '-.mp4'

title = """
    <div style="text-align: center; max-width: 700px; margin: 0 auto;">
        <div
        style="
            display: inline-flex;
            align-items: center;
            gap: 0.8rem;
            font-size: 1.75rem;
        "
        >
        <h1 style="font-weight: 900; margin-bottom: 7px; margin-top: 5px;">
            Intern·Vchitect (Text-to-Video)
        </h1>
        </div>
        <p style="margin-bottom: 10px; font-size: 94%">
        Apply Intern·Vchitect to generate a video 
        </p>
    </div>
"""

with gr.Blocks(css='style.css') as demo:
    gr.Markdown("<font color=red size=10><center>LaVie: Text-to-Video generation</center></font>")
    gr.Markdown(
        """<div style="text-align:center">
        [<a href="https://arxiv.org/abs/2309.15103">Arxiv Report</a>] | [<a href="https://vchitect.github.io/LaVie-project/">Project Page</a>] | [<a href="https://github.com/Vchitect/LaVie">Github</a>]</div>
        """
    )
    with gr.Column():
        with gr.Row(elem_id="col-container"):
            with gr.Column():
                prompt = gr.Textbox(value="a corgi walking in the park at sunrise, oil painting style", label="Prompt", placeholder="enter prompt", show_label=True, elem_id="prompt-in", min_width=200, lines=2)
                infer_type = gr.Dropdown(['ddpm', 'ddim', 'eulerdiscrete'], label='infer_type', value='ddim')
                ddim_steps = gr.Slider(label='Steps', minimum=50, maximum=300, value=50, step=1)
                seed_inp = gr.Slider(value=-1, label="seed (for random generation, use -1)", show_label=True, minimum=-1, maximum=2147483647)
                cfg = gr.Number(label="guidance_scale", value=7.5)

            with gr.Column():
                submit_btn = gr.Button("Generate video")
                video_out = gr.Video(label="Video result", elem_id="video-output")

            inputs = [prompt, seed_inp, ddim_steps, cfg, infer_type]
            outputs = [video_out]

        ex = gr.Examples(
            examples=[['a corgi walking in the park at sunrise, oil painting style', 400, 50, 7, 'ddim'],
                      ['a cute teddy bear reading a book in the park, oil painting style, high quality', 700, 50, 7, 'ddim'],
                      ['an epic tornado attacking above a glowing city at night, the tornado is made of smoke, highly detailed', 230, 50, 7, 'ddim'],
                      ['a jar filled with fire, 4K video, 3D rendered, well-rendered', 400, 50, 7, 'ddim'],
                      ['a teddy bear walking in the park, oil painting style, high quality', 400, 50, 7, 'ddim'],
                      ['a teddy bear walking on the street, 2k, high quality', 100, 50, 7, 'ddim'],
                      ['a panda taking a selfie, 2k, high quality', 400, 50, 7, 'ddim'],
                      ['a polar bear playing drum kit in NYC Times Square, 4k, high resolution', 400, 50, 7, 'ddim'],
                      ['jungle river at sunset, ultra quality', 400, 50, 7, 'ddim'],
                      ['a shark swimming in clear Carribean ocean, 2k, high quality', 400, 50, 7, 'ddim'],
                      ['A steam train moving on a mountainside by Vincent van Gogh', 230, 50, 7, 'ddim'],
                      ['a confused grizzly bear in calculus class', 1000, 50, 7, 'ddim']],
            fn=infer,
            inputs=[prompt, seed_inp, ddim_steps, cfg, infer_type],
            outputs=[video_out],
            cache_examples=True,
        )
        ex.dataset.headers = [""]

    submit_btn.click(infer, inputs, outputs)

demo.queue(max_size=12, api_open=False).launch(show_api=False)