From 24d16fc77a4803253863df7c05ff642225276148 Mon Sep 17 00:00:00 2001 From: huiwenshi Date: Sat, 14 Jun 2025 06:39:05 +0000 Subject: [PATCH] Upload hy3dpaint/textureGenPipeline.py with huggingface_hub --- hy3dpaint/textureGenPipeline.py | 192 ++++++++++++++++++++++++++++++++ 1 file changed, 192 insertions(+) create mode 100644 hy3dpaint/textureGenPipeline.py diff --git a/hy3dpaint/textureGenPipeline.py b/hy3dpaint/textureGenPipeline.py new file mode 100644 index 0000000..9396eea --- /dev/null +++ b/hy3dpaint/textureGenPipeline.py @@ -0,0 +1,192 @@ +# 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 os +import torch +import copy +import trimesh +import numpy as np +from PIL import Image +from typing import List +from DifferentiableRenderer.MeshRender import MeshRender +from utils.simplify_mesh_utils import remesh_mesh +from utils.multiview_utils import multiviewDiffusionNet +from utils.pipeline_utils import ViewProcessor +from utils.image_super_utils import imageSuperNet +from utils.uvwrap_utils import mesh_uv_wrap +from DifferentiableRenderer.mesh_utils import convert_obj_to_glb +import warnings + +warnings.filterwarnings("ignore") +from diffusers.utils import logging as diffusers_logging + +diffusers_logging.set_verbosity(50) + + +class Hunyuan3DPaintConfig: + def __init__(self, max_num_view, resolution): + self.device = "cuda" + + self.multiview_cfg_path = "cfgs/hunyuan-paint-pbr.yaml" + self.custom_pipeline = "hunyuanpaintpbr" + self.multiview_pretrained_path = "tencent/Hunyuan3D-2.1" + self.dino_ckpt_path = "facebook/dinov2-giant" + self.realesrgan_ckpt_path = "ckpt/RealESRGAN_x4plus.pth" + + self.raster_mode = "cr" + self.bake_mode = "back_sample" + self.render_size = 1024 * 2 + self.texture_size = 1024 * 4 + self.max_selected_view_num = max_num_view + self.resolution = resolution + self.bake_exp = 4 + self.merge_method = "fast" + + # view selection + 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] + + for azim in range(0, 360, 30): + self.candidate_camera_azims.append(azim) + self.candidate_camera_elevs.append(20) + self.candidate_view_weights.append(0.01) + + self.candidate_camera_azims.append(azim) + self.candidate_camera_elevs.append(-20) + self.candidate_view_weights.append(0.01) + + +class Hunyuan3DPaintPipeline: + + def __init__(self, config=None) -> None: + self.config = config if config is not None else Hunyuan3DPaintConfig() + self.models = {} + self.stats_logs = {} + self.render = MeshRender( + default_resolution=self.config.render_size, + texture_size=self.config.texture_size, + bake_mode=self.config.bake_mode, + raster_mode=self.config.raster_mode, + ) + self.view_processor = ViewProcessor(self.config, self.render) + self.load_models() + + def load_models(self): + torch.cuda.empty_cache() + self.models["super_model"] = imageSuperNet(self.config) + self.models["multiview_model"] = multiviewDiffusionNet(self.config) + print("Models Loaded.") + + @torch.no_grad() + def __call__(self, mesh_path=None, image_path=None, output_mesh_path=None, use_remesh=True, save_glb=True): + """Generate texture for 3D mesh using multiview diffusion""" + # Ensure image_prompt is a list + if isinstance(image_path, str): + image_prompt = Image.open(image_path) + elif isinstance(image_path, Image.Image): + image_prompt = image_path + if not isinstance(image_prompt, List): + image_prompt = [image_prompt] + else: + image_prompt = image_path + + # Process mesh + path = os.path.dirname(mesh_path) + if use_remesh: + processed_mesh_path = os.path.join(path, "white_mesh_remesh.obj") + remesh_mesh(mesh_path, processed_mesh_path) + else: + processed_mesh_path = mesh_path + + # Output path + if output_mesh_path is None: + output_mesh_path = os.path.join(path, f"textured_mesh.obj") + + # Load mesh + mesh = trimesh.load(processed_mesh_path) + mesh = mesh_uv_wrap(mesh) + self.render.load_mesh(mesh=mesh) + + ########### View Selection ######### + selected_camera_elevs, selected_camera_azims, selected_view_weights = self.view_processor.bake_view_selection( + self.config.candidate_camera_elevs, + self.config.candidate_camera_azims, + self.config.candidate_view_weights, + self.config.max_selected_view_num, + ) + + normal_maps = self.view_processor.render_normal_multiview( + selected_camera_elevs, selected_camera_azims, use_abs_coor=True + ) + position_maps = self.view_processor.render_position_multiview(selected_camera_elevs, selected_camera_azims) + + ########## Style ########### + image_caption = "high quality" + image_style = [] + for image in image_prompt: + image = image.resize((512, 512)) + if image.mode == "RGBA": + white_bg = Image.new("RGB", image.size, (255, 255, 255)) + white_bg.paste(image, mask=image.getchannel("A")) + image = white_bg + image_style.append(image) + image_style = [image.convert("RGB") for image in image_style] + + ########### Multiview ########## + multiviews_pbr = self.models["multiview_model"]( + image_style, + normal_maps + position_maps, + prompt=image_caption, + custom_view_size=self.config.resolution, + resize_input=True, + ) + ########### Enhance ########## + enhance_images = {} + enhance_images["albedo"] = copy.deepcopy(multiviews_pbr["albedo"]) + enhance_images["mr"] = copy.deepcopy(multiviews_pbr["mr"]) + + for i in range(len(enhance_images["albedo"])): + enhance_images["albedo"][i] = self.models["super_model"](enhance_images["albedo"][i]) + enhance_images["mr"][i] = self.models["super_model"](enhance_images["mr"][i]) + + ########### Bake ########## + for i in range(len(enhance_images)): + enhance_images["albedo"][i] = enhance_images["albedo"][i].resize( + (self.config.render_size, self.config.render_size) + ) + enhance_images["mr"][i] = enhance_images["mr"][i].resize((self.config.render_size, self.config.render_size)) + texture, mask = self.view_processor.bake_from_multiview( + enhance_images["albedo"], selected_camera_elevs, selected_camera_azims, selected_view_weights + ) + mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8) + texture_mr, mask_mr = self.view_processor.bake_from_multiview( + enhance_images["mr"], selected_camera_elevs, selected_camera_azims, selected_view_weights + ) + mask_mr_np = (mask_mr.squeeze(-1).cpu().numpy() * 255).astype(np.uint8) + + ########## inpaint ########### + texture = self.view_processor.texture_inpaint(texture, mask_np) + self.render.set_texture(texture, force_set=True) + if "mr" in enhance_images: + texture_mr = self.view_processor.texture_inpaint(texture_mr, mask_mr_np) + self.render.set_texture_mr(texture_mr) + + self.render.save_mesh(output_mesh_path, downsample=True) + + if save_glb: + convert_obj_to_glb(output_mesh_path, output_mesh_path.replace(".obj", ".glb")) + output_glb_path = output_mesh_path.replace(".obj", ".glb") + + return output_mesh_path