Spaces:
Sleeping
Sleeping
| 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'}) | |