app.py

import gradio as gr
from PIL import Image

import spaces

import io
import argparse
import os
import random
from typing import Dict, Optional, Tuple
from omegaconf import OmegaConf
import numpy as np

import torch
import torch.utils.checkpoint

from accelerate.logging import get_logger
from accelerate.utils import set_seed
from diffusers import AutoencoderKL, DDIMScheduler
from diffusers.utils import check_min_version
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor, CLIPVisionModelWithProjection
from torchvision import transforms

import sys

sys.path.append("2D_Stage")
sys.path.append("3D_Stage")
from tuneavideo.models.unet_mv2d_condition import UNetMV2DConditionModel
from tuneavideo.models.unet_mv2d_ref import UNetMV2DRefModel
from tuneavideo.models.PoseGuider import PoseGuider
from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline
from tuneavideo.util import shifted_noise
from einops import rearrange
import PIL
from PIL import Image
from torchvision.utils import save_image
import json
import cv2

import lrm
import trimesh
from lrm.utils.config import load_config
from refine import refine
from datetime import datetime
import gradio as gr
from pygltflib import GLTF2

import onnxruntime as rt
from huggingface_hub import hf_hub_download, list_repo_files
from rm_anime_bg.cli import get_mask, SCALE
import pymeshlab

# Patch for gradio_client bug: argument of type 'bool' is not iterable in get_type()
# JSON Schema allows additionalProperties to be a boolean, but gradio_client 5.14.x
# passes it directly to get_type() which expects a dict.
try:
    import gradio_client.utils as _gc_utils
    _orig_get_type = _gc_utils.get_type

    def _patched_get_type(schema):
        if isinstance(schema, bool):
            return "unknown"
        return _orig_get_type(schema)

    _gc_utils.get_type = _patched_get_type
except Exception:
    pass


def download_model_files():
    """モデルファイルをHuggingFace Hubからダウンロードする（初期化時に実行）"""
    repo_id = "zjpshadow/CharacterGen"
    all_files = list_repo_files(repo_id, revision="main")
    for file in all_files:
        if os.path.exists(file):
            continue
        if file.startswith("2D_Stage") or file.startswith("3D_Stage"):
            hf_hub_download(repo_id, file, local_dir=".")


check_min_version("0.24.0")

logger = get_logger(__name__, log_level="INFO")

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

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):
    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((512, 768), resample=PIL.Image.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))
    new_image.save("input.png")
    return totensor(image)

class rm_bg_api:

    def __init__(self, force_cpu: Optional[bool] = True):
        session_infer_path = hf_hub_download(
            repo_id="skytnt/anime-seg", filename="isnetis.onnx",
        )
        providers: list[str] = ["CPUExecutionProvider"]
        if not force_cpu and "CUDAExecutionProvider" in rt.get_available_providers():
            providers = ["CUDAExecutionProvider"]

        self.session_infer = rt.InferenceSession(
            session_infer_path, providers=providers,
        )

    def _remove_background_impl(
        self,
        imgs: list[np.ndarray],
        alpha_min: float,
        alpha_max: float,
    ) -> list:
        process_imgs = []
        for img in imgs:
            img = np.array(img)
            # CHANGE to RGB
            if img.shape[-1] == 4:
                img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
            mask = get_mask(self.session_infer, img)

            mask[mask < alpha_min] = 0.0  # type: ignore
            mask[mask > alpha_max] = 1.0  # type: ignore

            img_after = (mask * img).astype(np.uint8)  # type: ignore
            mask = (mask * SCALE).astype(np.uint8)  # type: ignore
            img_after = np.concatenate([img_after, mask], axis=2, dtype=np.uint8)
            mask = mask.repeat(3, axis=2)
            process_imgs.append(Image.fromarray(img_after))
        return process_imgs


