237 lines
11 KiB
Python
237 lines
11 KiB
Python
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()
|