import argparse import datetime import random import os import logging from omegaconf import OmegaConf import torch import diffusers from diffusers import AutoencoderKL, DDIMScheduler from transformers import CLIPTextModel, CLIPTokenizer from consisti2v.models.videoldm_unet import VideoLDMUNet3DConditionModel from consisti2v.pipelines.pipeline_autoregress_animation import AutoregressiveAnimationPipeline from consisti2v.utils.util import save_videos_grid from diffusers.utils.import_utils import is_xformers_available def main(args, config): logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) diffusers.utils.logging.set_verbosity_info() time_str = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") savedir = f"{config.output_dir}/{config.output_name}-{time_str}" os.makedirs(savedir) samples = [] sample_idx = 0 ### >>> create validation pipeline >>> ### if config.pipeline_pretrained_path is None: noise_scheduler = DDIMScheduler(**OmegaConf.to_container(config.noise_scheduler_kwargs)) tokenizer = CLIPTokenizer.from_pretrained(config.pretrained_model_path, subfolder="tokenizer", use_safetensors=True) text_encoder = CLIPTextModel.from_pretrained(config.pretrained_model_path, subfolder="text_encoder") vae = AutoencoderKL.from_pretrained(config.pretrained_model_path, subfolder="vae", use_safetensors=True) unet = VideoLDMUNet3DConditionModel.from_pretrained( config.pretrained_model_path, subfolder="unet", variant=config.unet_additional_kwargs['variant'], temp_pos_embedding=config.unet_additional_kwargs['temp_pos_embedding'], augment_temporal_attention=config.unet_additional_kwargs['augment_temporal_attention'], use_temporal=True, n_frames=config.sampling_kwargs['n_frames'], n_temp_heads=config.unet_additional_kwargs['n_temp_heads'], first_frame_condition_mode=config.unet_additional_kwargs['first_frame_condition_mode'], use_frame_stride_condition=config.unet_additional_kwargs['use_frame_stride_condition'], use_safetensors=True ) params_unet = [p.numel() for n, p in unet.named_parameters()] params_vae = [p.numel() for n, p in vae.named_parameters()] params_text_encoder = [p.numel() for n, p in text_encoder.named_parameters()] params = params_unet + params_vae + params_text_encoder print(f"### UNet Parameters: {sum(params) / 1e6} M") # 1. unet ckpt if config.unet_path is not None: if os.path.isdir(config.unet_path): unet_dict = VideoLDMUNet3DConditionModel.from_pretrained(config.unet_path) m, u = unet.load_state_dict(unet_dict.state_dict(), strict=False) assert len(u) == 0 del unet_dict else: checkpoint_dict = torch.load(config.unet_path, map_location="cpu") state_dict = checkpoint_dict["state_dict"] if "state_dict" in checkpoint_dict else checkpoint_dict if config.unet_ckpt_prefix is not None: state_dict = {k.replace(config.unet_ckpt_prefix, ''): v for k, v in state_dict.items()} m, u = unet.load_state_dict(state_dict, strict=False) assert len(u) == 0 if is_xformers_available() and int(torch.__version__.split(".")[0]) < 2: unet.enable_xformers_memory_efficient_attention() pipeline = AutoregressiveAnimationPipeline( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=noise_scheduler) else: pipeline = AutoregressiveAnimationPipeline.from_pretrained(config.pipeline_pretrained_path) pipeline.to("cuda") # (frameinit) initialize frequency filter for noise reinitialization ------------- if config.frameinit_kwargs.enable: pipeline.init_filter( width = config.sampling_kwargs.width, height = config.sampling_kwargs.height, video_length = config.sampling_kwargs.n_frames, filter_params = config.frameinit_kwargs.filter_params, ) # ------------------------------------------------------------------------------- ### <<< create validation pipeline <<< ### if args.prompt is not None: prompts = [args.prompt] n_prompts = [args.n_prompt] first_frame_paths = [args.path_to_first_frame] random_seeds = [int(args.seed)] if args.seed != "random" else "random" else: prompt_config = OmegaConf.load(args.prompt_config) prompts = prompt_config.prompts n_prompts = list(prompt_config.n_prompts) * len(prompts) if len(prompt_config.n_prompts) == 1 else prompt_config.n_prompts first_frame_paths = prompt_config.path_to_first_frames random_seeds = prompt_config.seeds if random_seeds == "random": random_seeds = [random.randint(0, 1e5) for _ in range(len(prompts))] else: random_seeds = [random_seeds] if isinstance(random_seeds, int) else list(random_seeds) random_seeds = random_seeds * len(prompts) if len(random_seeds) == 1 else random_seeds config.prompt_kwargs = OmegaConf.create({"random_seeds": [], "prompts": prompts, "n_prompts": n_prompts, "first_frame_paths": first_frame_paths}) for prompt_idx, (prompt, n_prompt, first_frame_path, random_seed) in enumerate(zip(prompts, n_prompts, first_frame_paths, random_seeds)): # manually set random seed for reproduction if random_seed != -1: torch.manual_seed(random_seed) else: torch.seed() config.prompt_kwargs.random_seeds.append(torch.initial_seed()) print(f"current seed: {torch.initial_seed()}") print(f"sampling {prompt} ...") sample = pipeline( prompt, negative_prompt = n_prompt, first_frame_paths = first_frame_path, num_inference_steps = config.sampling_kwargs.steps, guidance_scale_txt = config.sampling_kwargs.guidance_scale_txt, guidance_scale_img = config.sampling_kwargs.guidance_scale_img, width = config.sampling_kwargs.width, height = config.sampling_kwargs.height, video_length = config.sampling_kwargs.n_frames, noise_sampling_method = config.unet_additional_kwargs['noise_sampling_method'], noise_alpha = float(config.unet_additional_kwargs['noise_alpha']), eta = config.sampling_kwargs.ddim_eta, frame_stride = config.sampling_kwargs.frame_stride, guidance_rescale = config.sampling_kwargs.guidance_rescale, num_videos_per_prompt = config.sampling_kwargs.num_videos_per_prompt, autoregress_steps = config.sampling_kwargs.autoregress_steps, use_frameinit = config.frameinit_kwargs.enable, frameinit_noise_level = config.frameinit_kwargs.noise_level, ).videos samples.append(sample) prompt = "-".join((prompt.replace("/", "").split(" ")[:10])).replace(":", "") if sample.shape[0] > 1: for cnt, samp in enumerate(sample): save_videos_grid(samp.unsqueeze(0), f"{savedir}/sample/{sample_idx}-{cnt + 1}-{prompt}.{args.format}", format=args.format) else: save_videos_grid(sample, f"{savedir}/sample/{sample_idx}-{prompt}.{args.format}", format=args.format) print(f"save to {savedir}/sample/{prompt}.{args.format}") sample_idx += 1 samples = torch.concat(samples) save_videos_grid(samples, f"{savedir}/sample.{args.format}", n_rows=4, format=args.format) OmegaConf.save(config, f"{savedir}/config.yaml") if args.save_model: pipeline.save_pretrained(f"{savedir}/model") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--inference_config", type=str, default="configs/inference/inference_autoregress.yaml") parser.add_argument("--prompt", "-p", type=str, default=None) parser.add_argument("--n_prompt", "-n", type=str, default="") parser.add_argument("--seed", type=str, default="random") parser.add_argument("--path_to_first_frame", "-f", type=str, default=None) parser.add_argument("--prompt_config", type=str, default="configs/prompts/default.yaml") parser.add_argument("--format", type=str, default="gif", choices=["gif", "mp4"]) parser.add_argument("--save_model", action="store_true") parser.add_argument("optional_args", nargs='*', default=[]) args = parser.parse_args() config = OmegaConf.load(args.inference_config) if args.optional_args: modified_config = OmegaConf.from_dotlist(args.optional_args) config = OmegaConf.merge(config, modified_config) main(args, config)