import argparse import logging import os import warnings from pathlib import Path import json import wandb from datetime import datetime import diffusers import numpy as np import torch import torch.utils.checkpoint import transformers from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import DistributedDataParallelKwargs from diffusers import AutoencoderKL, DDIMScheduler from diffusers.utils import check_min_version from diffusers.utils.import_utils import is_xformers_available from omegaconf import OmegaConf from PIL import Image, ImageOps from transformers import CLIPVisionModelWithProjection from decord import VideoReader from pipelines.pipeline_vfi2video_share_vae import Pose2VideoPipeline from models.model_vfi import VFIShareVAEModel from models.pose_guider import PoseGuider from models.conv_in import ConvIn from models.unet_2d_condition import UNet2DConditionModel from models.unet_3d import UNet3DConditionModel from models.mutual_self_attention import ReferenceAttentionControl from utils.util import seed_everything from utils.video_utils import concat_pil warnings.filterwarnings("ignore") check_min_version("0.10.0.dev0") logger = get_logger(__name__, log_level="INFO") def save_videos_from_pil(pil_images, path, fps=8, bitrate=None): import av save_fmt = Path(path).suffix os.makedirs(os.path.dirname(path), exist_ok=True) width, height = pil_images[0].size if save_fmt == ".mp4": codec = "libx264" container = av.open(path, "w") stream = container.add_stream(codec, rate=fps, ) stream.width = width stream.height = height if bitrate is not None: stream.bit_rate = bitrate for pil_image in pil_images: # pil_image = Image.fromarray(image_arr).convert("RGB") av_frame = av.VideoFrame.from_image(pil_image) container.mux(stream.encode(av_frame)) container.mux(stream.encode()) container.close() elif save_fmt == ".gif": pil_images[0].save( fp=path, format="GIF", append_images=pil_images[1:], save_all=True, duration=(1 / fps * 1000), loop=0, ) else: raise ValueError("Unsupported file type. Use .mp4 or .gif.") def pad_to_square(image, size): width, height = image.size ratio = size / max(width, height) width, height = int(width * ratio), int(height * ratio) image = image.resize((width, height)) left, top = (size - width) // 2, (size - height) // 2 return ImageOps.expand(image, (left, top, size - width - left, size - height - top), fill=0) def log_validation( cfg, vae, image_enc, model, scheduler, accelerator, width, height, seed=42, dtype=torch.float32, ): logger.info("Running validation ...") unwrap_model = accelerator.unwrap_model(model).eval() denoising_unet = unwrap_model.denoising_unet reference_unet = unwrap_model.reference_unet pose_guider = unwrap_model.pose_guider seed_guider = unwrap_model.seed_guider generator = torch.manual_seed(seed) vae = vae.to(dtype=dtype) image_enc = image_enc.to(dtype=dtype) pipeline = Pose2VideoPipeline( vae=vae, image_encoder=image_enc, reference_unet=reference_unet, denoising_unet=denoising_unet, pose_guider=pose_guider, seed_guider=seed_guider, scheduler=scheduler, ) pipeline = pipeline.to(accelerator.device) vid_meta = [] for data_meta_path in cfg.data.meta_paths: vid_meta.extend(json.load(open(data_meta_path, "r"))) all_test_cases = [item for item in vid_meta if item.get("mode") == "test"] current_device_id = accelerator.process_index pil_images = [] # if cfg.wandb_project == "debug": # all_test_cases = all_test_cases[:2] + all_test_cases[-2:] # else: # all_test_cases = all_test_cases[:8] + all_test_cases[-8:] with torch.no_grad(): for idx, test_case in enumerate(all_test_cases): # if idx > 1: break if idx % accelerator.num_processes == current_device_id: video_path = test_case["video_path"] kps_pth = test_case["kps_path"] rgb_video = VideoReader(video_path) kps_video = VideoReader(kps_pth) # if idx >= len(all_test_cases) // 2: # # for dance video we pick from 8s # ref_id = 120 # else: ref_id = test_case.get("start_idx", 3 - cfg.data.seed_frames) + cfg.data.seed_frames start_g = ref_id - cfg.data.seed_frames end_s = ref_id + cfg.data.generated_frames end_g = ref_id + cfg.data.generated_frames + cfg.data.seed_frames ref_img = rgb_video[ref_id-1] ref_img_pil = Image.fromarray(ref_img.asnumpy()) w, h = ref_img_pil.size tgt_vidpil_lst = [] tgt_guid_vid_list = [] pose_guid_vid_list = [] for index in range(start_g, ref_id): img = rgb_video[index] pil_img = Image.fromarray(img.asnumpy()) tgt_guid_pil = pil_img.copy() pose_guid_pil = Image.new(pil_img.mode, pil_img.size, 0) pose_guid_vid_list.append(pose_guid_pil) tgt_vidpil_lst.append(pil_img) tgt_guid_vid_list.append(tgt_guid_pil) for index in range(ref_id, end_s): if index >= ref_id + cfg.data.generated_frames // 2 : img_dup = rgb_video[end_s] pil_img_dup = Image.fromarray(img_dup.asnumpy()) tgt_guid_vid_list.append(pil_img_dup) else: img_dup = rgb_video[ref_id-1] pil_img_dup = Image.fromarray(img_dup.asnumpy()) tgt_guid_vid_list.append(pil_img_dup) pose_guid_pil = Image.fromarray(kps_video[index].asnumpy()) pose_guid_vid_list.append(pose_guid_pil) img = rgb_video[index] pil_img = Image.fromarray(img.asnumpy()) tgt_vidpil_lst.append(pil_img) for index in range(end_s, end_g): img = rgb_video[index] pil_img = Image.fromarray(img.asnumpy()) tgt_guid_pil = pil_img.copy() pose_guid_pil = Image.new(pil_img.mode, pil_img.size, 0) pose_guid_vid_list.append(pose_guid_pil) tgt_vidpil_lst.append(pil_img) tgt_guid_vid_list.append(tgt_guid_pil) seed_mask = None if cfg.get("seed_input_channels", 4) == 5: w = h = cfg.data.image_size ref_img_pil = pad_to_square(ref_img_pil, w) tgt_guid_vid_list = [pad_to_square(image, w) for image in tgt_guid_vid_list] pose_guid_vid_list = [pad_to_square(image, w) for image in pose_guid_vid_list] tgt_vidpil_lst = [pad_to_square(image, w) for image in tgt_vidpil_lst] seed_mask = torch.zeros( (1, 1, cfg.data.sample_frames, h // 8, w // 8), device=accelerator.device ) seed_mask[:, :, :cfg.data.seed_frames] = 1 seed_mask[:, :, -cfg.data.seed_frames:] = 1 ori_w = w ori_h = h # generator = torch.manual_seed(seed) if cfg.get("seed_input_channels", 4) == 5: generator = torch.Generator(device=accelerator.device).manual_seed(seed) else: generator = torch.manual_seed(seed) video_tensor = pipeline( ref_img_pil, pose_guid_vid_list, tgt_guid_vid_list, w, h, cfg.data.sample_frames, cfg.validation.denoising_steps, cfg.validation.guidance_scale, generator=generator, seed_mask=seed_mask, ).videos video_tensor = video_tensor[0, ...].permute(1, 2, 3, 0).cpu().numpy() video_pil_lst = [] video_pil_lst_original_size = [] gap_w = (w - ori_w) // 2 gap_h = (h - ori_h) // 2 for frame_idx, image_tensor in enumerate(video_tensor): result_img_pil = Image.fromarray((image_tensor * 255).astype(np.uint8)) result_img_pil_2 = result_img_pil.copy() result_img_np = np.array(result_img_pil_2) result_img_np_ori = result_img_np[gap_h:gap_h+ori_h, gap_w:gap_w+ori_w] result_img_pil_ori = Image.fromarray(result_img_np_ori) # print(result_img_pil_ori.size) if frame_idx >= cfg.data.seed_frames and frame_idx < cfg.data.seed_frames+cfg.data.generated_frames: result_pil_lst = [pose_guid_vid_list[frame_idx], tgt_guid_vid_list[frame_idx], result_img_pil, tgt_vidpil_lst[frame_idx], result_img_pil] else: result_pil_lst = [pose_guid_vid_list[frame_idx], tgt_guid_vid_list[frame_idx], result_img_pil, tgt_vidpil_lst[frame_idx], tgt_vidpil_lst[frame_idx]] concated_pil = concat_pil(result_pil_lst) video_pil_lst_original_size.append(result_img_pil_ori) video_pil_lst.append(concated_pil) pil_images.append({"name": f"{test_case['video_path'][-14:-4]}", "video": video_pil_lst}) # clip = ImageSequenceClip(video_pil_lst_original_size, fps=24) # clip.write_videofile(video_path.replace("_rgb_", "_blend_res_"), codec='libx264', bitrate="4000k") output_path = video_path.replace("_rgb_", "_blend_res_") if cfg.get("seed_input_channels", 4) == 5 or output_path == video_path: output_path = str(Path(cfg.output_dir) / cfg.exp_name / f"{Path(video_path).stem}_blend_res.mp4") save_videos_from_pil(video_pil_lst_original_size, output_path, bitrate=4000000) unwrap_model.train() vae = vae.to(dtype=torch.float16) image_enc = image_enc.to(dtype=torch.float16) del pipeline torch.cuda.empty_cache() return pil_images def main(cfg): os.environ.setdefault("WANDB_MODE", "disabled") kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator( gradient_accumulation_steps=cfg.solver.gradient_accumulation_steps, mixed_precision=cfg.solver.mixed_precision, log_with="wandb", project_dir=cfg.wandb_log_dir, kwargs_handlers=[kwargs], ) logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) logger.info(accelerator.state, main_process_only=True) if accelerator.is_local_main_process: transformers.utils.logging.set_verbosity_warning() diffusers.utils.logging.set_verbosity_info() else: transformers.utils.logging.set_verbosity_error() diffusers.utils.logging.set_verbosity_error() if cfg.seed is not None: seed_everything(cfg.seed) sched_kwargs = OmegaConf.to_container(cfg.noise_scheduler_kwargs) if cfg.enable_zero_snr: sched_kwargs.update( rescale_betas_zero_snr=True, timestep_spacing="trailing", prediction_type="v_prediction", ) val_noise_scheduler = DDIMScheduler(**sched_kwargs) sched_kwargs.update({"beta_schedule": "scaled_linear"}) vae = AutoencoderKL.from_pretrained(cfg.vae_model_path).to("cuda", dtype=torch.float32) reference_unet = UNet2DConditionModel.from_pretrained( cfg.base_model_path, subfolder="unet", ).to(device="cuda") denoising_unet = UNet3DConditionModel.from_pretrained_2d( cfg.base_model_path, cfg.mm_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container( cfg.unet_additional_kwargs ), ).to(device="cuda") image_enc = CLIPVisionModelWithProjection.from_pretrained( cfg.image_encoder_path, ).to(dtype=torch.float32, device="cuda") pose_guider = PoseGuider( conditioning_embedding_channels=320, block_out_channels=(16, 32, 96, 256), conditioning_channels=3, ).to(device="cuda") seed_guider = ConvIn( conditioning_embedding_channels=320, conditioning_channels=cfg.get("seed_input_channels", 4) ) vae.requires_grad_(False) image_enc.requires_grad_(False) denoising_unet.requires_grad_(False) pose_guider.requires_grad_(False) seed_guider.requires_grad_(False) reference_control_writer = ReferenceAttentionControl( reference_unet, do_classifier_free_guidance=False, mode="write", fusion_blocks="full", ) reference_control_reader = ReferenceAttentionControl( denoising_unet, do_classifier_free_guidance=False, mode="read", fusion_blocks="full", ) model = VFIShareVAEModel( reference_unet, denoising_unet, seed_guider, pose_guider, reference_control_writer, reference_control_reader, ) checkpoint_path = cfg.get("checkpoint_path") or os.path.join( cfg.stage1_ckpt_dir, f"checkpoint-{cfg.stage1_ckpt_step}", "pytorch_model.bin" ) old_ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=True) load_result = model.load_state_dict(old_ckpt, strict=True) checkpoint_info = { "path": str(checkpoint_path), "strict": True, "missing_keys": load_result.missing_keys, "unexpected_keys": load_result.unexpected_keys, "seed_guider_weight_shape": list(old_ckpt["seed_guider.conv_in.weight"].shape), } del old_ckpt logger.info(f"Strictly loaded checkpoint {checkpoint_path}") model = accelerator.prepare( model ) if cfg.solver.enable_xformers_memory_efficient_attention: if is_xformers_available(): reference_unet.enable_xformers_memory_efficient_attention() denoising_unet.enable_xformers_memory_efficient_attention() else: raise ValueError( "xformers is not available. Make sure it is installed correctly" ) if accelerator.is_main_process: run_time = datetime.now().strftime("%Y%m%d-%H%M") accelerator.init_trackers( cfg.wandb_project, init_kwargs={"wandb": {"name": cfg.exp_name + "_" + run_time, "dir":cfg.wandb_log_dir}}, ) config_dict = OmegaConf.to_container(cfg) wandb.config.update(config_dict) accelerator.wait_for_everyone() sample_dicts = log_validation( cfg=cfg, vae=vae, image_enc=image_enc, model=model, scheduler=val_noise_scheduler, accelerator=accelerator, width=cfg.data.image_size, height=cfg.data.image_size, seed=cfg.seed, ) global_step = 0 image_save_dir = f"{cfg.wandb_log_dir}/{cfg.exp_name}/{global_step:06d}/" if not os.path.exists(image_save_dir): os.makedirs(image_save_dir, exist_ok=True) for sample_id, sample_dict in enumerate(sample_dicts): sample_name = sample_dict["name"] img = sample_dict["video"] if img[0].size[0] >= img[0].size[1]: img_path = f"{image_save_dir}{global_step:06d}-{sample_name}-higher.mp4" else: img_path = f"{image_save_dir}{global_step:06d}-{sample_name}-wider.mp4" save_videos_from_pil(img, img_path, fps=6) accelerator.wait_for_everyone() if accelerator.is_main_process: images_to_log_higher = [] images_to_log_wider = [] for img_path in os.listdir(image_save_dir): if img_path.endswith("-higher.mp4"): wandb_img = wandb.Video(image_save_dir+img_path, caption=f"{global_step:06d}-{img_path}") images_to_log_higher.append(wandb_img) elif img_path.endswith("-wider.mp4"): wandb_img = wandb.Video(image_save_dir+img_path, caption=f"{global_step:06d}-{img_path}") images_to_log_wider.append(wandb_img) wandb.log({ "images_higher": images_to_log_higher, "images_wider": images_to_log_wider, "iteration": global_step }) accelerator.wait_for_everyone() accelerator.end_training() if accelerator.is_main_process: receipt_dir = Path(cfg.output_dir) / cfg.exp_name receipt_dir.mkdir(parents=True, exist_ok=True) receipt = { "status": "complete", "checkpoint": checkpoint_info, "seed": cfg.seed, "data": OmegaConf.to_container(cfg.data), "validation": OmegaConf.to_container(cfg.validation), "vae_type": "AutoencoderKL", "vae_model_path": cfg.vae_model_path, "input_jsons": list(cfg.data.meta_paths), "samples": len(sample_dicts), } (receipt_dir / "inference_receipt.json").write_text(json.dumps(receipt, indent=2) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, default="diffusion/configs/inference.yaml") parser.add_argument("--json", type=str, default="") parser.add_argument("--checkpoint", type=str) parser.add_argument("--output-dir", type=str) parser.add_argument("--guidance-scale", type=float) args = parser.parse_args() if args.config[-5:] == ".yaml": config = OmegaConf.load(args.config) else: raise ValueError("Do not support this format config file") if args.json: config.data.meta_paths = [args.json] if args.checkpoint: config.checkpoint_path = args.checkpoint if args.output_dir: config.output_dir = args.output_dir config.wandb_log_dir = str(Path(args.output_dir) / "logs") if args.guidance_scale is not None: config.validation.guidance_scale = args.guidance_scale main(config)