class Inference2D_API:

    def __init__(self,
            pretrained_model_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"
        ):
        self.validation = validation
        self.use_noise = use_noise
        self.use_shifted_noise = use_shifted_noise
        self.unet_condition_type = unet_condition_type
        image_encoder_path = image_encoder_path.replace("./", "./2D_Stage/")
        ckpt_dir = ckpt_dir.replace("./", "./2D_Stage/")

        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("cuda")
        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)

        weight_dtype = torch.float16

        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)
        self.generator = torch.Generator(device=device)

    @torch.no_grad()
    def _inference_impl(self, input_image, val_width, val_height,
                    use_shifted_noise=False, crop=False, seed=100, timestep=20):
        set_seed(seed)
        totensor = transforms.ToTensor()

        metas = json.load(open("./2D_Stage/material/pose.json", "r"))
        cameras = []
        pose_images = []
        input_path = "./2D_Stage/material"
        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(input_path, lm[1])).resize(
                    (val_height, val_width), resample=PIL.Image.BICUBIC)).astype(np.float32) / 255.))
            else:
                pose_image = Image.open(os.path.join(input_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("cuda")
        pose_imgs_in = torch.stack(pose_images).to("cuda")
        prompts = "high quality, best quality"
        prompt_ids = self.tokenizer(
            prompts, 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
        weight_dtype = torch.bfloat16
        imgs_in = process_image(input_image, totensor)
        imgs_in = rearrange(imgs_in.unsqueeze(0).unsqueeze(0), "B Nv C H W -> (B Nv) C H W")

        with torch.autocast("cuda", dtype=weight_dtype):
            imgs_in = imgs_in.to("cuda")
            # B*Nv images
            out = self.validation_pipeline(prompt=prompts, image=imgs_in.to(weight_dtype), generator=self.generator,
                                        num_inference_steps=timestep,
                                        camera_matrixs=camera_matrixs.to(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) C H W", f=self.validation.video_length)

        image_outputs = []
        for bs in range(4):
            img_buf = io.BytesIO()
            save_image(out[bs], img_buf, format='PNG')
            img_buf.seek(0)
            img = Image.open(img_buf)
            image_outputs.append(img)
        torch.cuda.empty_cache()
        return image_outputs


def traverse(path, back_proj, smooth_iter):
    ms = pymeshlab.MeshSet()
    ms.load_new_mesh(f"{path}/model-00.obj")
    image = Image.open(f"{path}/{'refined_texture_kd.jpg' if back_proj else 'texture_kd.jpg'}")
    out_image_path = f"{path}/{'refined_texture_kd.png' if back_proj else 'texture_kd.png'}"
    image.save(out_image_path, 'PNG')
    ms.set_texture_per_mesh(textname=f"{path}/{'refined_texture_kd.png' if back_proj else 'texture_kd.png'}")
    ms.meshing_merge_close_vertices()
    ms.apply_coord_laplacian_smoothing(stepsmoothnum=smooth_iter)
    ms.save_current_mesh(f"{path}/temp-00.obj", save_vertex_normal=False, save_wedge_normal=False, save_vertex_color=False)

    mesh = trimesh.load(f"{path}/temp-00.obj", process=False)
    mesh.apply_transform(trimesh.transformations.rotation_matrix(np.radians(90.0), [-1, 0, 0]))
    mesh.apply_transform(trimesh.transformations.rotation_matrix(np.radians(180.0), [0, 1, 0]))

    mesh.export(f'{path}/output.glb', file_type='glb')

    image = Image.open(f"{path}/{'refined_texture_kd.png' if back_proj else 'texture_kd.png'}")
    texture = np.array(image)
    vertex_colors = np.zeros((mesh.vertices.shape[0], 4), dtype=np.uint8)

    for vertex_index in range(len(mesh.visual.uv)):
        uv = mesh.visual.uv[vertex_index]
        x = int(uv[0] * (texture.shape[1] - 1))
        y = int((1 - uv[1]) * (texture.shape[0] - 1))

        color = texture[y, x, :3]
        vertex_colors[vertex_index] = [color[0], color[1], color[2], 255]
    return trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, vertex_colors=vertex_colors, process=False)

class Inference3D_API:

    def __init__(self, device="cuda"):
        self.cfg = load_config("3D_Stage/configs/infer.yaml", makedirs=False)
        print("Loading system")
        self.device = device
        self.cfg.system.weights = self.cfg.system.weights.replace("./", "./3D_Stage/")
        self.cfg.system.image_tokenizer.pretrained_model_name_or_path = \
            self.cfg.system.image_tokenizer.pretrained_model_name_or_path.replace("./", "./3D_Stage/")
        self.cfg.system.renderer.tet_dir = self.cfg.system.renderer.tet_dir.replace("./", "./3D_Stage/")
        self.cfg.system.exporter.output_path = self.cfg.system.exporter.output_path.replace("./", "./3D_Stage/")
        self.system = lrm.find(self.cfg.system_cls)(self.cfg.system).to(self.device)
        self.system.eval()

    def _process_images_impl(self, img_input0, img_input1, img_input2, img_input3, back_proj, smooth_iter):
        meta = json.load(open("./3D_Stage/material/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 = []
        files = [img_input0, img_input1, img_input2, img_input3]
        new_images = []
        for file in files:
            image = np.array(file)
            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.save("input_3D.png")

            image = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2RGB)
            rgb = cv2.resize(image, (self.cfg.data.cond_width,
                                     self.cfg.data.cond_height)).astype(np.float32) / 255.0
            new_images.append(Image.fromarray(image.astype(np.uint8)).convert("RGB"))
            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)

        with torch.no_grad():
            scene_codes = self.system({"rgb_cond": rgb_cond, "c2w_cond": c2w_cond})
            exporter_output = self.system.exporter([f"{i:02d}" for i in range(rgb_cond.shape[0])], scene_codes)

        save_dir = os.path.join("./3D_Stage/outputs", datetime.now().strftime("@%Y%m%d-%H%M%S"))
        os.makedirs(save_dir, exist_ok=True)
        self.system.set_save_dir(save_dir)

        for out in exporter_output:
            save_func_name = f"save_{out.save_type}"
            save_func = getattr(self.system, save_func_name)
            save_func(f"{out.save_name}", **out.params)
        if back_proj:
            refine(save_dir, new_images[1], new_images[0], new_images[3], new_images[2])

        new_obj = traverse(save_dir, back_proj, smooth_iter)
        new_obj.export(f'{save_dir}/output.obj', file_type='obj')

        gltf = GLTF2().load(f'{save_dir}/output.glb')
        for material in gltf.materials:
            if material.pbrMetallicRoughness:
                material.pbrMetallicRoughness.baseColorFactor = [1.0, 1.0, 1.0, 100.0]
                material.pbrMetallicRoughness.metallicFactor = 0.0
                material.pbrMetallicRoughness.roughnessFactor = 1.0
        gltf.save(f'{save_dir}/output.glb')

        return save_dir, f"{save_dir}/output.obj", f"{save_dir}/output.glb"


# モジュールレベルのシングルトンインスタンス（遅延初期化）
_remove_api = None
_infer2dapi = None
_infer3dapi = None


@spaces.GPU
def run_remove_background(imgs, alpha_min, alpha_max):
    global _remove_api
    if _remove_api is None:
        _remove_api = rm_bg_api()
    return _remove_api._remove_background_impl(imgs, alpha_min, alpha_max)


@spaces.GPU
def run_inference2d(input_image, val_width, val_height,
                    use_shifted_noise=False, crop=False, seed=100, timestep=20):
    global _infer2dapi
    if _infer2dapi is None:
        download_model_files()
        _infer2dapi = Inference2D_API(**OmegaConf.load("./2D_Stage/configs/infer.yaml"))
    return _infer2dapi._inference_impl(input_image, val_width, val_height,
                                       use_shifted_noise, crop, seed, timestep)


@spaces.GPU
def run_inference3d(img_input0, img_input1, img_input2, img_input3, back_proj, smooth_iter):
    global _infer3dapi
    if _infer3dapi is None:
        download_model_files()
        _infer3dapi = Inference3D_API()
    return _infer3dapi._process_images_impl(img_input0, img_input1, img_input2, img_input3,
                                            back_proj, smooth_iter)



def gen4views(image, width, height, seed, timestep, remove_bg):
    if remove_bg:
        image = run_remove_background(
            imgs=[np.array(image)],
            alpha_min=0.1,
            alpha_max=0.9,
        )[0]
    return run_remove_background(
        imgs=run_inference2d(
            image, width, height, crop=True, seed=seed, timestep=timestep
        ), alpha_min=0.2, alpha_max=0.9)


def gen3dmesh(img_input0, img_input1, img_input2, img_input3, back_proj, smooth_iter):
    return run_inference3d(img_input0, img_input1, img_input2, img_input3, back_proj, smooth_iter)



with gr.Blocks() as demo:
    gr.Markdown("# [SIGGRAPH'24] CharacterGen: Efficient 3D Character Generation from Single Images with Multi-View Pose Calibration")
    with gr.Row():
        with gr.Column(variant="panel"):
            img_input = gr.Image(type="pil", label="Upload Image(without background)", image_mode="RGBA", width=768, height=512)
            with gr.Row():
                width_input = gr.Number(label="Width", value=512)
                height_input = gr.Number(label="Height", value=768)
                seed_input = gr.Number(label="Seed", value=2333)
                remove_bg = gr.Checkbox(label="Remove Background (with algorithm)", value=True)
        with gr.Column(variant="panel"):
            timestep = gr.Slider(minimum=10, maximum=70, step=1, value=40, label="Timesteps")
            button1 = gr.Button(value="Generate 4 Views")
            with gr.Row():
                img_input0 = gr.Image(type="pil", label="Back Image", image_mode="RGBA", width=256, height=384)
                img_input1 = gr.Image(type="pil", label="Front Image", image_mode="RGBA", width=256, height=384)
            with gr.Row():
                img_input2 = gr.Image(type="pil", label="Right Image", image_mode="RGBA", width=256, height=384)
                img_input3 = gr.Image(type="pil", label="Left Image", image_mode="RGBA", width=256, height=384)
        with gr.Column(variant="panel"):
            smooth_iter = gr.Slider(minimum=0, maximum=10, step=1, value=5, label="Laplacian Smoothing Iterations")
            with gr.Row():
                back_proj = gr.Checkbox(label="Back Projection")
                button2 = gr.Button(value="Generate 3D Mesh")
            output_dir = gr.Textbox(label="Output Directory")
            with gr.Row():
                with gr.Tab("GLB"):
                    output_model_glb = gr.Model3D( label="Output Model (GLB Format)", height=512)
                    gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
                with gr.Tab("OBJ"):
                    output_model_obj = gr.Model3D( label="Output Model (OBJ Format)")
                    gr.Markdown("Note: The model shown here's texture is mapped to vertex. Download to get correct results.")
    button1.click(
        fn=gen4views,
        inputs=[img_input, width_input, height_input, seed_input, timestep, remove_bg],
        outputs=[img_input2, img_input0, img_input3, img_input1],
        api_name="generate_views"
    )
    button2.click(
        gen3dmesh,
        inputs=[img_input0, img_input1, img_input2, img_input3, back_proj, smooth_iter],
        outputs=[output_dir, output_model_obj, output_model_glb],
        api_name="generate_3d_mesh"
    )

demo.queue()
demo.launch(show_api=False)