import argparse, os, sys, glob
import datetime, time
from omegaconf import OmegaConf
import torch
from decord import VideoReader, cpu
import torchvision
from pytorch_lightning import seed_everything
from lvdm.samplers.ddim import DDIMSampler
from lvdm.utils.common_utils import instantiate_from_config
from lvdm.utils.saving_utils import tensor_to_mp4
def get_filelist(data_dir, ext='*'):
file_list = glob.glob(os.path.join(data_dir, '*.%s'%ext))
file_list.sort()
return file_list
def load_model_checkpoint(model, ckpt, adapter_ckpt=None):
print('>>> Loading checkpoints ...')
if adapter_ckpt:
## main model
state_dict = torch.load(ckpt, map_location="cpu")
if "state_dict" in list(state_dict.keys()):
state_dict = state_dict["state_dict"]
model.load_state_dict(state_dict, strict=False)
print('@model checkpoint loaded.')
## adapter
state_dict = torch.load(adapter_ckpt, map_location="cpu")
if "state_dict" in list(state_dict.keys()):
state_dict = state_dict["state_dict"]
model.adapter.load_state_dict(state_dict, strict=True)
print('@adapter checkpoint loaded.')
else:
state_dict = torch.load(ckpt, map_location="cpu")
if "state_dict" in list(state_dict.keys()):
state_dict = state_dict["state_dict"]
model.load_state_dict(state_dict, strict=True)
print('@model checkpoint loaded.')
return model
def load_prompts(prompt_file):
f = open(prompt_file, 'r')
prompt_list = []
for idx, line in enumerate(f.readlines()):
l = line.strip()
if len(l) != 0:
prompt_list.append(l)
f.close()
return prompt_list
def load_video(filepath, frame_stride, video_size=(256,256), video_frames=16):
vidreader = VideoReader(filepath, ctx=cpu(0), width=video_size[1], height=video_size[0])
max_frames = len(vidreader)
temp_stride = max_frames // video_frames if frame_stride == -1 else frame_stride
if temp_stride * (video_frames-1) >= max_frames:
print(f'Warning: default frame stride is used because the input video clip {max_frames} is not long enough.')
temp_stride = max_frames // video_frames
frame_indices = [temp_stride*i for i in range(video_frames)]
frames = vidreader.get_batch(frame_indices)
## [t,h,w,c] -> [c,t,h,w]
frame_tensor = torch.tensor(frames.asnumpy()).permute(3, 0, 1, 2).float()
frame_tensor = (frame_tensor / 255. - 0.5) * 2
return frame_tensor
def save_results(prompt, samples, inputs, filename, realdir, fakedir, fps=10):
## save prompt
prompt = prompt[0] if isinstance(prompt, list) else prompt
path = os.path.join(realdir, "%s.txt"%filename)
with open(path, 'w') as f:
f.write(f'{prompt}')
f.close()
## save video
videos = [inputs, samples]
savedirs = [realdir, fakedir]
for idx, video in enumerate(videos):
if video is None:
continue
# b,c,t,h,w
video = video.detach().cpu()
video = torch.clamp(video.float(), -1., 1.)
n = video.shape[0]
video = video.permute(2, 0, 1, 3, 4) # t,n,c,h,w
frame_grids = [torchvision.utils.make_grid(framesheet, nrow=int(n)) for framesheet in video] #[3, 1*h, n*w]
grid = torch.stack(frame_grids, dim=0) # stack in temporal dim [t, 3, n*h, w]
grid = (grid + 1.0) / 2.0
grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1)
path = os.path.join(savedirs[idx], "%s.mp4"%filename)
torchvision.io.write_video(path, grid, fps=fps, video_codec='h264', options={'crf': '10'})
def adapter_guided_synthesis(model, prompts, videos, noise_shape, n_samples=1, ddim_steps=50, ddim_eta=1., \
unconditional_guidance_scale=1.0, unconditional_guidance_scale_temporal=None, **kwargs):
ddim_sampler = DDIMSampler(model)
batch_size = noise_shape[0]
## get condition embeddings (support single prompt only)
if isinstance(prompts, str):
prompts = [prompts]
cond = model.get_learned_conditioning(prompts)
if unconditional_guidance_scale != 1.0:
prompts = batch_size * [""]
uc = model.get_learned_conditioning(prompts)
else:
uc = None
## adapter features: process in 2D manner
b, c, t, h, w = videos.shape
extra_cond = model.get_batch_depth(videos, (h,w))
features_adapter = model.get_adapter_features(extra_cond)
batch_variants = []
for _ in range(n_samples):
if ddim_sampler is not None:
samples, _ = ddim_sampler.sample(S=ddim_steps,
conditioning=cond,
batch_size=noise_shape[0],
shape=noise_shape[1:],
verbose=False,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_conditioning=uc,
eta=ddim_eta,
temporal_length=noise_shape[2],
conditional_guidance_scale_temporal=unconditional_guidance_scale_temporal,
features_adapter=features_adapter,
**kwargs
)
## reconstruct from latent to pixel space
batch_images = model.decode_first_stage(samples, decode_bs=1, return_cpu=False)
batch_variants.append(batch_images)
## variants, batch, c, t, h, w
batch_variants = torch.stack(batch_variants)
return batch_variants.permute(1, 0, 2, 3, 4, 5), extra_cond
def run_inference(args, gpu_idx):
## model config
config = OmegaConf.load(args.base)
model_config = config.pop("model", OmegaConf.create())
model = instantiate_from_config(model_config)
model = model.cuda(gpu_idx)
assert os.path.exists(args.ckpt_path), "Error: checkpoint Not Found!"
model = load_model_checkpoint(model, args.ckpt_path, args.adapter_ckpt)
model.eval()
## run over data
assert (args.height % 16 == 0) and (args.width % 16 == 0), "Error: image size [h,w] should be multiples of 16!"
## latent noise shape
h, w = args.height // 8, args.width // 8
channels = model.channels
frames = model.temporal_length
noise_shape = [args.bs, channels, frames, h, w]
## inference
start = time.time()
prompt = args.prompt
video = load_video(args.video, args.frame_stride, video_size=(args.height, args.width), video_frames=16)
video = video.unsqueeze(0).to("cuda")
with torch.no_grad():
batch_samples, batch_conds = adapter_guided_synthesis(model, prompt, video, noise_shape, args.n_samples, args.ddim_steps, args.ddim_eta, \
args.unconditional_guidance_scale, args.unconditional_guidance_scale_temporal)
batch_samples = batch_samples[0]
os.makedirs(args.savedir, exist_ok=True)
filename = f"{args.prompt}_seed{args.seed}"
filename = filename.replace("/", "_slash_") if "/" in filename else filename
filename = filename.replace(" ", "_") if " " in filename else filename
tensor_to_mp4(video=batch_conds.detach().cpu(), savepath=os.path.join(args.savedir, f'{filename}_depth.mp4'), fps=10)
tensor_to_mp4(video=batch_samples.detach().cpu(), savepath=os.path.join(args.savedir, f'{filename}_sample.mp4'), fps=10)
print(f"Saved in {args.savedir}. Time used: {(time.time() - start):.2f} seconds")
def get_parser():
parser = argparse.ArgumentParser()
parser.add_argument("--savedir", type=str, default=None, help="results saving path")
parser.add_argument("--ckpt_path", type=str, default=None, help="checkpoint path")
parser.add_argument("--adapter_ckpt", type=str, default=None, help="adapter checkpoint path")
parser.add_argument("--base", type=str, help="config (yaml) path")
parser.add_argument("--prompt", type=str, default=None, help="prompt string")
parser.add_argument("--video", type=str, default=None, help="video path")
parser.add_argument("--n_samples", type=int, default=1, help="num of samples per prompt",)
parser.add_argument("--ddim_steps", type=int, default=50, help="steps of ddim if positive, otherwise use DDPM",)
parser.add_argument("--ddim_eta", type=float, default=1.0, help="eta for ddim sampling (0.0 yields deterministic sampling)",)
parser.add_argument("--bs", type=int, default=1, help="batch size for inference")
parser.add_argument("--height", type=int, default=512, help="image height, in pixel space")
parser.add_argument("--width", type=int, default=512, help="image width, in pixel space")
parser.add_argument("--frame_stride", type=int, default=-1, help="frame extracting from input video")
parser.add_argument("--unconditional_guidance_scale", type=float, default=1.0, help="prompt classifier-free guidance")
parser.add_argument("--unconditional_guidance_scale_temporal", type=float, default=None, help="temporal consistency guidance")
parser.add_argument("--seed", type=int, default=2023, help="seed for seed_everything")
return parser
if __name__ == '__main__':
now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
print("@CoVideoGen cond-Inference: %s"%now)
parser = get_parser()
args = parser.parse_args()
seed_everything(args.seed)
rank = 0
run_inference(args, rank)