# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""
Sample new images from a pre-trained DiT.
"""
import os
import sys
import math
import docx
try:
import utils
from diffusion import create_diffusion
from download import find_model
except:
# sys.path.append(os.getcwd())
sys.path.append(os.path.split(sys.path[0])[0])
# sys.path[0]
# os.path.split(sys.path[0])
import utils
from diffusion import create_diffusion
from download import find_model
import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
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 utils import mask_generation_before
from natsort import natsorted
from diffusers.utils.import_utils import is_xformers_available
# def get_input(args):
def get_input(path,args):
input_path = path
# input_path = args.input_path
transform_video = transforms.Compose([
video_transforms.ToTensorVideo(), # TCHW
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 = []
if args.mask_type.startswith('onelast'):
num = int(args.mask_type.split('onelast')[-1])
# get first and last frame
first_frame_path = os.path.join(input_path, natsorted(file_list)[0])
last_frame_path = os.path.join(input_path, natsorted(file_list)[-1])
first_frame = torch.as_tensor(np.array(Image.open(first_frame_path), dtype=np.uint8, copy=True)).unsqueeze(0)
last_frame = torch.as_tensor(np.array(Image.open(last_frame_path), dtype=np.uint8, copy=True)).unsqueeze(0)
for i in range(num):
video_frames.append(first_frame)
# add zeros to frames
num_zeros = args.num_frames-2*num
for i in range(num_zeros):
zeros = torch.zeros_like(first_frame)
video_frames.append(zeros)
for i in range(num):
video_frames.append(last_frame)
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)
else:
for file in file_list:
if file.endswith('jpg') or file.endswith('png'):
image = torch.as_tensor(np.array(Image.open(os.path.join(input_path,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 == '.jpg' or extention == '.png':
# raise TypeError('a single image is not supported yet!!')
print("reading video from a image")
video_frames = []
num = int(args.mask_type.split('first')[-1])
first_frame = torch.as_tensor(np.array(Image.open(input_path), dtype=np.uint8, copy=True)).unsqueeze(0)
for i in range(num):
video_frames.append(first_frame)
num_zeros = args.num_frames-num
for i in range(num_zeros):
zeros = torch.zeros_like(first_frame)
video_frames.append(zeros)
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
else:
raise TypeError(f'{extention} is not supported !!')
else:
raise ValueError('Please check your path input!!')
else:
# raise ValueError('Need to give a video or some images')
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
# prepare inputs
if args.use_fp16:
z = torch.randn(1, 4, args.num_frames, args.latent_h, args.latent_w, dtype=torch.float16, device=device) # b,c,f,h,w
masked_video = masked_video.to(dtype=torch.float16)
mask = mask.to(dtype=torch.float16)
else:
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)
# classifier_free_guidance
if args.do_classifier_free_guidance:
masked_video = torch.cat([masked_video] * 2)
mask = torch.cat([mask] * 2)
z = torch.cat([z] * 2)
prompt_all = [prompt] + [args.negative_prompt]
else:
masked_video = masked_video
mask = mask
z = z
prompt_all = [prompt]
text_prompt = text_encoder(text_prompts=prompt_all, train=False)
model_kwargs = dict(encoder_hidden_states=text_prompt,
class_labels=None,
cfg_scale=args.cfg_scale,
use_fp16=args.use_fp16,) # tav unet
# Sample images:
if args.sample_method == 'ddim':
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_mask
)
elif args.sample_method == 'ddpm':
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_mask
)
samples, _ = samples.chunk(2, dim=0) # [1, 4, 16, 32, 32]
if args.use_fp16:
samples = samples.to(dtype=torch.float16)
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):
# Setup PyTorch:
if args.seed:
torch.manual_seed(args.seed)
torch.set_grad_enabled(False)
device = "cuda" if torch.cuda.is_available() else "cpu"
# device = "cpu"
if args.ckpt is None:
raise ValueError("Please specify a checkpoint path using --ckpt <path>")
# 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('loading model')
model = get_models(args).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()
else:
raise ValueError("xformers is not available. Make sure it is installed correctly")
# load model
ckpt_path = args.ckpt
state_dict = torch.load(ckpt_path, map_location=lambda storage, loc: storage)['ema']
model.load_state_dict(state_dict)
print('loading succeed')
model.eval() # important!
pretrained_model_path = args.pretrained_model_path
diffusion = create_diffusion(str(args.num_sampling_steps))
vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").to(device)
text_encoder = TextEmbedder(pretrained_model_path).to(device)
if args.use_fp16:
print('Warnning: using half percision for inferencing!')
vae.to(dtype=torch.float16)
model.to(dtype=torch.float16)
text_encoder.to(dtype=torch.float16)
# Labels to condition the model with (feel free to change):
prompt = args.text_prompt
if prompt ==[]:
prompt = args.input_path.split('/')[-1].split('.')[0].replace('_', ' ')
else:
prompt = prompt[0]
prompt_base = prompt.replace(' ','_')
prompt = prompt + args.additional_prompt
if not os.path.exists(os.path.join(args.save_img_path)):
os.makedirs(os.path.join(args.save_img_path))
for file in os.listdir(args.img_path):
video_input, reserve_frames = get_input(os.path.join(args.img_path,file),args)
video_input = video_input.to(device).unsqueeze(0)
mask = mask_generation_before(args.mask_type, video_input.shape, video_input.dtype, device)
masked_video = video_input * (mask == 0)
prompt = "tilt up, high quality, stable "
prompt = prompt + args.additional_prompt
video_clip = auto_inpainting(args, video_input, masked_video, mask, prompt, vae, text_encoder, diffusion, model, device,)
video_ = ((video_clip * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1)
torchvision.io.write_video(os.path.join(args.save_img_path, prompt[0:20]+file+ '.mp4'), video_, fps=8)
# video_input, researve_frames = get_input(args) # f,c,h,w
# video_input = video_input.to(device).unsqueeze(0) # b,f,c,h,w
# mask = mask_generation_before(args.mask_type, video_input.shape, video_input.dtype, device) # b,f,c,h,w
# # TODO: change the first3 to last3
# masked_video = video_input * (mask == 0)
# video_clip = auto_inpainting(args, video_input, masked_video, mask, prompt, vae, text_encoder, diffusion, model, device,)
# video_ = ((video_clip * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1)
# torchvision.io.write_video(os.path.join(args.save_img_path, prompt_base+ '.mp4'), video_, fps=8)
print(f'save in {args.save_img_path}')
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="./configs/sample_mask.yaml")
parser.add_argument("--run-time", type=int, default=0)
args = parser.parse_args()
omega_conf = OmegaConf.load(args.config)
omega_conf.run_time = args.run_time
main(omega_conf)