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

Running on CUDA 😀

" elif hasattr(torch, "xpu") and torch.xpu.is_available(): DESCRIPTION += "\n

Running on XPU 🤓

" else: DESCRIPTION += "\n

Running on CPU 🥶 This demo does not work on CPU.

" 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()