hy3dgen/texgen/pipelines.py

# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
# except for the third-party components listed below.
# Hunyuan 3D does not impose any additional limitations beyond what is outlined
# in the repsective licenses of these third-party components.
# Users must comply with all terms and conditions of original licenses of these third-party
# components and must ensure that the usage of the third party components adheres to
# all relevant laws and regulations.

# For avoidance of doubts, Hunyuan 3D means the large language models and
# their software and algorithms, including trained model weights, parameters (including
# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
# fine-tuning enabling code and other elements of the foregoing made publicly available
# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.


import logging
import numpy as np
import os
import torch
from PIL import Image
from typing import Union, Optional
from pathlib import Path

from .differentiable_renderer.mesh_render import MeshRender
from .utils.dehighlight_utils import Light_Shadow_Remover
from .utils.multiview_utils import Multiview_Diffusion_Net
# from .utils.imagesuper_utils import Image_Super_Net
from .utils.uv_warp_utils import mesh_uv_wrap

logger = logging.getLogger(__name__)

# -------------------------------------------
# Device Selection (Global clean handling)
# -------------------------------------------

def get_best_device():
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


class Hunyuan3DTexGenConfig:
    def __init__(self, light_remover_ckpt_path, multiview_ckpt_path):

        # Old: self.device = 'cuda'
        self.device = get_best_device()

        self.light_remover_ckpt_path = light_remover_ckpt_path
        self.multiview_ckpt_path = multiview_ckpt_path

        self.candidate_camera_azims = [0, 90, 180, 270, 0, 180]
        self.candidate_camera_elevs = [0, 0, 0, 0, 90, -90]
        self.candidate_view_weights = [1, 0.1, 0.5, 0.1, 0.05, 0.05]

        self.render_size = 2048
        self.texture_size = 2048
        self.bake_exp = 4
        self.merge_method = 'fast'


