Ji4chenLi
update app.py to link space and models
e13087b
import os
import uuid
from omegaconf import OmegaConf
import spaces
import random
import imageio
import torch
import torchvision
import gradio as gr
import numpy as np
from gradio.components import Textbox, Video
from huggingface_hub import hf_hub_download
from utils.common_utils import load_model_checkpoint
from utils.utils import instantiate_from_config
from scheduler.t2v_turbo_scheduler import T2VTurboScheduler
from pipeline.t2v_turbo_vc2_pipeline import T2VTurboVC2Pipeline
DESCRIPTION = """# T2V-Turbo 🚀
Our model is distilled from [VideoCrafter2](https://ailab-cvc.github.io/videocrafter2/).
T2V-Turbo learns a LoRA on top of the base model by aligning to the reward feedback from [HPSv2.1](https://github.com/tgxs002/HPSv2/tree/master) and [InternVid2 Stage 2 Model](https://huggingface.co/OpenGVLab/InternVideo2-Stage2_1B-224p-f4).
T2V-Turbo-v2 optimizes the training techniques by finetuning the full base model and further aligns to [CLIPScore](https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K)
T2V-Turbo trains on pure WebVid-10M data, whereas T2V-Turbo-v2 carufully optimizes different learning objectives with a mixutre of VidGen-1M and WebVid-10M data.
Moreover, T2V-Turbo-v2 supports to distill motion priors from the training videos.
[Project page for T2V-Turbo](https://t2v-turbo.github.io) 🥳
[Project page for T2V-Turbo-v2](https://t2v-turbo-v2.github.io) 🤓
"""
if torch.cuda.is_available():
DESCRIPTION += "\n<p>Running on CUDA 😀</p>"
elif hasattr(torch, "xpu") and torch.xpu.is_available():
DESCRIPTION += "\n<p>Running on XPU 🤓</p>"
else:
DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
MAX_SEED = np.iinfo(np.int32).max
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
if randomize_seed:
seed = random.randint(0, MAX_SEED)
return seed
def save_video(video_array, video_save_path, fps: int = 16):
video = video_array.detach().cpu()
video = torch.clamp(video.float(), -1.0, 1.0)
video = video.permute(1, 0, 2, 3) # t,c,h,w
video = (video + 1.0) / 2.0
video = (video * 255).to(torch.uint8).permute(0, 2, 3, 1)
torchvision.io.write_video(
video_save_path, video, fps=fps, video_codec="h264", options={"crf": "10"}
)
example_txt = [
"An astronaut riding a horse.",
"Darth vader surfing in waves.",
"light wind, feathers moving, she moves her gaze, 4k",
"a girl floating underwater.",
"Pikachu snowboarding.",
"Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
"A musician strums his guitar, serenading the moonlit night.",
]
examples = [[i, 7.5, 0.5, 16, 16, 0, True, "bf16"] for i in example_txt]
@spaces.GPU(duration=120)
@torch.inference_mode()
def generate(
prompt: str,
guidance_scale: float = 7.5,
percentage: float = 0.5,
num_inference_steps: int = 4,
num_frames: int = 16,
seed: int = 0,
randomize_seed: bool = False,
param_dtype="bf16",
motion_gs: float = 0.05,
fps: int = 8,
):
seed = randomize_seed_fn(seed, randomize_seed)
torch.manual_seed(seed)
if param_dtype == "bf16":
dtype = torch.bfloat16
unet.dtype = torch.bfloat16
elif param_dtype == "fp16":
dtype = torch.float16
unet.dtype = torch.float16
elif param_dtype == "fp32":
dtype = torch.float32
unet.dtype = torch.float32
else:
raise ValueError(f"Unknown dtype: {param_dtype}")
pipeline.unet.to(device, dtype)
pipeline.text_encoder.to(device, dtype)
pipeline.vae.to(device, dtype)
pipeline.to(device, dtype)
result = pipeline(
prompt=prompt,
frames=num_frames,
fps=fps,
guidance_scale=guidance_scale,
motion_gs=motion_gs,
use_motion_cond=True,
percentage=percentage,
num_inference_steps=num_inference_steps,
lcm_origin_steps=200,
num_videos_per_prompt=1,
)
torch.cuda.empty_cache()
tmp_save_path = "tmp.mp4"
root_path = "./videos/"
os.makedirs(root_path, exist_ok=True)
video_save_path = os.path.join(root_path, tmp_save_path)
save_video(result[0], video_save_path, fps=fps)
display_model_info = f"Video size: {num_frames}x320x512, Sampling Step: {num_inference_steps}, Guidance Scale: {guidance_scale}"
return video_save_path, prompt, display_model_info, seed
block_css = """
#buttons button {
min-width: min(120px,100%);
}
"""
if __name__ == "__main__":
device = torch.device("cuda:0")
config = OmegaConf.load("configs/inference_t2v_512_v2.0.yaml")
model_config = config.pop("model", OmegaConf.create())
pretrained_t2v = instantiate_from_config(model_config)
pretrained_path = hf_hub_download("VideoCrafter/VideoCrafter2", filename="model.ckpt")
pretrained_t2v = load_model_checkpoint(pretrained_t2v, pretrained_path)
unet_config = model_config["params"]["unet_config"]
unet_config["params"]["use_checkpoint"] = False
unet_config["params"]["time_cond_proj_dim"] = 256
unet_config["params"]["motion_cond_proj_dim"] = 256
unet = instantiate_from_config(unet_config)
unet_path = hf_hub_download(repo_id="jiachenli-ucsb/T2V-Turbo-v2", filename="unet_mg.pt")
unet.load_state_dict(torch.load(unet_path, map_location=device))
unet.eval()
pretrained_t2v.model.diffusion_model = unet
scheduler = T2VTurboScheduler(
linear_start=model_config["params"]["linear_start"],
linear_end=model_config["params"]["linear_end"],
)
pipeline = T2VTurboVC2Pipeline(pretrained_t2v, scheduler, model_config)
pipeline.to(device)
demo = gr.Interface(
fn=generate,
inputs=[
Textbox(label="", placeholder="Please enter your prompt. \n"),
gr.Slider(
label="Guidance scale",
minimum=2,
maximum=14,
step=0.1,
value=7.5,
),
gr.Slider(
label="Percentage of steps to apply motion guidance (v2 w/ MG only)",
minimum=0.0,
maximum=0.5,
step=0.05,
value=0.5,
),
gr.Slider(
label="Number of inference steps",
minimum=4,
maximum=50,
step=1,
value=16,
),
gr.Slider(
label="Number of Video Frames",
minimum=16,
maximum=48,
step=8,
value=16,
),
gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
randomize=True,
),
gr.Checkbox(label="Randomize seed", value=True),
gr.Radio(
["bf16", "fp16", "fp32"],
label="torch.dtype",
value="bf16",
interactive=True,
info="Dtype for inference. Default is bf16.",
)
],
outputs=[
gr.Video(label="Generated Video", width=512, height=320, interactive=False, autoplay=True),
Textbox(label="input prompt"),
Textbox(label="model info"),
gr.Slider(label="seed"),
],
description=DESCRIPTION,
theme=gr.themes.Default(),
css=block_css,
examples=examples,
cache_examples=False,
concurrency_limit=10,
)
demo.launch()