EX4L/T2V-Turbo

1import os2import uuid3from omegaconf import OmegaConf4import spaces5 6import random7 8import imageio9import torch10import torchvision11import gradio as gr12import numpy as np13 14from gradio.components import Textbox, Video15from huggingface_hub import hf_hub_download16 17from utils.common_utils import load_model_checkpoint18from utils.utils import instantiate_from_config19from scheduler.t2v_turbo_scheduler import T2VTurboScheduler20from pipeline.t2v_turbo_vc2_pipeline import T2VTurboVC2Pipeline21 22DESCRIPTION = """# T2V-Turbo 🚀23 24Our model is distilled from [VideoCrafter2](https://ailab-cvc.github.io/videocrafter2/).25 26T2V-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).27 28T2V-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)29 30T2V-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.31 32Moreover, T2V-Turbo-v2 supports to distill motion priors from the training videos. 33 34[Project page for T2V-Turbo](https://t2v-turbo.github.io) 🥳35 36[Project page for T2V-Turbo-v2](https://t2v-turbo-v2.github.io) 🤓37"""38if torch.cuda.is_available():39    DESCRIPTION += "\n<p>Running on CUDA 😀</p>"40elif hasattr(torch, "xpu") and torch.xpu.is_available():41    DESCRIPTION += "\n<p>Running on XPU 🤓</p>"42else:43    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"44 45MAX_SEED = np.iinfo(np.int32).max46 47 48def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:49    if randomize_seed:50        seed = random.randint(0, MAX_SEED)51    return seed52 53 54def save_video(video_array, video_save_path, fps: int = 16):55    video = video_array.detach().cpu()56    video = torch.clamp(video.float(), -1.0, 1.0)57    video = video.permute(1, 0, 2, 3)  # t,c,h,w58    video = (video + 1.0) / 2.059    video = (video * 255).to(torch.uint8).permute(0, 2, 3, 1)60 61    torchvision.io.write_video(62        video_save_path, video, fps=fps, video_codec="h264", options={"crf": "10"}63    )64 65example_txt = [66    "An astronaut riding a horse.",67    "Darth vader surfing in waves.",68    "light wind, feathers moving, she moves her gaze, 4k",69    "a girl floating underwater.",70    "Pikachu snowboarding.",71    "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",72    "A musician strums his guitar, serenading the moonlit night.",73]74 75examples = [[i, 7.5, 0.5, 16, 16, 0, True, "bf16"] for i in example_txt]76 77@spaces.GPU(duration=120)78@torch.inference_mode()79def generate(80    prompt: str,81    guidance_scale: float = 7.5,82    percentage: float = 0.5,83    num_inference_steps: int = 4,84    num_frames: int = 16,85    seed: int = 0,86    randomize_seed: bool = False,87    param_dtype="bf16",88    motion_gs: float = 0.05,89    fps: int = 8,90):91 92    seed = randomize_seed_fn(seed, randomize_seed)93    torch.manual_seed(seed)94 95    if param_dtype == "bf16":96        dtype = torch.bfloat1697        unet.dtype = torch.bfloat1698    elif param_dtype == "fp16":99        dtype = torch.float16100        unet.dtype = torch.float16101    elif param_dtype == "fp32":102        dtype = torch.float32103        unet.dtype = torch.float32104    else:105        raise ValueError(f"Unknown dtype: {param_dtype}")106 107    pipeline.unet.to(device, dtype)108    pipeline.text_encoder.to(device, dtype)109    pipeline.vae.to(device, dtype)110    pipeline.to(device, dtype)111 112    result = pipeline(113        prompt=prompt,114        frames=num_frames,115        fps=fps,116        guidance_scale=guidance_scale,117        motion_gs=motion_gs,118        use_motion_cond=True,119        percentage=percentage,120        num_inference_steps=num_inference_steps,121        lcm_origin_steps=200,122        num_videos_per_prompt=1,123    )124 125    torch.cuda.empty_cache()126    tmp_save_path = "tmp.mp4"127    root_path = "./videos/"128    os.makedirs(root_path, exist_ok=True)129    video_save_path = os.path.join(root_path, tmp_save_path)130 131    save_video(result[0], video_save_path, fps=fps)132    display_model_info = f"Video size: {num_frames}x320x512, Sampling Step: {num_inference_steps}, Guidance Scale: {guidance_scale}"133    return video_save_path, prompt, display_model_info, seed134 135 136block_css = """137#buttons button {138    min-width: min(120px,100%);139}140"""141 142 143if __name__ == "__main__":144    device = torch.device("cuda:0")145 146    config = OmegaConf.load("configs/inference_t2v_512_v2.0.yaml")147    model_config = config.pop("model", OmegaConf.create())148    pretrained_t2v = instantiate_from_config(model_config)149 150    pretrained_path = hf_hub_download("VideoCrafter/VideoCrafter2", filename="model.ckpt")151    pretrained_t2v = load_model_checkpoint(pretrained_t2v, pretrained_path)152    153    unet_config = model_config["params"]["unet_config"]154    unet_config["params"]["use_checkpoint"] = False155    unet_config["params"]["time_cond_proj_dim"] = 256156    unet_config["params"]["motion_cond_proj_dim"] = 256157 158    unet = instantiate_from_config(unet_config)159 160    unet_path = hf_hub_download(repo_id="jiachenli-ucsb/T2V-Turbo-v2", filename="unet_mg.pt")161    unet.load_state_dict(torch.load(unet_path, map_location=device))162    unet.eval()163 164    pretrained_t2v.model.diffusion_model = unet165    scheduler = T2VTurboScheduler(166        linear_start=model_config["params"]["linear_start"],167        linear_end=model_config["params"]["linear_end"],168    )169    pipeline = T2VTurboVC2Pipeline(pretrained_t2v, scheduler, model_config)170    pipeline.to(device)171 172    demo = gr.Interface(173        fn=generate,174        inputs=[175            Textbox(label="", placeholder="Please enter your prompt. \n"),176            gr.Slider(177                label="Guidance scale",178                minimum=2,179                maximum=14,180                step=0.1,181                value=7.5,182            ),183            gr.Slider(184                label="Percentage of steps to apply motion guidance (v2 w/ MG only)",185                minimum=0.0,186                maximum=0.5,187                step=0.05,188                value=0.5,189            ),190            gr.Slider(191                label="Number of inference steps",192                minimum=4,193                maximum=50,194                step=1,195                value=16,196            ),197            gr.Slider(198                label="Number of Video Frames",199                minimum=16,200                maximum=48,201                step=8,202                value=16,203            ),204            gr.Slider(205                label="Seed",206                minimum=0,207                maximum=MAX_SEED,208                step=1,209                value=0,210                randomize=True,211            ),212            gr.Checkbox(label="Randomize seed", value=True),213            gr.Radio(214                ["bf16", "fp16", "fp32"],215                label="torch.dtype",216                value="bf16",217                interactive=True,218                info="Dtype for inference. Default is bf16.",219            )220        ],221        outputs=[222            gr.Video(label="Generated Video", width=512, height=320, interactive=False, autoplay=True),223            Textbox(label="input prompt"),224            Textbox(label="model info"),225            gr.Slider(label="seed"),226        ],227        description=DESCRIPTION,228        theme=gr.themes.Default(),229        css=block_css,230        examples=examples,231        cache_examples=False,232        concurrency_limit=10,233    )234    demo.launch()235