class Hunyuan3DPaintPipeline:
    @classmethod
    def from_pretrained(cls, model_path):
        original_model_path = model_path
        print(f"原始路径 original_model_path: {model_path}")

        if not os.path.exists(model_path):

            print(f"不存在原始路径: {model_path}")
            base_dir = os.environ.get('HY3DGEN_MODELS', '~/.cache/hy3dgen')
            model_path = os.path.expanduser(os.path.join(base_dir, model_path))

            delight_model_path = os.path.join(model_path, 'hunyuan3d-delight-v2-0')
            multiview_model_path = os.path.join(model_path, 'hunyuan3d-paint-v2-0')

            if not os.path.exists(delight_model_path) or not os.path.exists(multiview_model_path):
                try:
                    import huggingface_hub

                    model_path = huggingface_hub.snapshot_download(
                        repo_id=original_model_path,
                        allow_patterns=["hunyuan3d-delight-v2-0/*"]
                    )
                    model_path = huggingface_hub.snapshot_download(
                        repo_id=original_model_path,
                        allow_patterns=["hunyuan3d-paint-v2-0/*"]
                    )

                    delight_model_path = os.path.join(model_path, 'hunyuan3d-delight-v2-0')
                    multiview_model_path = os.path.join(model_path, 'hunyuan3d-paint-v2-0')

                    return cls(Hunyuan3DTexGenConfig(delight_model_path, multiview_model_path))

                except Exception as e:
                    print("构造 Hunyuan3DPaintPipeline 实例时出错：", e)
                    raise

            else:
                return cls(Hunyuan3DTexGenConfig(delight_model_path, multiview_model_path))

        raise FileNotFoundError(f"Model path {original_model_path} not found and Hub download failed.")


    def __init__(self, config):
        self.config = config
        self.models = {}

        self.render = MeshRender(
            default_resolution=self.config.render_size,
            texture_size=self.config.texture_size
        )

        self.load_models()


    # -------------------------------------------
    # Load Models — Dynamic CUDA handling
    # -------------------------------------------
    def load_models(self):

        # Originally forced CUDA:
        # torch.cuda.empty_cache()

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        self.models['delight_model'] = Light_Shadow_Remover(self.config)
        self.models['multiview_model'] = Multiview_Diffusion_Net(self.config)
        # self.models['super_model'] = Image_Super_Net(self.config)


    def enable_model_cpu_offload(
        self, 
        gpu_id: Optional[int] = None, 
        device: Union[torch.device, str] = None
    ):
        if device is None:
            device = self.config.device

        if hasattr(self.models['delight_model'], "pipeline"):
            self.models['delight_model'].pipeline.enable_model_cpu_offload(
                gpu_id=gpu_id, device=device
            )

        if hasattr(self.models['multiview_model'], "pipeline"):
            self.models['multiview_model'].pipeline.enable_model_cpu_offload(
                gpu_id=gpu_id, device=device
            )


    # -------------------------------------------
    # Rendering functions unchanged
    # -------------------------------------------
    def render_normal_multiview(self, camera_elevs, camera_azims, use_abs_coor=True):
        normal_maps = []
        for elev, azim in zip(camera_elevs, camera_azims):
            normal_map = self.render.render_normal(
                elev, azim, use_abs_coor=use_abs_coor, return_type='pl')
            normal_maps.append(normal_map)

        return normal_maps


    def render_position_multiview(self, camera_elevs, camera_azims):
        position_maps = []
        for elev, azim in zip(camera_elevs, camera_azims):
            position_map = self.render.render_position(
                elev, azim, return_type='pl')
            position_maps.append(position_map)
        return position_maps


    def bake_from_multiview(self, views, camera_elevs,
                            camera_azims, view_weights, method='graphcut'):
        project_textures, project_weighted_cos_maps = [], []
        project_boundary_maps = []

        for view, camera_elev, camera_azim, weight in zip(
            views, camera_elevs, camera_azims, view_weights
        ):
            project_texture, project_cos_map, project_boundary_map = self.render.back_project(
                view, camera_elev, camera_azim
            )

            project_cos_map = weight * (project_cos_map ** self.config.bake_exp)
            project_textures.append(project_texture)
            project_weighted_cos_maps.append(project_cos_map)
            project_boundary_maps.append(project_boundary_map)

        if method == 'fast':
            texture, ori_trust_map = self.render.fast_bake_texture(
                project_textures, project_weighted_cos_maps)
        else:
            raise f'no method {method}'

        return texture, ori_trust_map > 1E-8


    def texture_inpaint(self, texture, mask):

        texture_np = self.render.uv_inpaint(texture, mask)
        texture = torch.tensor(texture_np / 255).float().to(texture.device)

        return texture


    def recenter_image(self, image, border_ratio=0.2):
        if image.mode == 'RGB':
            return image
        elif image.mode == 'L':
            return image.convert('RGB')

        alpha = np.array(image)[:, :, 3]
        non_zero = np.argwhere(alpha > 0)
        if non_zero.size == 0:
            raise ValueError("Image fully transparent")

        min_row, min_col = non_zero.min(axis=0)
        max_row, max_col = non_zero.max(axis=0)

        cropped = image.crop((min_col, min_row, max_col + 1, max_row + 1))

        w, h = cropped.size
        bw = int(w * border_ratio)
        bh = int(h * border_ratio)

        new_w = w + 2 * bw
        new_h = h + 2 * bh
        sq = max(new_w, new_h)

        new_img = Image.new('RGBA', (sq, sq), (255, 255, 255, 0))
        new_img.paste(cropped, ((sq - new_w) // 2 + bw, (sq - new_h) // 2 + bh))

        return new_img


    @torch.no_grad()
    def __call__(self, mesh, image):

        if isinstance(image, str):
            image_prompt = Image.open(image)
        else:
            image_prompt = image

        image_prompt = self.recenter_image(image_prompt)

        # delight
        image_prompt = self.models['delight_model'](image_prompt)

        mesh = mesh_uv_wrap(mesh)
        self.render.load_mesh(mesh)

        elevs = self.config.candidate_camera_elevs
        azims = self.config.candidate_camera_azims
        weights = self.config.candidate_view_weights

        normal_maps = self.render_normal_multiview(elevs, azims)
        position_maps = self.render_position_multiview(elevs, azims)

        camera_info = [
            (((azim // 30) + 9) % 12) //
            {-20: 1, 0: 1, 20: 1, -90: 3, 90: 3}[elev] +
            {-20: 0, 0: 12, 20: 24, -90: 36, 90: 40}[elev]
            for azim, elev in zip(azims, elevs)
        ]

        multiviews = self.models['multiview_model'](
            image_prompt, normal_maps + position_maps, camera_info
        )

        for i in range(len(multiviews)):
            multiviews[i] = multiviews[i].resize(
                (self.config.render_size, self.config.render_size)
            )

        texture, mask = self.bake_from_multiview(
            multiviews, elevs, azims, weights, method=self.config.merge_method
        )

        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)

        texture = self.texture_inpaint(texture, mask_np)
        self.render.set_texture(texture)

        return self.render.save_mesh()