Haiyang Liu

Add dated historical interpolation weights and matching runtime

272e50d verified 10 days ago

18.7 kB

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)

	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)