yslan/ObjCtrl-2.5D

1import argparse2import json3import os4import torch5import numpy as np6from tqdm import tqdm7from omegaconf import OmegaConf8from PIL import Image9from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection10from diffusers import AutoencoderKLTemporalDecoder, EulerDiscreteScheduler11from diffusers.utils.import_utils import is_xformers_available12from packaging import version as pver13 14from cameractrl.pipelines.pipeline_animation import StableVideoDiffusionPipelinePoseCond15from cameractrl.models.unet import UNetSpatioTemporalConditionModelPoseCond16from cameractrl.models.pose_adaptor import CameraPoseEncoder17from cameractrl.utils.util import save_videos_grid18 19 20class Camera(object):21    def __init__(self, entry):22        fx, fy, cx, cy = entry[1:5]23        self.fx = fx24        self.fy = fy25        self.cx = cx26        self.cy = cy27        w2c_mat = np.array(entry[7:]).reshape(3, 4)28        w2c_mat_4x4 = np.eye(4)29        w2c_mat_4x4[:3, :] = w2c_mat30        self.w2c_mat = w2c_mat_4x431        self.c2w_mat = np.linalg.inv(w2c_mat_4x4)32 33 34def setup_for_distributed(is_master):35    """36    This function disables printing when not in master process37    """38    import builtins as __builtin__39    builtin_print = __builtin__.print40 41    def print(*args, **kwargs):42        force = kwargs.pop('force', False)43        if is_master or force:44            builtin_print(*args, **kwargs)45 46    __builtin__.print = print47 48 49def custom_meshgrid(*args):50    # ref: https://pytorch.org/docs/stable/generated/torch.meshgrid.html?highlight=meshgrid#torch.meshgrid51    if pver.parse(torch.__version__) < pver.parse('1.10'):52        return torch.meshgrid(*args)53    else:54        return torch.meshgrid(*args, indexing='ij')55 56 57def get_relative_pose(cam_params, zero_first_frame_scale):58    abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]59    abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]60    source_cam_c2w = abs_c2ws[0]61    if zero_first_frame_scale:62        cam_to_origin = 063    else:64        cam_to_origin = np.linalg.norm(source_cam_c2w[:3, 3])65    target_cam_c2w = np.array([66        [1, 0, 0, 0],67        [0, 1, 0, -cam_to_origin],68        [0, 0, 1, 0],69        [0, 0, 0, 1]70    ])71    abs2rel = target_cam_c2w @ abs_w2cs[0]72    ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]73    ret_poses = np.array(ret_poses, dtype=np.float32)74    return ret_poses75 76 77def ray_condition(K, c2w, H, W, device):78    # c2w: B, V, 4, 479    # K: B, V, 480 81    B = K.shape[0]82 83    j, i = custom_meshgrid(84        torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype),85        torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype),86    )87    i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]88    j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]89 90    fx, fy, cx, cy = K.chunk(4, dim=-1)  # B,V, 191 92    zs = torch.ones_like(i)  # [B, HxW]93    xs = (i - cx) / fx * zs94    ys = (j - cy) / fy * zs95    zs = zs.expand_as(ys)96 97    directions = torch.stack((xs, ys, zs), dim=-1)  # B, V, HW, 398    directions = directions / directions.norm(dim=-1, keepdim=True)  # B, V, HW, 399 100    rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2)  # B, V, 3, HW101    rays_o = c2w[..., :3, 3]  # B, V, 3102    rays_o = rays_o[:, :, None].expand_as(rays_d)  # B, V, 3, HW103    # c2w @ dirctions104    rays_dxo = torch.linalg.cross(rays_o, rays_d)105    plucker = torch.cat([rays_dxo, rays_d], dim=-1)106    plucker = plucker.reshape(B, c2w.shape[1], H, W, 6)  # B, V, H, W, 6107    return plucker108 109 110def get_pipeline(ori_model_path, unet_subfolder, down_block_types, up_block_types, pose_encoder_kwargs,111                 attention_processor_kwargs, pose_adaptor_ckpt, enable_xformers, device):112    noise_scheduler = EulerDiscreteScheduler.from_pretrained(ori_model_path, subfolder="scheduler")113    feature_extractor = CLIPImageProcessor.from_pretrained(ori_model_path, subfolder="feature_extractor")114    image_encoder = CLIPVisionModelWithProjection.from_pretrained(ori_model_path, subfolder="image_encoder")115    vae = AutoencoderKLTemporalDecoder.from_pretrained(ori_model_path, subfolder="vae")116    unet = UNetSpatioTemporalConditionModelPoseCond.from_pretrained(ori_model_path,117                                                                    subfolder=unet_subfolder,118                                                                    down_block_types=down_block_types,119                                                                    up_block_types=up_block_types)120    pose_encoder = CameraPoseEncoder(**pose_encoder_kwargs)121    print("Setting the attention processors")122    unet.set_pose_cond_attn_processor(enable_xformers=(enable_xformers and is_xformers_available()), **attention_processor_kwargs)123    print(f"Loading weights of camera encoder and attention processor from {pose_adaptor_ckpt}")124    ckpt_dict = torch.load(pose_adaptor_ckpt, map_location=unet.device)125    pose_encoder_state_dict = ckpt_dict['pose_encoder_state_dict']126    pose_encoder_m, pose_encoder_u = pose_encoder.load_state_dict(pose_encoder_state_dict)127    assert len(pose_encoder_m) == 0 and len(pose_encoder_u) == 0128    attention_processor_state_dict = ckpt_dict['attention_processor_state_dict']129    _, attention_processor_u = unet.load_state_dict(attention_processor_state_dict, strict=False)130    assert len(attention_processor_u) == 0131    print("Loading done")132    vae.to(device)133    image_encoder.to(device)134    unet.to(device)135    pipeline = StableVideoDiffusionPipelinePoseCond(136        vae=vae,137        image_encoder=image_encoder,138        unet=unet,139        scheduler=noise_scheduler,140        feature_extractor=feature_extractor,141        pose_encoder=pose_encoder142    )143    pipeline = pipeline.to(device)144    return pipeline145 146 147def main(args):148    os.makedirs(os.path.join(args.out_root, 'generated_videos'), exist_ok=True)149    os.makedirs(os.path.join(args.out_root, 'reference_images'), exist_ok=True)150    rank = args.local_rank151    setup_for_distributed(rank == 0)152    gpu_id = rank % torch.cuda.device_count()153    model_configs = OmegaConf.load(args.model_config)154    device = f"cuda:{gpu_id}"155    print(f'Constructing pipeline')156    pipeline = get_pipeline(args.ori_model_path, model_configs['unet_subfolder'], model_configs['down_block_types'],157                            model_configs['up_block_types'], model_configs['pose_encoder_kwargs'],158                            model_configs['attention_processor_kwargs'], args.pose_adaptor_ckpt, args.enable_xformers, device)159    160    print('Done')161 162    print('Loading K, R, t matrix')163    with open(args.trajectory_file, 'r') as f:164        poses = f.readlines()165    poses = [pose.strip().split(' ') for pose in poses[1:]]166    cam_params = [[float(x) for x in pose] for pose in poses]167    cam_params = [Camera(cam_param) for cam_param in cam_params]168 169    sample_wh_ratio = args.image_width / args.image_height170    pose_wh_ratio = args.original_pose_width / args.original_pose_height171    if pose_wh_ratio > sample_wh_ratio:172        resized_ori_w = args.image_height * pose_wh_ratio173        for cam_param in cam_params:174            cam_param.fx = resized_ori_w * cam_param.fx / args.image_width175    else:176        resized_ori_h = args.image_width / pose_wh_ratio177        for cam_param in cam_params:178            cam_param.fy = resized_ori_h * cam_param.fy / args.image_height179    intrinsic = np.asarray([[cam_param.fx * args.image_width,180                             cam_param.fy * args.image_height,181                             cam_param.cx * args.image_width,182                             cam_param.cy * args.image_height]183                            for cam_param in cam_params], dtype=np.float32)184    K = torch.as_tensor(intrinsic)[None]  # [1, 1, 4]185    c2ws = get_relative_pose(cam_params, zero_first_frame_scale=True)186    c2ws = torch.as_tensor(c2ws)[None]  # [1, n_frame, 4, 4]187    plucker_embedding = ray_condition(K, c2ws, args.image_height, args.image_width, device='cpu')       # b f h w 6188    plucker_embedding = plucker_embedding.permute(0, 1, 4, 2, 3).contiguous().to(device=device)189 190    prompt_dict = json.load(open(args.prompt_file, 'r'))191    prompt_images = prompt_dict['image_paths']192    prompt_captions = prompt_dict['captions']193    N = int(len(prompt_images) // args.n_procs)194    remainder = int(len(prompt_images) % args.n_procs)195    prompts_per_gpu = [N + 1 if gpu_id < remainder else N for gpu_id in range(args.n_procs)]196    low_idx = sum(prompts_per_gpu[:gpu_id])197    high_idx = low_idx + prompts_per_gpu[gpu_id]198    prompt_images = prompt_images[low_idx: high_idx]199    prompt_captions = prompt_captions[low_idx: high_idx]200    print(f"rank {rank} / {torch.cuda.device_count()}, number of prompts: {len(prompt_images)}")201 202    generator = torch.Generator(device=device)203    generator.manual_seed(42)204 205    for prompt_image, prompt_caption in tqdm(zip(prompt_images, prompt_captions)):206        # save_name = "_".join(prompt_caption.split(" "))207        save_name = prompt_caption.split('.')[0]208        condition_image = Image.open(prompt_image)209        with torch.no_grad():210            sample = pipeline(211                image=condition_image,212                pose_embedding=plucker_embedding,213                height=args.image_height,214                width=args.image_width,215                num_frames=args.num_frames,216                num_inference_steps=args.num_inference_steps,217                min_guidance_scale=args.min_guidance_scale,218                max_guidance_scale=args.max_guidance_scale,219                do_image_process=True,220                generator=generator,221                output_type='pt'222            ).frames[0].transpose(0, 1).cpu()      # [3, f, h, w] 0-1223        resized_condition_image = condition_image.resize((args.image_width, args.image_height))224        # save_videos_grid(sample[None], f"{os.path.join(args.out_root, 'generated_videos')}/{save_name}.mp4", rescale=False)225        save_videos_grid(sample[None], f"{os.path.join(args.out_root, 'generated_videos')}/{save_name}.gif", rescale=False)226        resized_condition_image.save(os.path.join(args.out_root, 'reference_images', f'{save_name}.png'))227 228 229if __name__ == '__main__':230    parser = argparse.ArgumentParser()231    parser.add_argument("--out_root", type=str)232    parser.add_argument("--image_height", type=int, default=320)233    parser.add_argument("--image_width", type=int, default=576)234    parser.add_argument("--num_frames", type=int, default=14, help="14 for svd and 25 for svd-xt", choices=[14, 25])235    parser.add_argument("--ori_model_path", type=str)236    parser.add_argument("--unet_subfolder", type=str, default='unet')237    parser.add_argument("--enable_xformers", action='store_true')238    parser.add_argument("--pose_adaptor_ckpt", default=None)239    parser.add_argument("--num_inference_steps", type=int, default=25)240    parser.add_argument("--min_guidance_scale", type=float, default=1.0)241    parser.add_argument("--max_guidance_scale", type=float, default=3.0)242    parser.add_argument("--prompt_file", required=True, help='prompts path, json or txt')243    parser.add_argument("--trajectory_file", required=True)244    parser.add_argument("--original_pose_width", type=int, default=1280)245    parser.add_argument("--original_pose_height", type=int, default=720)246    parser.add_argument("--model_config", required=True)247    parser.add_argument("--n_procs", type=int, default=8)248 249    # DDP args250    parser.add_argument("--world_size", default=1, type=int,251                        help="number of the distributed processes.")252    parser.add_argument('--local-rank', type=int, default=-1,253                        help='Replica rank on the current node. This field is required '254                             'by `torch.distributed.launch`.')255    args = parser.parse_args()256    main(args)257