-
Notifications
You must be signed in to change notification settings - Fork 56
/
inference.py
153 lines (129 loc) · 7.65 KB
/
inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
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
import torchvision
from controlnet_aux.processor import Processor
from models.pipeline_controlvideo import ControlVideoPipeline
from models.util import save_videos_grid, read_video
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_version = {
"v10":{
"openpose": "checkpoints/sd-controlnet-openpose",
"depth_midas": "checkpoints/sd-controlnet-depth",
"canny": "checkpoints/sd-controlnet-canny",
},
"v11": {
"softedge_pidinet": "checkpoints/control_v11p_sd15_softedge",
"softedge_pidsafe": "checkpoints/control_v11p_sd15_softedge",
"softedge_hed": "checkpoints/control_v11p_sd15_softedge",
"softedge_hedsafe": "checkpoints/control_v11p_sd15_softedge",
"scribble_hed": "checkpoints/control_v11p_sd15_scribble",
"scribble_pidinet": "checkpoints/control_v11p_sd15_scribble",
"lineart_anime": "checkpoints/control_v11p_sd15_lineart_anime",
"lineart_coarse": "checkpoints/control_v11p_sd15_lineart",
"lineart_realistic": "checkpoints/control_v11p_sd15_lineart",
"depth_midas": "checkpoints/control_v11f1p_sd15_depth",
"depth_leres": "checkpoints/control_v11f1p_sd15_depth",
"depth_leres++": "checkpoints/control_v11f1p_sd15_depth",
"depth_zoe": "checkpoints/control_v11f1p_sd15_depth",
"canny": "checkpoints/control_v11p_sd15_canny",
"openpose": "checkpoints/control_v11p_sd15_openpose",
"openpose_face": "checkpoints/control_v11p_sd15_openpose",
"openpose_faceonly": "checkpoints/control_v11p_sd15_openpose",
"openpose_full": "checkpoints/control_v11p_sd15_openpose",
"openpose_hand": "checkpoints/control_v11p_sd15_openpose",
"normal_bae": "checkpoints/control_v11p_sd15_normalbae"
}
}
# load processor from processor_id
# options are:
# ["canny", "depth_leres", "depth_leres++", "depth_midas", "depth_zoe", "lineart_anime",
# "lineart_coarse", "lineart_realistic", "mediapipe_face", "mlsd", "normal_bae", "normal_midas",
# "openpose", "openpose_face", "openpose_faceonly", "openpose_full", "openpose_hand",
# "scribble_hed, "scribble_pidinet", "shuffle", "softedge_hed", "softedge_hedsafe",
# "softedge_pidinet", "softedge_pidsafe"]
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 digits, 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")
parser.add_argument("--version", type=str, default='v10', choices=["v10", "v11"], help="Version of ControlNet")
parser.add_argument("--frame_rate", type=int, default=None, help="The frame rate of loading input video. Default rate is computed according to video length.")
parser.add_argument("--temp_video_name", type=str, default=None, help="Default video name")
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
processor = Processor(args.condition)
controlnet_dict = controlnet_dict_version[args.version]
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, frame_rate=args.frame_rate)
# 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
t2i_transform = torchvision.transforms.ToPILImage()
pil_annotation = []
for frame in video:
pil_frame = t2i_transform(frame)
pil_annotation.append(processor(pil_frame, to_pil=True))
# 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=8)
# Reduce memory (optional)
del processor; 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=50, 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=50, smooth_steps=args.smoother_steps,
generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT,
width=args.width, height=args.height
).videos
args.temp_video_name = args.prompt if args.temp_video_name is None else args.temp_video_name
save_videos_grid(sample, f"{args.output_path}/{args.temp_video_name}.mp4")