ControlVideo / inference.py
fffiloni's picture
Update inference.py
8a72154
raw
history blame
6.68 kB
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 = "checkpoints/stable-diffusion-v1-5"
inter_path = "checkpoints/flownet.pkl"
controlnet_dict = {
"pose": "checkpoints/sd-controlnet-openpose",
"depth": "checkpoints/sd-controlnet-depth",
"canny": "checkpoints/sd-controlnet-canny",
}
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("--inference_steps", type=int, default=25, help="Number of inference steps")
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("--fps", type=int, default=8, help="FPS for final output")
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")
parser.add_argument("--temp_chunk_path", type=str, required=True, help="Path to temporary chunks")
args = parser.parse_args()
return args
if __name__ == "__main__":
args = get_args()
os.makedirs(args.output_path, exist_ok=True)
# Height and width should be a multiple of 32
args.height = (args.height // 32) * 32
args.width = (args.width // 32) * 32
if args.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[args.condition](body_estimation)
else:
annotator = controlnet_parser_dict[args.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(args.seed)
# Step 1. Read a video
video = read_video(video_path=args.video_path, video_length=args.video_length, width=args.width, height=args.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(args.output_path, "source_video.mp4"), rescale=True)
# Step 2. Parse a video to conditional frames
pil_annotation = get_annotation(video, annotator)
if args.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(args.output_path, f"{args.condition}_condition.mp4"), video_cond, fps=args.fps)
# Reduce memory (optional)
#del annotator; torch.cuda.empty_cache()
# Step 3. inference
if args.is_long_video:
window_size = int(np.sqrt(args.video_length))
sample = pipe.generate_long_video(args.prompt + POS_PROMPT, video_length=args.video_length, frames=pil_annotation,
num_inference_steps=args.inference_steps, smooth_steps=args.smoother_steps, window_size=window_size,
generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT,
width=args.width, height=args.height
).videos
else:
sample = pipe(args.prompt + POS_PROMPT, video_length=args.video_length, frames=pil_annotation,
num_inference_steps=args.inference_steps, smooth_steps=args.smoother_steps,
generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT,
width=args.width, height=args.height
).videos
save_videos_grid(sample, f"{args.output_path}/{args.temp_chunk_path}.mp4")
del pipe
torch.cuda.empty_cache()