xinyun/ComfyUI-Hunyuan3DWrapper
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-Hunyuan3DWrapper.git synced 2026-01-23 16:34:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

PBR texture saving

Browse Source
This commit is contained in:
kijai 2025-01-25 14:55:45 +02:00
parent 57c84169cd
commit 699bb03e2c
2 changed files with 102 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

125
nodes.py
View File

@ -1,5 +1,6 @@
import os
import torch
import torchvision.transforms as transforms
from PIL import Image
from pathlib import Path
import numpy as np
@ -14,7 +15,7 @@ from comfy.utils import load_torch_file, ProgressBar
script_directory = os.path.dirname(os.path.abspath(__file__))
from .utils import log, print_memory, pil_list_to_torch_batch
from .utils import log, print_memory
class ComfyProgressCallback:
def __init__(self, total_steps):
@ -138,28 +139,6 @@ class DownloadAndLoadHy3DDelightModel:
delight_pipe.enable_model_cpu_offload()
return (delight_pipe,)
class Hy3DLoadMesh:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"glb_path": ("STRING", {"default": "", "tooltip": "The glb path with mesh to load."}),
}
}
RETURN_TYPES = ("HY3DMESH",)
RETURN_NAMES = ("mesh",)
OUTPUT_TOOLTIPS = ("The glb model with mesh to texturize.",)
FUNCTION = "load"
CATEGORY = "Hunyuan3DWrapper"
DESCRIPTION = "Loads a glb model from the given path."
def load(self, glb_path):
mesh = trimesh.load(glb_path, force="mesh")
return (mesh,)
class Hy3DDelightImage:
@classmethod
@ -249,6 +228,7 @@ class DownloadAndLoadHy3DPaintModel:
pipeline.enable_model_cpu_offload()
return (pipeline,)
#region Texture
class Hy3DCameraConfig:
@classmethod
def INPUT_TYPES(s):
@ -661,7 +641,31 @@ class Hy3DApplyTexture:
textured_mesh = self.render.save_mesh()
return (textured_mesh,)
#region Mesh
class Hy3DLoadMesh:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"glb_path": ("STRING", {"default": "", "tooltip": "The glb path with mesh to load."}),
}
}
RETURN_TYPES = ("HY3DMESH",)
RETURN_NAMES = ("mesh",)
OUTPUT_TOOLTIPS = ("The glb model with mesh to texturize.",)
FUNCTION = "load"
CATEGORY = "Hunyuan3DWrapper"
DESCRIPTION = "Loads a glb model from the given path."
def load(self, glb_path):
mesh = trimesh.load(glb_path, force="mesh")
return (mesh,)
class Hy3DGenerateMesh:
@classmethod
def INPUT_TYPES(s):
@ -754,6 +758,77 @@ class Hy3DPostprocessMesh:
return (mesh, )
class Hy3DGetMeshPBRTextures:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mesh": ("HY3DMESH",),
"texture" : (["base_color", "emissive", "metallic_roughness", "normal", "occlusion"], ),
},
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image",)
FUNCTION = "get_textures"
CATEGORY = "Hunyuan3DWrapper"
def get_textures(self, mesh, texture):
TEXTURE_MAPPING = {
'base_color': ('baseColorTexture', "Base color"),
'emissive': ('emissiveTexture', "Emissive"),
'metallic_roughness': ('metallicRoughnessTexture', "Metallic roughness"),
'normal': ('normalTexture', "Normal"),
'occlusion': ('occlusionTexture', "Occlusion"),
}
texture_attr, texture_name = TEXTURE_MAPPING[texture]
texture_data = getattr(mesh.visual.material, texture_attr)
if texture_data is None:
raise ValueError(f"{texture_name} texture not found")
to_tensor = transforms.ToTensor()
return (to_tensor(texture_data).unsqueeze(0).permute(0, 2, 3, 1).cpu().float(),)
class Hy3DSetMeshPBRTextures:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mesh": ("HY3DMESH",),
"image": ("IMAGE", ),
"texture" : (["base_color", "emissive", "metallic_roughness", "normal", "occlusion"], ),
},
}
RETURN_TYPES = ("HY3DMESH", )
RETURN_NAMES = ("mesh",)
FUNCTION = "set_textures"
CATEGORY = "Hunyuan3DWrapper"
def set_textures(self, mesh, image, texture):
TEXTURE_MAPPING = {
'base_color': ('baseColorTexture', "Base color"),
'emissive': ('emissiveTexture', "Emissive"),
'metallic_roughness': ('metallicRoughnessTexture', "Metallic roughness"),
'normal': ('normalTexture', "Normal"),
'occlusion': ('occlusionTexture', "Occlusion"),
}
new_mesh = mesh.copy()
texture_attr, texture_name = TEXTURE_MAPPING[texture]
image_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
if image_np.shape[2] == 4: # RGBA
pil_image = Image.fromarray(image_np, 'RGBA')
else: # RGB
pil_image = Image.fromarray(image_np, 'RGB')
setattr(new_mesh.visual.material, texture_attr, pil_image)
return (new_mesh,)
class Hy3DExportMesh:
@classmethod
def INPUT_TYPES(s):
@ -798,6 +873,8 @@ NODE_CLASS_MAPPINGS = {
"Hy3DApplyTexture": Hy3DApplyTexture,
"CV2InpaintTexture": CV2InpaintTexture,
"Hy3DRenderMultiViewDepth": Hy3DRenderMultiViewDepth,
"Hy3DGetMeshPBRTextures": Hy3DGetMeshPBRTextures,
"Hy3DSetMeshPBRTextures": Hy3DSetMeshPBRTextures
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Hy3DModelLoader": "Hy3DModelLoader",
@ -818,4 +895,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"Hy3DApplyTexture": "Hy3D Apply Texture",
"CV2InpaintTexture": "CV2 Inpaint Texture",
"Hy3DRenderMultiViewDepth": "Hy3D Render MultiView Depth",
"Hy3DGetMeshPBRTextures": "Hy3D Get Mesh PBR Textures",
"Hy3DSetMeshPBRTextures": "Hy3D Set Mesh PBR Textures"
}

11
utils.py
View File

@ -23,14 +23,3 @@ def print_memory(device):
log.info(f"Max reserved memory: {max_reserved=:.3f} GB")
#memory_summary = torch.cuda.memory_summary(device=device, abbreviated=False)
#log.info(f"Memory Summary:\n{memory_summary}")
def pil_list_to_torch_batch(normal_maps):
# Convert PIL images to numpy arrays and stack
arrays = [np.array(img) for img in normal_maps]
batch = np.stack(arrays, axis=0)
# Convert to torch tensor, ensure float32
tensor = torch.from_numpy(batch).float() / 255.0
# Tensor is now in B,H,W,C format
return tensor