Spaces:
Sleeping
Sleeping
| import torch | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| import os | |
| import sys | |
| try: | |
| import utils | |
| from diffusion import create_diffusion | |
| except: | |
| sys.path.append(os.path.split(sys.path[0])[0]) | |
| import utils | |
| from diffusion import create_diffusion | |
| import argparse | |
| import torchvision | |
| from PIL import Image | |
| from einops import rearrange | |
| from models import get_models | |
| from diffusers.models import AutoencoderKL | |
| from models.clip import TextEmbedder | |
| from omegaconf import OmegaConf | |
| from pytorch_lightning import seed_everything | |
| from utils import mask_generation_before | |
| from diffusers.utils.import_utils import is_xformers_available | |
| from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor | |
| from vlogger.videofusion import fusion | |
| from vlogger.videocaption import captioning | |
| from vlogger.videoaudio import make_audio, merge_video_audio, concatenate_videos | |
| from vlogger.STEB.model_transform import ip_scale_set, ip_transform_model, tca_transform_model | |
| from vlogger.planning_utils.gpt4_utils import (readscript, | |
| readtimescript, | |
| readprotagonistscript, | |
| readreferencescript, | |
| readzhscript) | |
| def auto_inpainting(args, | |
| video_input, | |
| masked_video, | |
| mask, | |
| prompt, | |
| image, | |
| vae, | |
| text_encoder, | |
| image_encoder, | |
| diffusion, | |
| model, | |
| device, | |
| ): | |
| image_prompt_embeds = None | |
| if prompt is None: | |
| prompt = "" | |
| if image is not None: | |
| clip_image = CLIPImageProcessor()(images=image, return_tensors="pt").pixel_values | |
| clip_image_embeds = image_encoder(clip_image.to(device)).image_embeds | |
| uncond_clip_image_embeds = torch.zeros_like(clip_image_embeds).to(device) | |
| image_prompt_embeds = torch.cat([clip_image_embeds, uncond_clip_image_embeds], dim=0) | |
| image_prompt_embeds = rearrange(image_prompt_embeds, '(b n) c -> b n c', b=2).contiguous() | |
| model = ip_scale_set(model, args.ref_cfg_scale) | |
| if args.use_fp16: | |
| image_prompt_embeds = image_prompt_embeds.to(dtype=torch.float16) | |
| b, f, c, h, w = video_input.shape | |
| latent_h = video_input.shape[-2] // 8 | |
| latent_w = video_input.shape[-1] // 8 | |
| if args.use_fp16: | |
| z = torch.randn(1, 4, 16, latent_h, 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, 16, latent_h, 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) | |
| mask = torch.cat([mask] * 2) | |
| z = torch.cat([z] * 2) | |
| prompt_all = [prompt] + [args.negative_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, | |
| ip_hidden_states=image_prompt_embeds) | |
| # Sample images: | |
| 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=True, | |
| ) | |
| 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" | |
| seed_everything(args.seed) | |
| model = get_models(args).to(device) | |
| model = tca_transform_model(model).to(device) | |
| model = ip_transform_model(model).to(device) | |
| 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") | |
| if args.use_compile: | |
| model = torch.compile(model) | |
| ckpt_path = args.ckpt | |
| state_dict = torch.load(ckpt_path, map_location=lambda storage, loc: storage)['ema'] | |
| model_dict = model.state_dict() | |
| pretrained_dict = {} | |
| for k, v in state_dict.items(): | |
| if k in model_dict: | |
| pretrained_dict[k] = v | |
| model_dict.update(pretrained_dict) | |
| model.load_state_dict(model_dict) | |
| model.eval() # important! | |
| diffusion = create_diffusion(str(args.num_sampling_steps)) | |
| vae = AutoencoderKL.from_pretrained(args.pretrained_model_path, subfolder="vae").to(device) | |
| text_encoder = text_encoder = TextEmbedder(args.pretrained_model_path).to(device) | |
| image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.image_encoder_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) | |
| print("model ready!\n", flush=True) | |
| # load protagonist script | |
| character_places = readprotagonistscript(args.protagonist_file_path) | |
| print("protagonists ready!", flush=True) | |
| # load script | |
| video_list = readscript(args.script_file_path) | |
| print("video script ready!", flush=True) | |
| # load reference script | |
| reference_lists = readreferencescript(video_list, character_places, args.reference_file_path) | |
| print("reference script ready!", flush=True) | |
| # load zh script | |
| zh_video_list = readzhscript(args.zh_script_file_path) | |
| print("zh script ready!", flush=True) | |
| # load time script | |
| key_list = [] | |
| for key, value in character_places.items(): | |
| key_list.append(key) | |
| time_list = readtimescript(args.time_file_path) | |
| print("time script ready!", flush=True) | |
| # generation begin | |
| sample_list = [] | |
| for i, text_prompt in enumerate(video_list): | |
| sample_list.append([]) | |
| for time in range(time_list[i]): | |
| if time == 0: | |
| print('Generating the ({}) prompt'.format(text_prompt), flush=True) | |
| if reference_lists[i][0] == 0 or reference_lists[i][0] > len(key_list): | |
| pil_image = None | |
| else: | |
| pil_image = Image.open(args.reference_image_path[reference_lists[i][0] - 1]) | |
| pil_image.resize((256, 256)) | |
| video_input = torch.zeros([1, 16, 3, args.image_size[0], args.image_size[1]]).to(device) | |
| mask = mask_generation_before("first0", video_input.shape, video_input.dtype, device) # b,f,c,h,w | |
| masked_video = video_input * (mask == 0) | |
| samples = auto_inpainting(args, | |
| video_input, | |
| masked_video, | |
| mask, | |
| text_prompt, | |
| pil_image, | |
| vae, | |
| text_encoder, | |
| image_encoder, | |
| diffusion, | |
| model, | |
| device, | |
| ) | |
| sample_list[i].append(samples) | |
| else: | |
| if sum(video.shape[0] for video in sample_list[i]) / args.fps >= time_list[i]: | |
| break | |
| print('Generating the ({}) prompt'.format(text_prompt), flush=True) | |
| if reference_lists[i][0] == 0 or reference_lists[i][0] > len(key_list): | |
| pil_image = None | |
| else: | |
| pil_image = Image.open(args.reference_image_path[reference_lists[i][0] - 1]) | |
| pil_image.resize((256, 256)) | |
| pre_video = sample_list[i][-1][-args.researve_frame:] | |
| f, c, h, w = pre_video.shape | |
| lat_video = torch.zeros(args.num_frames - args.researve_frame, c, h, w).to(device) | |
| video_input = torch.concat([pre_video, lat_video], dim=0) | |
| 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) | |
| video_clip = auto_inpainting(args, | |
| video_input, | |
| masked_video, | |
| mask, | |
| text_prompt, | |
| pil_image, | |
| vae, | |
| text_encoder, | |
| image_encoder, | |
| diffusion, | |
| model, | |
| device, | |
| ) | |
| sample_list[i].append(video_clip[args.researve_frame:]) | |
| print(video_clip[args.researve_frame:].shape) | |
| # transition | |
| if args.video_transition and i != 0: | |
| video_1 = sample_list[i - 1][-1][-1:] | |
| video_2 = sample_list[i][0][:1] | |
| f, c, h, w = video_1.shape | |
| video_middle = torch.zeros(args.num_frames - 2, c, h, w).to(device) | |
| video_input = torch.concat([video_1, video_middle, video_2], dim=0) | |
| video_input = video_input.to(device).unsqueeze(0) | |
| mask = mask_generation_before("onelast1", video_input.shape, video_input.dtype, device) | |
| masked_video = masked_video = video_input * (mask == 0) | |
| video_clip = auto_inpainting(args, | |
| video_input, | |
| masked_video, | |
| mask, | |
| "smooth transition, slow motion, slow changing.", | |
| pil_image, | |
| vae, | |
| text_encoder, | |
| image_encoder, | |
| diffusion, | |
| model, | |
| device, | |
| ) | |
| sample_list[i].insert(0, video_clip[1:-1]) | |
| # save videos | |
| samples = torch.concat(sample_list[i], dim=0) | |
| samples = samples[0: time_list[i] * args.fps] | |
| if not os.path.exists(args.save_origin_video_path): | |
| os.makedirs(args.save_origin_video_path) | |
| video_ = ((samples * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous() | |
| torchvision.io.write_video(args.save_origin_video_path + "/" + f"{i}" + '.mp4', video_, fps=args.fps) | |
| # post processing | |
| fusion(args.save_origin_video_path) | |
| captioning(args.script_file_path, args.zh_script_file_path, args.save_origin_video_path, args.save_caption_video_path) | |
| fusion(args.save_caption_video_path) | |
| make_audio(args.script_file_path, args.save_audio_path) | |
| merge_video_audio(args.save_caption_video_path, args.save_audio_path, args.save_audio_caption_video_path) | |
| concatenate_videos(args.save_audio_caption_video_path) | |
| print('final video save path {}'.format(args.save_audio_caption_video_path)) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--config", type=str, default="configs/vlog_read_script_sample.yaml") | |
| args = parser.parse_args() | |
| omega_conf = OmegaConf.load(args.config) | |
| save_path = omega_conf.save_path | |
| save_origin_video_path = os.path.join(save_path, "origin_video") | |
| save_caption_video_path = os.path.join(save_path.rsplit('/', 1)[0], "caption_video") | |
| save_audio_path = os.path.join(save_path.rsplit('/', 1)[0], "audio") | |
| save_audio_caption_video_path = os.path.join(save_path.rsplit('/', 1)[0], "audio_caption_video") | |
| if omega_conf.sample_num is not None: | |
| for i in range(omega_conf.sample_num): | |
| omega_conf.save_origin_video_path = save_origin_video_path + f'-{i}' | |
| omega_conf.save_caption_video_path = save_caption_video_path + f'-{i}' | |
| omega_conf.save_audio_path = save_audio_path + f'-{i}' | |
| omega_conf.save_audio_caption_video_path = save_audio_caption_video_path + f'-{i}' | |
| omega_conf.seed += i | |
| main(omega_conf) | |
| else: | |
| omega_conf.save_origin_video_path = save_origin_video_path | |
| omega_conf.save_caption_video_path = save_caption_video_path | |
| omega_conf.save_audio_path = save_audio_path | |
| omega_conf.save_audio_caption_video_path = save_audio_caption_video_path | |
| main(omega_conf) | |