import sys import os from os.path import dirname from typing import Dict, List import numpy as np import torch import torch.utils.checkpoint from diffusers import AutoencoderKL, DDIMScheduler from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor, CLIPVisionModelWithProjection from torchvision import transforms from CharacterGen.Stage_2D.tuneavideo.models.unet_mv2d_condition import UNetMV2DConditionModel from CharacterGen.Stage_2D.tuneavideo.models.unet_mv2d_ref import UNetMV2DRefModel from CharacterGen.Stage_2D.tuneavideo.models.PoseGuider import PoseGuider from CharacterGen.Stage_2D.tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline from einops import rearrange from PIL import Image import json import trimesh from CharacterGen.Stage_3D import lrm from datetime import datetime from pygltflib import GLTF2 import pymeshlab CHARACTER_GEN_ROOT_ABS_PATH = dirname(__file__) CHARACTER_GEN_2D_DATA_ABS_PATH = os.path.join(CHARACTER_GEN_ROOT_ABS_PATH, "Stage_2D/material") CHARACTER_GEN_3D_ABS_PATH = os.path.join(CHARACTER_GEN_ROOT_ABS_PATH, "Stage_3D") CHARACTER_GEN_3D_DATA_ABS_PATH = os.path.join(CHARACTER_GEN_3D_ABS_PATH, "material") sys.path.append(CHARACTER_GEN_3D_ABS_PATH) def get_bg_color(bg_color): if bg_color == 'white': bg_color = np.array([1., 1., 1.], dtype=np.float32) elif bg_color == 'black': bg_color = np.array([0., 0., 0.], dtype=np.float32) elif bg_color == 'gray': bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32) elif bg_color == 'random': bg_color = np.random.rand(3) elif isinstance(bg_color, float): bg_color = np.array([bg_color] * 3, dtype=np.float32) else: raise NotImplementedError return bg_color def process_image(image, totensor, w=512, h=768): if not image.mode == "RGBA": image = image.convert("RGBA") # Find non-transparent pixels non_transparent = np.nonzero(np.array(image)[..., 3]) min_x, max_x = non_transparent[1].min(), non_transparent[1].max() min_y, max_y = non_transparent[0].min(), non_transparent[0].max() image = image.crop((min_x, min_y, max_x, max_y)) # paste to center max_dim = max(image.width, image.height) max_height = max_dim max_width = int(max_dim / 3 * 2) new_image = Image.new("RGBA", (max_width, max_height)) left = (max_width - image.width) // 2 top = (max_height - image.height) // 2 new_image.paste(image, (left, top)) image = new_image.resize((w, h), resample=Image.Resampling.BICUBIC) image = np.array(image) image = image.astype(np.float32) / 255. assert image.shape[-1] == 4 # RGBA alpha = image[..., 3:4] bg_color = get_bg_color("gray") image = image[..., :3] * alpha + bg_color * (1 - alpha) # save image new_image = Image.fromarray((image * 255).astype(np.uint8)) return totensor(image) class Inference2D_API: def __init__(self, pretrained_model_path: str, checkpoint_root_path: str, image_encoder_path: str, ckpt_dir: str, validation: Dict, local_crossattn: bool = True, unet_from_pretrained_kwargs=None, unet_condition_type=None, use_pose_guider=False, use_shifted_noise=False, use_noise=True, device="cuda", weight_dtype=torch.float16 ): self.validation = validation self.use_noise = use_noise self.use_shifted_noise = use_shifted_noise self.unet_condition_type = unet_condition_type self.device = device self.weight_dtype = weight_dtype image_encoder_path = os.path.join(checkpoint_root_path, image_encoder_path) ckpt_dir = os.path.join(checkpoint_root_path, ckpt_dir) self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer") text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder") image_encoder = CLIPVisionModelWithProjection.from_pretrained(image_encoder_path) feature_extractor = CLIPImageProcessor() vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae") unet = UNetMV2DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", local_crossattn=local_crossattn, **unet_from_pretrained_kwargs) ref_unet = UNetMV2DRefModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", local_crossattn=local_crossattn, **unet_from_pretrained_kwargs) if use_pose_guider: pose_guider = PoseGuider(noise_latent_channels=4).to(device) else: pose_guider = None unet_params = torch.load(os.path.join(ckpt_dir, "pytorch_model.bin"), map_location="cpu") if use_pose_guider: pose_guider_params = torch.load(os.path.join(ckpt_dir, "pytorch_model_1.bin"), map_location="cpu") ref_unet_params = torch.load(os.path.join(ckpt_dir, "pytorch_model_2.bin"), map_location="cpu") pose_guider.load_state_dict(pose_guider_params) else: ref_unet_params = torch.load(os.path.join(ckpt_dir, "pytorch_model_1.bin"), map_location="cpu") unet.load_state_dict(unet_params) ref_unet.load_state_dict(ref_unet_params) text_encoder.to(device, dtype=weight_dtype) image_encoder.to(device, dtype=weight_dtype) vae.to(device, dtype=weight_dtype) ref_unet.to(device, dtype=weight_dtype) unet.to(device, dtype=weight_dtype) vae.requires_grad_(False) unet.requires_grad_(False) ref_unet.requires_grad_(False) noise_scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler") self.validation_pipeline = TuneAVideoPipeline( vae=vae, text_encoder=text_encoder, tokenizer=self.tokenizer, unet=unet, ref_unet=ref_unet,feature_extractor=feature_extractor,image_encoder=image_encoder, scheduler=noise_scheduler ) self.validation_pipeline.enable_vae_slicing() self.validation_pipeline.set_progress_bar_config(disable=True) @torch.no_grad() def inference( self, input_image, val_width, val_height, prompt="high quality, best quality", prompt_neg=None, guidance_scale=7.5, num_inference_steps=40, seed=100, use_shifted_noise=False, crop=False ): generator = torch.Generator(device=self.device).manual_seed(seed) totensor = transforms.ToTensor() metas = json.load(open(os.path.join(CHARACTER_GEN_2D_DATA_ABS_PATH, "pose.json"), "r")) cameras = [] pose_images = [] for lm in metas: cameras.append(torch.tensor(np.array(lm[0]).reshape(4, 4).transpose(1,0)[:3, :4]).reshape(-1)) if not crop: pose_images.append(totensor(np.asarray(Image.open(os.path.join(CHARACTER_GEN_2D_DATA_ABS_PATH, lm[1])).resize( (val_height, val_width), resample=Image.Resampling.BICUBIC)).astype(np.float32) / 255.)) else: pose_image = Image.open(os.path.join(CHARACTER_GEN_2D_DATA_ABS_PATH, lm[1])) crop_area = (128, 0, 640, 768) pose_images.append(totensor(np.array(pose_image.crop(crop_area)).astype(np.float32)) / 255.) camera_matrixs = torch.stack(cameras).unsqueeze(0).to(self.device) pose_imgs_in = torch.stack(pose_images).to(self.device) prompt_ids = self.tokenizer( prompt, max_length=self.tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt" ).input_ids[0] # (B*Nv, 3, H, W) B==1 imgs_in = process_image(input_image, totensor, val_width, val_height) imgs_in = rearrange(imgs_in.unsqueeze(0).unsqueeze(0), "B Nv C H W -> (B Nv) C H W") with torch.autocast(self.device, dtype=self.weight_dtype): imgs_in = imgs_in.to(self.device) # B*Nv images out = self.validation_pipeline(prompt=prompt, negative_prompt=prompt_neg, image=imgs_in.to(self.weight_dtype), generator=generator, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, camera_matrixs=camera_matrixs.to(self.weight_dtype), prompt_ids=prompt_ids, height=val_height, width=val_width, unet_condition_type=self.unet_condition_type, pose_guider=None, pose_image=pose_imgs_in, use_noise=self.use_noise, use_shifted_noise=use_shifted_noise, **self.validation).videos out = rearrange(out, "B C f H W -> (B f) H W C", f=self.validation.video_length) torch.cuda.empty_cache() return out class Inference3D_API: def __init__(self, checkpoint_root_path, cfg, device="cuda"): self.cfg = cfg self.device = device self.cfg.system.weights = os.path.join(checkpoint_root_path, self.cfg.system.weights) self.cfg.system.image_tokenizer.pretrained_model_name_or_path = os.path.join(checkpoint_root_path, self.cfg.system.image_tokenizer.pretrained_model_name_or_path) self.cfg.system.renderer.tet_dir = os.path.join(CHARACTER_GEN_3D_DATA_ABS_PATH, self.cfg.system.renderer.tet_dir) self.system = lrm.find(self.cfg.system_cls)(self.cfg.system).to(self.device) self.system.eval() def inference(self, mv_imgs: List[Image.Image]): meta = json.load(open(os.path.join(CHARACTER_GEN_3D_DATA_ABS_PATH, "meta.json"))) c2w_cond = [np.array(loc["transform_matrix"]) for loc in meta["locations"]] c2w_cond = torch.from_numpy(np.stack(c2w_cond, axis=0)).float()[None].to(self.device) rgb_cond = [] new_images = [] for pil_img in mv_imgs: image = np.array(pil_img) image = Image.fromarray(image) if image.width != image.height: max_dim = max(image.width, image.height) new_image = Image.new("RGBA", (max_dim, max_dim)) left = (max_dim - image.width) // 2 top = (max_dim - image.height) // 2 new_image.paste(image, (left, top)) image = new_image image = image.convert('RGB') rgb = image.resize((self.cfg.data.cond_width, self.cfg.data.cond_height), Image.Resampling.BILINEAR) rgb = np.array(rgb).astype(np.float32) / 255.0 new_images.append(image) rgb_cond.append(rgb) assert len(rgb_cond) == 4, "Please provide 4 images" rgb_cond = torch.from_numpy(np.stack(rgb_cond, axis=0)).float()[None].to(self.device) scene_code = self.system({"rgb_cond": rgb_cond, "c2w_cond": c2w_cond})[0] mesh = self.system.get_geometry(scene_code) # Flip mesh horizontally #mesh.v_pos[:, 0] = -mesh.v_pos[:, 0] return mesh.v_pos.detach(), mesh.t_pos_idx.detach()