import argparse, os, sys, glob, yaml, math, random
import datetime, time
import numpy as np
from omegaconf import OmegaConf
from tqdm import trange, tqdm
from einops import repeat
from collections import OrderedDict
from decord import VideoReader, cpu
import torch
import torchvision
sys.path.insert(1, os.path.join(sys.path[0], '..', '..'))
from lvdm.models.samplers.ddim import DDIMSampler
def batch_ddim_sampling(model, cond, noise_shape, n_samples=1, ddim_steps=50, ddim_eta=1.0,\
cfg_scale=1.0, temporal_cfg_scale=None, **kwargs):
ddim_sampler = DDIMSampler(model)
uncond_type = model.uncond_type
batch_size = noise_shape[0]
## construct unconditional guidance
if cfg_scale != 1.0:
if isinstance(cond, dict):
c_cat, text_emb = cond["c_concat"][0], cond["c_crossattn"][0]
else:
text_emb = cond
if uncond_type == "empty_seq":
prompts = batch_size * [""]
uc = model.get_learned_conditioning(prompts)
elif uncond_type == "zero_embed":
uc = torch.zeros_like(text_emb)
else:
raise NotImplementedError
## hybrid case
if isinstance(cond, dict):
uc_hybrid = {"c_concat": [c_cat], "c_crossattn": [uc]}
if 'c_adm' in cond:
uc_hybrid.update({'c_adm': cond['c_adm']})
uc = uc_hybrid
else:
uc = None
## sampling
batch_variants = []
for _ in range(n_samples):
if ddim_sampler is not None:
kwargs.update({"clean_cond": True})
samples, _ = ddim_sampler.sample(S=ddim_steps,
conditioning=cond,
batch_size=noise_shape[0],
shape=noise_shape[1:],
verbose=False,
unconditional_guidance_scale=cfg_scale,
unconditional_conditioning=uc,
eta=ddim_eta,
temporal_length=noise_shape[2],
conditional_guidance_scale_temporal=temporal_cfg_scale,
x_T=None,
**kwargs
)
## reconstruct from latent to pixel space
batch_images = model.decode_first_stage(samples)
batch_variants.append(batch_images)
## batch, <samples>, c, t, h, w
batch_variants = torch.stack(batch_variants, dim=1)
return batch_variants
def batch_sliding_interpolation(model, cond, base_videos, base_stride, noise_shape, n_samples=1,\
ddim_steps=50, ddim_eta=1.0, cfg_scale=1.0, temporal_cfg_scale=None, **kwargs):
'''
Current implementation has a flaw: the inter-episode keyframe is used as pre-last and cur-first, so keyframe repeated.
For example, cond_frames=[0,4,7], model.temporal_length=8, base_stride=4, then
base frame : 0 4 8 12 16 20 24 28
interplation: (0~7) (8~15) (16~23) (20~27)
'''
b,c,t,h,w = noise_shape
base_z0 = model.encode_first_stage(base_videos)
unit_length = model.temporal_length
n_base_frames = base_videos.shape[2]
n_refs = len(model.cond_frames)
sliding_steps = (n_base_frames-1) // (n_refs-1)
sliding_steps = sliding_steps+1 if (n_base_frames-1) % (n_refs-1) > 0 else sliding_steps
cond_mask = model.cond_mask.to("cuda")
proxy_z0 = torch.zeros((b,c,unit_length,h,w), dtype=torch.float32).to("cuda")
batch_samples = None
last_offset = None
for idx in range(sliding_steps):
base_idx = idx * (n_refs-1)
## check index overflow
if base_idx+n_refs > n_base_frames:
last_offset = base_idx - (n_base_frames - n_refs)
base_idx = n_base_frames - n_refs
cond_z0 = base_z0[:,:,base_idx:base_idx+n_refs,:,:]
proxy_z0[:,:,model.cond_frames,:,:] = cond_z0
if isinstance(cond, dict):
c_cat, text_emb = cond["c_concat"][0], cond["c_crossattn"][0]
episode_idx = idx * unit_length
if last_offset is not None:
episode_idx = episode_idx - last_offset * base_stride
cond_idx = {"c_concat": [c_cat[:,:,episode_idx:episode_idx+unit_length,:,:]], "c_crossattn": [text_emb]}
else:
cond_idx = cond
noise_shape_idx = [b,c,unit_length,h,w]
## batch, <samples>, c, t, h, w
batch_idx = batch_ddim_sampling(model, cond_idx, noise_shape_idx, n_samples, ddim_steps, ddim_eta, cfg_scale, \
temporal_cfg_scale, mask=cond_mask, x0=proxy_z0, **kwargs)
if batch_samples is None:
batch_samples = batch_idx
else:
## b,s,c,t,h,w
if last_offset is None:
batch_samples = torch.cat([batch_samples[:,:,:,:-1,:,:], batch_idx], dim=3)
else:
batch_samples = torch.cat([batch_samples[:,:,:,:-1,:,:], batch_idx[:,:,:,last_offset * base_stride:,:,:]], dim=3)
return batch_samples
def get_filelist(data_dir, ext='*'):
file_list = glob.glob(os.path.join(data_dir, '*.%s'%ext))
file_list.sort()
return file_list
def get_dirlist(path):
list = []
if (os.path.exists(path)):
files = os.listdir(path)
for file in files:
m = os.path.join(path,file)
if (os.path.isdir(m)):
list.append(m)
list.sort()
return list
def load_model_checkpoint(model, ckpt, adapter_ckpt=None):
def load_checkpoint(model, ckpt, full_strict):
state_dict = torch.load(ckpt, map_location="cpu")
try:
## deepspeed
new_pl_sd = OrderedDict()
for key in state_dict['module'].keys():
new_pl_sd[key[16:]]=state_dict['module'][key]
model.load_state_dict(new_pl_sd, strict=full_strict)
except:
if "state_dict" in list(state_dict.keys()):
state_dict = state_dict["state_dict"]
model.load_state_dict(state_dict, strict=full_strict)
return model
if adapter_ckpt:
## main model
load_checkpoint(model, ckpt, full_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:
load_checkpoint(model, ckpt, full_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_batch(filepath_list, frame_stride, video_size=(256,256), video_frames=16):
'''
Notice about some special cases:
1. video_frames=-1 means to take all the frames (with fs=1)
2. when the total video frames is less than required, padding strategy will be used (repreated last frame)
'''
fps_list = []
batch_tensor = []
assert frame_stride > 0, "valid frame stride should be a positive interge!"
for filepath in filepath_list:
padding_num = 0
vidreader = VideoReader(filepath, ctx=cpu(0), width=video_size[1], height=video_size[0])
fps = vidreader.get_avg_fps()
total_frames = len(vidreader)
max_valid_frames = (total_frames-1) // frame_stride + 1
if video_frames < 0:
## all frames are collected: fs=1 is a must
required_frames = total_frames
frame_stride = 1
else:
required_frames = video_frames
query_frames = min(required_frames, max_valid_frames)
frame_indices = [frame_stride*i for i in range(query_frames)]
## [t,h,w,c] -> [c,t,h,w]
frames = vidreader.get_batch(frame_indices)
frame_tensor = torch.tensor(frames.asnumpy()).permute(3, 0, 1, 2).float()
frame_tensor = (frame_tensor / 255. - 0.5) * 2
if max_valid_frames < required_frames:
padding_num = required_frames - max_valid_frames
frame_tensor = torch.cat([frame_tensor, *([frame_tensor[:,-1:,:,:]]*padding_num)], dim=1)
print(f'{os.path.split(filepath)[1]} is not long enough: {padding_num} frames padded.')
batch_tensor.append(frame_tensor)
sample_fps = int(fps/frame_stride)
fps_list.append(sample_fps)
return torch.stack(batch_tensor, dim=0)
def save_videos(batch_tensors, savedir, filenames, fps=10):
# b,samples,c,t,h,w
n_samples = batch_tensors.shape[1]
for idx, vid_tensor in enumerate(batch_tensors):
video = vid_tensor.detach().cpu()
video = torch.clamp(video.float(), -1., 1.)
video = video.permute(2, 0, 1, 3, 4) # t,n,c,h,w
frame_grids = [torchvision.utils.make_grid(framesheet, nrow=int(n_samples)) 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)
savepath = os.path.join(savedir, f"{filenames[idx]}.mp4")
torchvision.io.write_video(savepath, grid, fps=fps, video_codec='h264', options={'crf': '10'})