import gradio import os import numpy as np import argparse import imageio import torch from einops import rearrange from diffusers import DDIMScheduler, AutoencoderKL from transformers import CLIPTextModel, CLIPTokenizer # from annotator.canny import CannyDetector # from annotator.openpose import OpenposeDetector # from annotator.midas import MidasDetector # import sys # sys.path.insert(0, ".") from huggingface_hub import hf_hub_download import controlnet_aux from controlnet_aux import OpenposeDetector, CannyDetector, MidasDetector from controlnet_aux.open_pose.body import Body from models.pipeline_controlvideo import ControlVideoPipeline from models.util import save_videos_grid, read_video, get_annotation from models.unet import UNet3DConditionModel from models.controlnet import ControlNetModel3D from models.RIFE.IFNet_HDv3 import IFNet device = "cuda" sd_path = "runwayml/stable-diffusion-v1-5" inter_path = "checkpoints/flownet.pkl" controlnet_dict = { "pose": "lllyasviel/control_v11p_sd15_openpose", "depth": "lllyasviel/control_v11p_sd15_canny", "canny": "lllyasviel/control_v11f1p_sd15_depth", } controlnet_parser_dict = { "pose": OpenposeDetector, "depth": MidasDetector, "canny": CannyDetector, } POS_PROMPT = " ,best quality, extremely detailed, HD, ultra-realistic, 8K, HQ, masterpiece, trending on artstation, art, smooth" NEG_PROMPT = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer difits, cropped, worst quality, low quality, deformed body, bloated, ugly, unrealistic" def get_args(): parser = argparse.ArgumentParser() parser.add_argument("--prompt", type=str, required=True, help="Text description of target video") parser.add_argument("--video_path", type=str, required=True, help="Path to a source video") parser.add_argument("--output_path", type=str, default="./outputs", help="Directory of output") parser.add_argument("--condition", type=str, default="depth", help="Condition of structure sequence") parser.add_argument("--video_length", type=int, default=15, help="Length of synthesized video") parser.add_argument("--height", type=int, default=512, help="Height of synthesized video, and should be a multiple of 32") parser.add_argument("--width", type=int, default=512, help="Width of synthesized video, and should be a multiple of 32") parser.add_argument("--smoother_steps", nargs='+', default=[19, 20], type=int, help="Timesteps at which using interleaved-frame smoother") parser.add_argument("--is_long_video", action='store_true', help="Whether to use hierarchical sampler to produce long video") parser.add_argument("--seed", type=int, default=42, help="Random seed of generator") args = parser.parse_args() return args def infer(prompt, video_path, condition, video_length, is_long_video): #args = get_args() #os.makedirs(args.output_path, exist_ok=True) # Height and width should be a multiple of 32 output_path = "" height, width = 512 height = (height // 32) * 32 width = (width // 32) * 32 smoother_steps = [19, 20] is_long_video = False seed = 42 if condition == "pose": pretrained_model_or_path = "lllyasviel/ControlNet" body_model_path = hf_hub_download(pretrained_model_or_path, "annotator/ckpts/body_pose_model.pth", cache_dir="checkpoints") body_estimation = Body(body_model_path) annotator = controlnet_parser_dict[condition](body_estimation) else: annotator = controlnet_parser_dict[condition]() tokenizer = CLIPTokenizer.from_pretrained(sd_path, subfolder="tokenizer") text_encoder = CLIPTextModel.from_pretrained(sd_path, subfolder="text_encoder").to(dtype=torch.float16) vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae").to(dtype=torch.float16) unet = UNet3DConditionModel.from_pretrained_2d(sd_path, subfolder="unet").to(dtype=torch.float16) controlnet = ControlNetModel3D.from_pretrained_2d(controlnet_dict[args.condition]).to(dtype=torch.float16) interpolater = IFNet(ckpt_path=inter_path).to(dtype=torch.float16) scheduler=DDIMScheduler.from_pretrained(sd_path, subfolder="scheduler") pipe = ControlVideoPipeline( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, controlnet=controlnet, interpolater=interpolater, scheduler=scheduler, ) pipe.enable_vae_slicing() pipe.enable_xformers_memory_efficient_attention() pipe.to(device) generator = torch.Generator(device="cuda") generator.manual_seed(seed) # Step 1. Read a video video = read_video(video_path=video_path, video_length=video_length, width=width, height=height) # Save source video original_pixels = rearrange(video, "(b f) c h w -> b c f h w", b=1) save_videos_grid(original_pixels, os.path.join(output_path, "source_video.mp4"), rescale=True) # Step 2. Parse a video to conditional frames pil_annotation = get_annotation(video, annotator) if condition == "depth" and controlnet_aux.__version__ == '0.0.1': pil_annotation = [pil_annot[0] for pil_annot in pil_annotation] # Save condition video video_cond = [np.array(p).astype(np.uint8) for p in pil_annotation] imageio.mimsave(os.path.join(output_path, f"{condition}_condition.mp4"), video_cond, fps=8) # Reduce memory (optional) del annotator; torch.cuda.empty_cache() # Step 3. inference if is_long_video: window_size = int(np.sqrt(video_length)) sample = pipe.generate_long_video(prompt + POS_PROMPT, video_length=video_length, frames=pil_annotation, num_inference_steps=50, smooth_steps=args.smoother_steps, window_size=window_size, generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT, width=width, height=height ).videos else: sample = pipe(prompt + POS_PROMPT, video_length=video_length, frames=pil_annotation, num_inference_steps=50, smooth_steps=args.smoother_steps, generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT, width=width, height=height ).videos save_videos_grid(sample, f"{output_path}/{prompt}.mp4") return f"{output_path}/{prompt}.mp4" with gr.Blocks() as demo: with gr.Column(): prompt = gr.Textbox(label="prompt") video_path = gr.Video(source="upload", type="filepath") condition = gr.Textbox(label="Condition", value="depth") video_length = gr.Slider(label="video length", minimum=1, maximum=15, step=1, value=2) seed = gr.Number(label="seed", valie=42) submit_btn = gr.Button("Submit") video_res = gr.Video(label="result") submit_btn.click(fn=infer, inputs=[prompt, video_path, condition, video_length, seed, ], outputs=[video_res])