""" we introduce a temporal interpolation network to enhance the smoothness of generated videos and synthesize richer temporal details. This network takes a 16-frame base video as input and produces an upsampled output consisting of 61 frames. """ import os import sys import math try: import utils from diffusion import create_diffusion from download import find_model except: sys.path.append(os.path.split(sys.path[0])[0]) import utils from diffusion import create_diffusion from download import find_model import torch import argparse import torchvision from einops import rearrange from models import get_models from torchvision.utils import save_image from diffusers.models import AutoencoderKL from models.clip import TextEmbedder from omegaconf import OmegaConf from PIL import Image import numpy as np from torchvision import transforms sys.path.append("..") from datasets import video_transforms from decord import VideoReader from utils import mask_generation, mask_generation_before from natsort import natsorted from diffusers.utils.import_utils import is_xformers_available torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True def get_input(args): input_path = args.input_path transform_video = transforms.Compose([ video_transforms.ToTensorVideo(), # TCHW # video_transforms.CenterCropResizeVideo((args.image_h, args.image_w)), video_transforms.ResizeVideo((args.image_h, args.image_w)), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True) ]) temporal_sample_func = video_transforms.TemporalRandomCrop(args.num_frames * args.frame_interval) if input_path is not None: print(f'loading video from {input_path}') if os.path.isdir(input_path): file_list = os.listdir(input_path) video_frames = [] for file in file_list: if file.endswith('jpg') or file.endswith('png'): image = torch.as_tensor(np.array(Image.open(file), dtype=np.uint8, copy=True)).unsqueeze(0) video_frames.append(image) else: continue n = 0 video_frames = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2) # f,c,h,w video_frames = transform_video(video_frames) return video_frames, n elif os.path.isfile(input_path): _, full_file_name = os.path.split(input_path) file_name, extention = os.path.splitext(full_file_name) if extention == '.mp4': video_reader = VideoReader(input_path) total_frames = len(video_reader) start_frame_ind, end_frame_ind = temporal_sample_func(total_frames) frame_indice = np.linspace(start_frame_ind, end_frame_ind-1, args.num_frames, dtype=int) video_frames = torch.from_numpy(video_reader.get_batch(frame_indice).asnumpy()).permute(0, 3, 1, 2).contiguous() video_frames = transform_video(video_frames) n = 3 del video_reader return video_frames, n else: raise TypeError(f'{extention} is not supported !!') else: raise ValueError('Please check your path input!!') else: print('given video is None, using text to video') video_frames = torch.zeros(16,3,args.latent_h,args.latent_w,dtype=torch.uint8) args.mask_type = 'all' video_frames = transform_video(video_frames) n = 0 return video_frames, n def auto_inpainting(args, video_input, masked_video, mask, prompt, vae, text_encoder, diffusion, model, device,): b,f,c,h,w=video_input.shape latent_h = args.image_size[0] // 8 latent_w = args.image_size[1] // 8 z = torch.randn(1, 4, args.num_frames, args.latent_h, args.latent_w, device=device) # b,c,f,h,w masked_video = rearrange(masked_video, 'b f c h w -> (b f) c h w').contiguous() masked_video = vae.encode(masked_video).latent_dist.sample().mul_(0.18215) masked_video = rearrange(masked_video, '(b f) c h w -> b c f h w', b=b).contiguous() mask = torch.nn.functional.interpolate(mask[:,:,0,:], size=(latent_h, latent_w)).unsqueeze(1) masked_video = torch.cat([masked_video] * 2) if args.do_classifier_free_guidance else masked_video mask = torch.cat([mask] * 2) if args.do_classifier_free_guidance else mask z = torch.cat([z] * 2) if args.do_classifier_free_guidance else z prompt_all = [prompt] + [args.negative_prompt] if args.do_classifier_free_guidance else [prompt] text_prompt = text_encoder(text_prompts=prompt_all, train=False) model_kwargs = dict(encoder_hidden_states=text_prompt, class_labels=None) if args.use_ddim_sample_loop: samples = diffusion.ddim_sample_loop( model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, \ progress=True, device=device, mask=mask, x_start=masked_video, use_concat=args.use_concat ) else: samples = diffusion.p_sample_loop( model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, \ progress=True, device=device, mask=mask, x_start=masked_video, use_concat=args.use_concat ) # torch.Size([2, 4, 16, 32, 32]) samples, _ = samples.chunk(2, dim=0) # [1, 4, 16, 32, 32] video_clip = samples[0].permute(1, 0, 2, 3).contiguous() # [16, 4, 32, 32] video_clip = vae.decode(video_clip / 0.18215).sample # [16, 3, 256, 256] return video_clip def auto_inpainting_copy_no_mask(args, video_input, prompt, vae, text_encoder, diffusion, model, device,): b,f,c,h,w=video_input.shape latent_h = args.image_size[0] // 8 latent_w = args.image_size[1] // 8 video_input = rearrange(video_input, 'b f c h w -> (b f) c h w').contiguous() video_input = vae.encode(video_input).latent_dist.sample().mul_(0.18215) video_input = rearrange(video_input, '(b f) c h w -> b c f h w', b=b).contiguous() lr_indice = torch.IntTensor([i for i in range(0,62,4)]).to(device) copied_video = torch.index_select(video_input, 2, lr_indice) copied_video = torch.repeat_interleave(copied_video, 4, dim=2) copied_video = copied_video[:,:,1:-2,:,:] copied_video = torch.cat([copied_video] * 2) if args.do_classifier_free_guidance else copied_video torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) z = torch.randn(1, 4, args.num_frames, args.latent_h, args.latent_w, device=device) # b,c,f,h,w z = torch.cat([z] * 2) if args.do_classifier_free_guidance else z prompt_all = [prompt] + [args.negative_prompt] if args.do_classifier_free_guidance else [prompt] text_prompt = text_encoder(text_prompts=prompt_all, train=False) model_kwargs = dict(encoder_hidden_states=text_prompt, class_labels=None) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) if args.use_ddim_sample_loop: samples = diffusion.ddim_sample_loop( model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, \ progress=True, device=device, mask=None, x_start=copied_video, use_concat=args.use_concat, copy_no_mask=args.copy_no_mask, ) else: raise ValueError(f'We only have ddim sampling implementation for now') samples, _ = samples.chunk(2, dim=0) # [1, 4, 16, 32, 32] video_clip = samples[0].permute(1, 0, 2, 3).contiguous() # [16, 4, 32, 32] video_clip = vae.decode(video_clip / 0.18215).sample # [16, 3, 256, 256] return video_clip def main(args): for seed in args.seed_list: args.seed = seed torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) # print(f'torch.seed() = {torch.seed()}') print('sampling begins') torch.set_grad_enabled(False) device = "cuda" if torch.cuda.is_available() else "cpu" # device = "cpu" ckpt_path = args.pretrained_path + "/lavie_interpolation.pt" sd_path = args.pretrained_path + "/stable-diffusion-v1-4" for ckpt in [ckpt_path]: ckpt_num = str(ckpt_path).zfill(7) # Load model: latent_h = args.image_size[0] // 8 latent_w = args.image_size[1] // 8 args.image_h = args.image_size[0] args.image_w = args.image_size[1] args.latent_h = latent_h args.latent_w = latent_w print(f'args.copy_no_mask = {args.copy_no_mask}') model = get_models(args, sd_path).to(device) if args.use_compile: model = torch.compile(model) if args.enable_xformers_memory_efficient_attention: if is_xformers_available(): model.enable_xformers_memory_efficient_attention() # model.enable_vae_slicing() # ziqi added else: raise ValueError("xformers is not available. Make sure it is installed correctly") # Auto-download a pre-trained model or load a custom checkpoint from train.py: print(f'loading model from {ckpt_path}') # load ckpt state_dict = find_model(ckpt_path) print(f'state_dict["conv_in.weight"].shape = {state_dict["conv_in.weight"].shape}') # [320, 8, 3, 3] print('loading succeed') # model.load_state_dict(state_dict) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) model.eval() # important! diffusion = create_diffusion(str(args.num_sampling_steps)) vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae").to(device) text_encoder = TextEmbedder(sd_path).to(device) video_list = os.listdir(args.input_folder) args.input_path_list = [os.path.join(args.input_folder, video) for video in video_list] for input_path in args.input_path_list: args.input_path = input_path print(f'=======================================') if not args.input_path.endswith('.mp4'): print(f'Skipping {args.input_path}') continue print(f'args.input_path = {args.input_path}') torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) # Labels to condition the model with (feel free to change): video_name = args.input_path.split('/')[-1].split('.mp4')[0] args.prompt = [video_name] print(f'args.prompt = {args.prompt}') prompts = args.prompt class_name = [p + args.additional_prompt for p in prompts] if not os.path.exists(os.path.join(args.output_folder)): os.makedirs(os.path.join(args.output_folder)) video_input, researve_frames = get_input(args) # f,c,h,w video_input = video_input.to(device).unsqueeze(0) # b,f,c,h,w if args.copy_no_mask: pass else: mask = mask_generation_before(args.mask_type, video_input.shape, video_input.dtype, device) # b,f,c,h,w if args.copy_no_mask: pass else: if args.mask_type == 'tsr': masked_video = video_input * (mask == 0) else: masked_video = video_input * (mask == 0) all_video = [] if researve_frames != 0: all_video.append(video_input) for idx, prompt in enumerate(class_name): if idx == 0: if args.copy_no_mask: video_clip = auto_inpainting_copy_no_mask(args, video_input, prompt, vae, text_encoder, diffusion, model, device,) video_clip_ = video_clip.unsqueeze(0) all_video.append(video_clip_) else: video_clip = auto_inpainting(args, video_input, masked_video, mask, prompt, vae, text_encoder, diffusion, model, device,) video_clip_ = video_clip.unsqueeze(0) all_video.append(video_clip_) else: raise NotImplementedError masked_video = video_input * (mask == 0) video_clip = auto_inpainting_copy_no_mask(args, video_clip.unsqueeze(0), masked_video, mask, prompt, vae, text_encoder, diffusion, model, device,) video_clip_ = video_clip.unsqueeze(0) all_video.append(video_clip_[:, 3:]) video_ = ((video_clip * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1) for fps in args.fps_list: save_path = args.output_folder if not os.path.exists(os.path.join(save_path)): os.makedirs(os.path.join(save_path)) local_save_path = os.path.join(save_path, f'{video_name}.mp4') print(f'save in {local_save_path}') torchvision.io.write_video(local_save_path, video_, fps=fps) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, required=True) args = parser.parse_args() main(**OmegaConf.load(args.config))