2239 lines
95 KiB
Python
2239 lines
95 KiB
Python
from PIL import Image, ImageSequence, ImageOps
|
|
import torch
|
|
import shutil
|
|
import argparse
|
|
import copy
|
|
import os
|
|
import time
|
|
import re
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
import trimesh as Trimesh
|
|
import gc
|
|
import json
|
|
from typing import Union, Optional, Tuple, List, Any, Callable
|
|
from pathlib import Path
|
|
|
|
from spandrel import ModelLoader, ImageModelDescriptor
|
|
|
|
import folder_paths
|
|
import node_helpers
|
|
import hashlib
|
|
|
|
import comfy.model_management as mm
|
|
from comfy.utils import load_torch_file, ProgressBar
|
|
import comfy.utils
|
|
|
|
|
|
class _LazyImport:
|
|
"""Defer a heavy import until the symbol is first actually used.
|
|
|
|
Importing the shape/paint stacks at module scope made ComfyUI's startup scan
|
|
of this pack cost 10.7s (now 0.6s): diffusers, transformers, meshlib, plus
|
|
two separate paths into onnxruntime -- rembg, and textureGenPipeline ->
|
|
multiview_utils -> pytorch_lightning -> torchmetrics.functional.audio.
|
|
None of it is needed until a node actually runs, so we bind a proxy here
|
|
and resolve it on first use.
|
|
|
|
Note this only moves the cost off startup; whichever pack loads first and
|
|
eagerly imports the same shared stack (ComfyUI-Hunyuan3DWrapper does) will
|
|
pay it instead.
|
|
|
|
Only safe for names that are exclusively called or attribute-accessed;
|
|
anything used with isinstance() or subclassed must stay a real import.
|
|
"""
|
|
|
|
__slots__ = ("_module", "_attr", "_resolved")
|
|
|
|
def __init__(self, module, attr):
|
|
self._module = module
|
|
self._attr = attr
|
|
self._resolved = None
|
|
|
|
def _get(self):
|
|
if self._resolved is None:
|
|
from importlib import import_module
|
|
self._resolved = getattr(import_module(self._module, __package__), self._attr)
|
|
return self._resolved
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return self._get()(*args, **kwargs)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._get(), name)
|
|
|
|
|
|
#shape
|
|
Hunyuan3DDiTFlowMatchingPipeline = _LazyImport(".hy3dshape.hy3dshape.pipelines", "Hunyuan3DDiTFlowMatchingPipeline")
|
|
FaceReducer = _LazyImport(".hy3dshape.hy3dshape.postprocessors", "FaceReducer")
|
|
FloaterRemover = _LazyImport(".hy3dshape.hy3dshape.postprocessors", "FloaterRemover")
|
|
DegenerateFaceRemover = _LazyImport(".hy3dshape.hy3dshape.postprocessors", "DegenerateFaceRemover")
|
|
BackgroundRemover = _LazyImport(".hy3dshape.hy3dshape.rembg", "BackgroundRemover")
|
|
ShapeVAE = _LazyImport(".hy3dshape.hy3dshape.models.autoencoders", "ShapeVAE")
|
|
postprocessmesh = _LazyImport(".hy3dshape.hy3dshape.meshlib", "postprocessmesh")
|
|
|
|
#painting
|
|
mesh_uv_wrap = _LazyImport(".hy3dpaint.utils.uvwrap_utils", "mesh_uv_wrap")
|
|
Hunyuan3DPaintPipeline = _LazyImport(".hy3dpaint.textureGenPipeline", "Hunyuan3DPaintPipeline")
|
|
Hunyuan3DPaintConfig = _LazyImport(".hy3dpaint.textureGenPipeline", "Hunyuan3DPaintConfig")
|
|
|
|
script_directory = os.path.dirname(os.path.abspath(__file__))
|
|
comfy_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
|
diffusions_dir = os.path.join(comfy_path, "models", "diffusers")
|
|
|
|
def parse_string_to_int_list(number_string):
|
|
"""
|
|
Parses a string containing comma-separated numbers into a list of integers.
|
|
|
|
Args:
|
|
number_string: A string containing comma-separated numbers (e.g., "20000,10000,5000").
|
|
|
|
Returns:
|
|
A list of integers parsed from the input string.
|
|
Returns an empty list if the input string is empty or None.
|
|
"""
|
|
if not number_string:
|
|
return []
|
|
|
|
try:
|
|
# Split the string by comma and convert each part to an integer
|
|
int_list = [int(num.strip()) for num in number_string.split(',')]
|
|
return int_list
|
|
except ValueError as e:
|
|
print(f"Error converting string to integer: {e}. Please ensure all values are valid numbers.")
|
|
return []
|
|
|
|
def hy3dpaintimages_to_tensor(images):
|
|
tensors = []
|
|
for pil_img in images:
|
|
np_img = np.array(pil_img).astype(np.uint8)
|
|
np_img = np_img / 255.0
|
|
tensor_img = torch.from_numpy(np_img).float()
|
|
tensors.append(tensor_img)
|
|
tensors = torch.stack(tensors)
|
|
return tensors
|
|
|
|
def get_picture_files(folder_path):
|
|
"""
|
|
Retrieves all picture files (based on common extensions) from a given folder.
|
|
|
|
Args:
|
|
folder_path (str): The path to the folder to search.
|
|
|
|
Returns:
|
|
list: A list of full paths to the picture files found.
|
|
"""
|
|
picture_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp')
|
|
picture_files = []
|
|
|
|
if not os.path.isdir(folder_path):
|
|
print(f"Error: Folder '{folder_path}' not found.")
|
|
return []
|
|
|
|
for entry_name in os.listdir(folder_path):
|
|
full_path = os.path.join(folder_path, entry_name)
|
|
|
|
# Check if the entry is actually a file (and not a sub-directory)
|
|
if os.path.isfile(full_path):
|
|
file_name, file_extension = os.path.splitext(entry_name)
|
|
if file_extension.lower().endswith(picture_extensions):
|
|
picture_files.append(full_path)
|
|
return picture_files
|
|
|
|
def get_mesh_files(folder_path, name_filter = None):
|
|
"""
|
|
Retrieves all picture files (based on common extensions) from a given folder.
|
|
|
|
Args:
|
|
folder_path (str): The path to the folder to search.
|
|
|
|
Returns:
|
|
list: A list of full paths to the picture files found.
|
|
"""
|
|
mesh_extensions = ('.obj', '.glb')
|
|
mesh_files = []
|
|
|
|
if not os.path.isdir(folder_path):
|
|
print(f"Error: Folder '{folder_path}' not found.")
|
|
return []
|
|
|
|
for entry_name in os.listdir(folder_path):
|
|
full_path = os.path.join(folder_path, entry_name)
|
|
|
|
# Check if the entry is actually a file (and not a sub-directory)
|
|
if os.path.isfile(full_path):
|
|
file_name, file_extension = os.path.splitext(entry_name)
|
|
if file_extension.lower().endswith(mesh_extensions):
|
|
if name_filter is None or name_filter.lower() in file_name.lower():
|
|
mesh_files.append(full_path)
|
|
return mesh_files
|
|
|
|
def get_filename_without_extension_os_path(full_file_path):
|
|
"""
|
|
Extracts the filename without its extension from a full file path using os.path.
|
|
|
|
Args:
|
|
full_file_path (str): The complete path to the file.
|
|
|
|
Returns:
|
|
str: The filename without its extension.
|
|
"""
|
|
# 1. Get the base name (filename with extension)
|
|
base_name = os.path.basename(full_file_path)
|
|
|
|
# 2. Split the base name into root (filename without ext) and extension
|
|
file_name_without_ext, _ = os.path.splitext(base_name)
|
|
|
|
return file_name_without_ext
|
|
|
|
def _convert_texture_format(tex: Union[np.ndarray, torch.Tensor, Image.Image],
|
|
texture_size: Tuple[int, int], device: str, force_set: bool = False) -> torch.Tensor:
|
|
"""Unified texture format conversion logic."""
|
|
if not force_set:
|
|
if isinstance(tex, np.ndarray):
|
|
tex = Image.fromarray((tex * 255).astype(np.uint8))
|
|
elif isinstance(tex, torch.Tensor):
|
|
tex_np = tex.cpu().numpy()
|
|
|
|
# 2. Handle potential batch dimension (B, C, H, W) or (B, H, W, C)
|
|
if tex_np.ndim == 4:
|
|
if tex_np.shape[0] == 1:
|
|
tex_np = tex_np.squeeze(0)
|
|
else:
|
|
tex_np = tex_np[0]
|
|
|
|
# 3. Handle data type and channel order for PIL
|
|
if tex_np.ndim == 3:
|
|
if tex_np.shape[0] in [1, 3, 4] and tex_np.shape[0] < tex_np.shape[1] and tex_np.shape[0] < tex_np.shape[2]:
|
|
tex_np = np.transpose(tex_np, (1, 2, 0))
|
|
elif tex_np.shape[2] in [1, 3, 4] and tex_np.shape[0] > 4 and tex_np.shape[1] > 4:
|
|
pass
|
|
else:
|
|
raise ValueError(f"Unsupported 3D tensor shape after squeezing batch and moving to CPU. "
|
|
f"Expected (C, H, W) or (H, W, C) but got {tex_np.shape}")
|
|
|
|
if tex_np.shape[2] == 1:
|
|
tex_np = tex_np.squeeze(2) # Remove the channel dimension
|
|
|
|
elif tex_np.ndim == 2:
|
|
pass
|
|
else:
|
|
raise ValueError(f"Unsupported tensor dimension after squeezing batch and moving to CPU: {tex_np.ndim} "
|
|
f"with shape {tex_np.shape}. Expected 2D or 3D image data.")
|
|
|
|
tex_np_uint8 = (tex_np * 255).astype(np.uint8)
|
|
|
|
tex = Image.fromarray(tex_np_uint8)
|
|
|
|
|
|
tex = tex.resize(texture_size).convert("RGB")
|
|
tex = np.array(tex) / 255.0
|
|
return torch.from_numpy(tex).to(device).float()
|
|
else:
|
|
if isinstance(tex, np.ndarray):
|
|
tex = torch.from_numpy(tex)
|
|
return tex.to(device).float()
|
|
|
|
def convert_ndarray_to_pil(texture):
|
|
texture_size = len(texture)
|
|
tex = _convert_texture_format(texture,(texture_size, texture_size),"cuda")
|
|
tex = tex.cpu().numpy()
|
|
processed_texture = (tex * 255).astype(np.uint8)
|
|
pil_texture = Image.fromarray(processed_texture)
|
|
return pil_texture
|
|
|
|
def get_filename_list(folder_name: str):
|
|
files = [f for f in os.listdir(folder_name)]
|
|
return files
|
|
|
|
# Tensor to PIL
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
# PIL to Tensor
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def numpy2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def convert_pil_images_to_tensor(images):
|
|
tensor_array = []
|
|
|
|
for image in images:
|
|
tensor_array.append(pil2tensor(image))
|
|
|
|
return tensor_array
|
|
|
|
def convert_tensor_images_to_pil(images):
|
|
pil_array = []
|
|
|
|
for image in images:
|
|
pil_array.append(tensor2pil(image))
|
|
|
|
return pil_array
|
|
|
|
class MetaData:
|
|
def __init__(self):
|
|
self.camera_config = None
|
|
self.albedos = None
|
|
self.mrs = None
|
|
self.albedos_upscaled = None
|
|
self.mrs_upscaled = None
|
|
self.mesh_file = None
|
|
|
|
|
|
class Hy3D21MeshGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder"}),
|
|
"image": ("IMAGE", {"tooltip": "Image to generate mesh from"}),
|
|
"steps": ("INT", {"default": 50, "min": 1, "max": 100, "step": 1, "tooltip": "Number of diffusion steps"}),
|
|
"guidance_scale": ("FLOAT", {"default": 5.0, "min": 1, "max": 30, "step": 0.1, "tooltip": "Guidance scale"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"attention_mode": (["sdpa", "sageattn"], {"default": "sdpa"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DLATENT",)
|
|
RETURN_NAMES = ("latents",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def loadmodel(self, model, image, steps, guidance_scale, seed, attention_mode):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
seed = seed % (2**32)
|
|
|
|
#from .hy3dshape.hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline
|
|
#from .hy3dshape.hy3dshape.rembg import BackgroundRemover
|
|
#import torchvision.transforms as T
|
|
|
|
model_path = folder_paths.get_full_path("diffusion_models", model)
|
|
|
|
pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_single_file(
|
|
config_path=os.path.join(script_directory, 'configs', 'dit_config_2_1.yaml'),
|
|
ckpt_path=model_path,
|
|
offload_device=offload_device,
|
|
attention_mode=attention_mode)
|
|
|
|
# to_pil = T.ToPILImage()
|
|
# image = to_pil(image[0].permute(2, 0, 1))
|
|
|
|
# if image.mode == 'RGB':
|
|
# rembg = BackgroundRemover()
|
|
# image = rembg(image)
|
|
|
|
image = tensor2pil(image)
|
|
|
|
latents = pipeline(
|
|
image=image,
|
|
num_inference_steps=steps,
|
|
guidance_scale=guidance_scale,
|
|
generator=torch.manual_seed(seed)
|
|
)
|
|
|
|
del pipeline
|
|
#del vae
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return (latents,)
|
|
|
|
# class Hy3D21MultiViewsMeshGenerator:
|
|
# @classmethod
|
|
# def INPUT_TYPES(s):
|
|
# return {
|
|
# "required": {
|
|
# "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder"}),
|
|
# "steps": ("INT", {"default": 50, "min": 1, "max": 100, "step": 1, "tooltip": "Number of diffusion steps"}),
|
|
# "guidance_scale": ("FLOAT", {"default": 5.0, "min": 1, "max": 30, "step": 0.1, "tooltip": "Guidance scale"}),
|
|
# "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
# "attention_mode": (["sdpa", "sageattn"], {"default": "sdpa"}),
|
|
# },
|
|
# "optional":{
|
|
# "front": ("IMAGE", {"tooltip": "front image"}),
|
|
# "left": ("IMAGE", {"tooltip": "left image"}),
|
|
# "back": ("IMAGE", {"tooltip": "back image"}),
|
|
# "right": ("IMAGE", {"tooltip": "right image"}),
|
|
# },
|
|
# }
|
|
|
|
# RETURN_TYPES = ("HY3DLATENT",)
|
|
# RETURN_NAMES = ("latents",)
|
|
# FUNCTION = "loadmodel"
|
|
# CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
# def loadmodel(self, model, steps, guidance_scale, seed, attention_mode, front = None, left = None, back = None, right = None):
|
|
# device = mm.get_torch_device()
|
|
# offload_device=mm.unet_offload_device()
|
|
|
|
# seed = seed % (2**32)
|
|
|
|
# #from .hy3dshape.hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline
|
|
# #from .hy3dshape.hy3dshape.rembg import BackgroundRemover
|
|
# #import torchvision.transforms as T
|
|
|
|
# model_path = folder_paths.get_full_path("diffusion_models", model)
|
|
|
|
# pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_single_file(
|
|
# config_path=os.path.join(script_directory, 'configs', 'dit_config_2_1_mv.yaml'),
|
|
# ckpt_path=model_path,
|
|
# offload_device=offload_device,
|
|
# attention_mode=attention_mode)
|
|
|
|
# # to_pil = T.ToPILImage()
|
|
# # image = to_pil(image[0].permute(2, 0, 1))
|
|
|
|
# # if image.mode == 'RGB':
|
|
# # rembg = BackgroundRemover()
|
|
# # image = rembg(image)
|
|
|
|
# if front is not None:
|
|
# front = tensor2pil(front)
|
|
# if left is not None:
|
|
# left = tensor2pil(left)
|
|
# if right is not None:
|
|
# right = tensor2pil(right)
|
|
# if back is not None:
|
|
# back = tensor2pil(back)
|
|
|
|
# view_dict = {
|
|
# 'front': front,
|
|
# 'left': left,
|
|
# 'right': right,
|
|
# 'back': back
|
|
# }
|
|
|
|
# latents = pipeline(
|
|
# image=view_dict,
|
|
# num_inference_steps=steps,
|
|
# guidance_scale=guidance_scale,
|
|
# generator=torch.manual_seed(seed)
|
|
# )
|
|
|
|
# del pipeline
|
|
# #del vae
|
|
|
|
# mm.soft_empty_cache()
|
|
# torch.cuda.empty_cache()
|
|
# gc.collect()
|
|
|
|
# return (latents,)
|
|
|
|
class Hy3DMultiViewsGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"camera_config": ("HY3D21CAMERA",),
|
|
"view_size": ("INT", {"default": 512, "min": 512, "max":1024, "step":256}),
|
|
"image": ("IMAGE", {"tooltip": "Image to generate mesh from"}),
|
|
"steps": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1, "tooltip": "Number of steps"}),
|
|
"guidance_scale": ("FLOAT", {"default": 3.0, "min": 1, "max": 10, "step": 0.1, "tooltip": "Guidance scale"}),
|
|
"texture_size": ("INT", {"default":1024,"min":512,"max":4096,"step":512}),
|
|
"unwrap_mesh": ("BOOLEAN", {"default":True}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DPIPELINE", "IMAGE","IMAGE","IMAGE","IMAGE","HY3D21CAMERA","HY3D21METADATA",)
|
|
RETURN_NAMES = ("pipeline", "albedo","mr","positions","normals","camera_config", "metadata")
|
|
FUNCTION = "genmultiviews"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def genmultiviews(self, trimesh, camera_config, view_size, image, steps, guidance_scale, texture_size, unwrap_mesh, seed):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
seed = seed % (2**32)
|
|
|
|
conf = Hunyuan3DPaintConfig(view_size, camera_config["selected_camera_azims"], camera_config["selected_camera_elevs"], camera_config["selected_view_weights"], camera_config["ortho_scale"], texture_size)
|
|
paint_pipeline = Hunyuan3DPaintPipeline(conf)
|
|
|
|
image = tensor2pil(image)
|
|
|
|
temp_folder_path = os.path.join(comfy_path, "temp")
|
|
os.makedirs(temp_folder_path, exist_ok=True)
|
|
temp_output_path = os.path.join(temp_folder_path, "textured_mesh.obj")
|
|
|
|
albedo, mr, normal_maps, position_maps = paint_pipeline(mesh=trimesh, image_path=image, output_mesh_path=temp_output_path, num_steps=steps, guidance_scale=guidance_scale, unwrap=unwrap_mesh, seed=seed)
|
|
|
|
albedo_tensor = hy3dpaintimages_to_tensor(albedo)
|
|
mr_tensor = hy3dpaintimages_to_tensor(mr)
|
|
normals_tensor = hy3dpaintimages_to_tensor(normal_maps)
|
|
positions_tensor = hy3dpaintimages_to_tensor(position_maps)
|
|
|
|
return (paint_pipeline, albedo_tensor, mr_tensor, positions_tensor, normals_tensor, camera_config,)
|
|
|
|
class Hy3DBakeMultiViews:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HY3DPIPELINE", ),
|
|
"camera_config": ("HY3D21CAMERA", ),
|
|
"albedo": ("IMAGE", ),
|
|
"mr": ("IMAGE", )
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DPIPELINE", "NPARRAY", "NPARRAY", "NPARRAY", "NPARRAY", "IMAGE", "IMAGE",)
|
|
RETURN_NAMES = ("pipeline", "albedo", "albedo_mask", "mr", "mr_mask", "albedo_texture", "mr_texture",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, pipeline, camera_config, albedo, mr):
|
|
albedo = convert_tensor_images_to_pil(albedo)
|
|
mr = convert_tensor_images_to_pil(mr)
|
|
|
|
texture, mask, texture_mr, mask_mr = pipeline.bake_from_multiview(albedo,mr,camera_config["selected_camera_elevs"], camera_config["selected_camera_azims"], camera_config["selected_view_weights"])
|
|
|
|
texture_pil = convert_ndarray_to_pil(texture)
|
|
#mask_pil = convert_ndarray_to_pil(mask)
|
|
texture_mr_pil = convert_ndarray_to_pil(texture_mr)
|
|
#mask_mr_pil = convert_ndarray_to_pil(mask_mr)
|
|
|
|
texture_tensor = pil2tensor(texture_pil)
|
|
#mask_tensor = pil2tensor(mask_pil)
|
|
texture_mr_tensor = pil2tensor(texture_mr_pil)
|
|
#mask_mr_tensor = pil2tensor(mask_mr_pil)
|
|
|
|
return (pipeline, texture, mask, texture_mr, mask_mr, texture_tensor, texture_mr_tensor)
|
|
|
|
class Hy3DInPaint:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HY3DPIPELINE", ),
|
|
"albedo": ("NPARRAY", ),
|
|
"albedo_mask": ("NPARRAY", ),
|
|
"mr": ("NPARRAY", ),
|
|
"mr_mask": ("NPARRAY",),
|
|
"output_mesh_name": ("STRING",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","IMAGE","TRIMESH", "STRING",)
|
|
RETURN_NAMES = ("albedo", "mr", "trimesh", "output_glb_path")
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, pipeline, albedo, albedo_mask, mr, mr_mask, output_mesh_name):
|
|
|
|
#albedo = tensor2pil(albedo)
|
|
#albedo_mask = tensor2pil(albedo_mask)
|
|
#mr = tensor2pil(mr)
|
|
#mr_mask = tensor2pil(mr_mask)
|
|
|
|
vertex_inpaint = True
|
|
method = "NS"
|
|
|
|
albedo, mr = pipeline.inpaint(albedo, albedo_mask, mr, mr_mask, vertex_inpaint, method)
|
|
|
|
pipeline.set_texture_albedo(albedo)
|
|
pipeline.set_texture_mr(mr)
|
|
|
|
temp_folder_path = os.path.join(comfy_path, "temp")
|
|
os.makedirs(temp_folder_path, exist_ok=True)
|
|
output_mesh_path = os.path.join(temp_folder_path, f"{output_mesh_name}.obj")
|
|
output_temp_path = pipeline.save_mesh(output_mesh_path)
|
|
|
|
output_glb_path = os.path.join(comfy_path, "output", f"{output_mesh_name}.glb")
|
|
shutil.copyfile(output_temp_path, output_glb_path)
|
|
|
|
trimesh = Trimesh.load(output_glb_path, force="mesh")
|
|
|
|
texture_pil = convert_ndarray_to_pil(albedo)
|
|
texture_mr_pil = convert_ndarray_to_pil(mr)
|
|
texture_tensor = pil2tensor(texture_pil)
|
|
texture_mr_tensor = pil2tensor(texture_mr_pil)
|
|
|
|
output_glb_path = f"{output_mesh_name}.glb"
|
|
|
|
pipeline.clean_memory()
|
|
|
|
del pipeline
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return (texture_tensor, texture_mr_tensor, trimesh, output_glb_path)
|
|
|
|
class Hy3D21CameraConfig:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"camera_azimuths": ("STRING", {"default": "0, 90, 180, 270, 0, 180", "multiline": False}),
|
|
"camera_elevations": ("STRING", {"default": "0, 0, 0, 0, 90, -90", "multiline": False}),
|
|
"view_weights": ("STRING", {"default": "1, 0.1, 0.5, 0.1, 0.05, 0.05", "multiline": False}),
|
|
"ortho_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 2.0, "step": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3D21CAMERA",)
|
|
RETURN_NAMES = ("camera_config",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, camera_azimuths, camera_elevations, view_weights, ortho_scale):
|
|
angles_list = list(map(int, camera_azimuths.replace(" ", "").split(',')))
|
|
elevations_list = list(map(int, camera_elevations.replace(" ", "").split(',')))
|
|
weights_list = list(map(float, view_weights.replace(" ", "").split(',')))
|
|
|
|
camera_config = {
|
|
"selected_camera_azims": angles_list,
|
|
"selected_camera_elevs": elevations_list,
|
|
"selected_view_weights": weights_list,
|
|
"ortho_scale": ortho_scale,
|
|
}
|
|
|
|
return (camera_config,)
|
|
|
|
class Hy3D21VAELoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}),
|
|
},
|
|
"optional":{
|
|
"vae_config": ("HY3D21VAECONFIG",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DVAE",)
|
|
RETURN_NAMES = ("vae",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def loadmodel(self, model_name, vae_config=None):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
model_path = folder_paths.get_full_path("vae", model_name)
|
|
|
|
vae_sd = load_torch_file(model_path)
|
|
|
|
if(vae_config==None):
|
|
vae_config = {
|
|
'num_latents': 4096,
|
|
'embed_dim': 64,
|
|
'num_freqs': 8,
|
|
'include_pi': False,
|
|
'heads': 16,
|
|
'width': 1024,
|
|
'num_encoder_layers': 8,
|
|
'num_decoder_layers': 16,
|
|
'qkv_bias': False,
|
|
'qk_norm': True,
|
|
'scale_factor': 1.0039506158752403,
|
|
'geo_decoder_mlp_expand_ratio': 4,
|
|
'geo_decoder_downsample_ratio': 1,
|
|
'geo_decoder_ln_post': True,
|
|
'point_feats': 4,
|
|
'pc_size': 81920,
|
|
'pc_sharpedge_size': 0
|
|
}
|
|
|
|
vae = ShapeVAE(**vae_config)
|
|
vae.load_state_dict(vae_sd)
|
|
vae.eval().to(torch.float16)
|
|
|
|
return (vae,)
|
|
|
|
class Hy3D21VAEConfig:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"num_latents": ("INT", {"default": 4096, "min":0, "max":256000}),
|
|
"embed_dim": ("INT", {"default": 64, "min":0, "max":256000}),
|
|
"num_freqs": ("INT", {"default": 8, "min":0, "max":256000}),
|
|
"include_pi": ("BOOLEAN", {"default": False}),
|
|
"heads": ("INT", {"default":16, "min":0, "max":256000}),
|
|
"width": ("INT", {"default":1024, "min":0, "max":256000}),
|
|
"num_encoder_layers": ("INT", {"default":8, "min":0, "max":256000}),
|
|
"num_decoder_layers": ("INT", {"default":16, "min":0, "max":256000}),
|
|
"qkv_bias": ("BOOLEAN", {"default":False}),
|
|
"qk_norm": ("BOOLEAN", {"default":True}),
|
|
"scale_factor": ("FLOAT", {"default":1.0039506158752403}),
|
|
"geo_decoder_mlp_expand_ratio": ("INT", {"default":4, "min":0, "max":256000}),
|
|
"geo_decoder_downsample_ratio": ("INT", {"default":1, "min":0, "max":256000}),
|
|
"geo_decoder_ln_post": ("BOOLEAN", {"default":True}),
|
|
"point_feats": ("INT", {"default":4, "min":0, "max":256000}),
|
|
"pc_size": ("INT", {"default":81920, "min":0, "max":256000}),
|
|
"pc_sharpedge_size": ("INT", {"default":0, "min":0, "max":256000}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3D21VAECONFIG",)
|
|
RETURN_NAMES = ("vae_config",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, num_latents, embed_dim, num_freqs, include_pi, heads, width, num_encoder_layers, num_decoder_layers, qkv_bias, qk_norm, scale_factor, geo_decoder_mlp_expand_ratio, geo_decoder_downsample_ratio, geo_decoder_ln_post, point_feats, pc_size, pc_sharpedge_size):
|
|
vae_config = {
|
|
"num_latents": num_latents,
|
|
"embed_dim": embed_dim,
|
|
"num_freqs": num_freqs,
|
|
"include_pi": include_pi,
|
|
"heads":heads,
|
|
"width":width,
|
|
"num_encoder_layers":num_encoder_layers,
|
|
"num_decoder_layers":num_decoder_layers,
|
|
"qkv_bias":qkv_bias,
|
|
"qk_norm":qk_norm,
|
|
"scale_factor":scale_factor,
|
|
"geo_decoder_mlp_expand_ratio":geo_decoder_mlp_expand_ratio,
|
|
"geo_decoder_downsample_ratio":geo_decoder_downsample_ratio,
|
|
"geo_decoder_ln_post":geo_decoder_ln_post,
|
|
"point_feats":point_feats,
|
|
"pc_size":pc_size,
|
|
"pc_sharpedge_size":pc_sharpedge_size
|
|
}
|
|
|
|
return (vae_config,)
|
|
|
|
class Hy3D21VAEDecode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"vae": ("HY3DVAE",),
|
|
"latents": ("HY3DLATENT", ),
|
|
"box_v": ("FLOAT", {"default": 1.01, "min": -10.0, "max": 10.0, "step": 0.001}),
|
|
"octree_resolution": ("INT", {"default": 384, "min": 8, "max": 4096, "step": 8}),
|
|
"num_chunks": ("INT", {"default": 8000, "min": 1, "max": 10000000, "step": 1, "tooltip": "Number of chunks to process at once, higher values use more memory, but make the process faster"}),
|
|
"mc_level": ("FLOAT", {"default": 0, "min": -1.0, "max": 1.0, "step": 0.0001}),
|
|
"mc_algo": (["mc", "dmc"], {"default": "mc"}),
|
|
},
|
|
"optional": {
|
|
"enable_flash_vdm": ("BOOLEAN", {"default": True}),
|
|
"force_offload": ("BOOLEAN", {"default": False, "tooltip": "Offloads the model to the offload device once the process is done."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, vae, latents, box_v, octree_resolution, mc_level, num_chunks, mc_algo, enable_flash_vdm, force_offload):
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
|
|
vae.to(device)
|
|
|
|
vae.enable_flashvdm_decoder(enabled=enable_flash_vdm, mc_algo=mc_algo)
|
|
|
|
latents = vae.decode(latents)
|
|
outputs = vae.latents2mesh(
|
|
latents,
|
|
output_type='trimesh',
|
|
bounds=box_v,
|
|
mc_level=mc_level,
|
|
num_chunks=num_chunks,
|
|
octree_resolution=octree_resolution,
|
|
mc_algo=mc_algo,
|
|
enable_pbar=True
|
|
)[0]
|
|
|
|
if force_offload==True:
|
|
vae.to(offload_device)
|
|
|
|
outputs.mesh_f = outputs.mesh_f[:, ::-1]
|
|
mesh_output = Trimesh.Trimesh(outputs.mesh_v, outputs.mesh_f)
|
|
print(f"Decoded mesh with {mesh_output.vertices.shape[0]} vertices and {mesh_output.faces.shape[0]} faces")
|
|
|
|
#del pipeline
|
|
del vae
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return (mesh_output, )
|
|
|
|
class Hy3D21ResizeImages:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"width": ("INT", {"default":1024, "min":16, "max":8192} ),
|
|
"height": ("INT", {"default":1024, "min":16, "max":8192} ),
|
|
"sampling": (["NEAREST","LANCZOS","BILINEAR","BICUBIC","BOX","HAMMING"], {"default":"BICUBIC"})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, images, width, height, sampling):
|
|
if sampling=='NEAREST':
|
|
resampling = Image.Resampling.NEAREST
|
|
elif sampling=='LANCZOS':
|
|
resampling = Image.Resampling.LANCZOS
|
|
elif sampling=='BILINEAR':
|
|
resampling = Image.Resampling.BILINEAR
|
|
elif sampling=='BICUBIC':
|
|
resampling = Image.Resampling.BICUBIC
|
|
elif sampling=='BOX':
|
|
resampling = Image.Resampling.BOX
|
|
elif sampling=='HAMMING':
|
|
resampling = Image.Resampling.HAMMING
|
|
else:
|
|
raise Exception('Unknown sampling')
|
|
|
|
if isinstance(images, List):
|
|
for i in range(len(images)):
|
|
if isinstance(images[i], torch.Tensor):
|
|
images[i] = tensor2pil(images[i])
|
|
images[i] = images[i].resize((width,height), resampling)
|
|
images[i] = pil2tensor(images[i])
|
|
elif isinstance(images, torch.Tensor):
|
|
pil_images = convert_tensor_images_to_pil(images)
|
|
for index, img in enumerate(pil_images):
|
|
img = img.resize((width,height), resampling)
|
|
pil_images[index] = img
|
|
tensors = hy3dpaintimages_to_tensor(pil_images)
|
|
return (tensors,)
|
|
elif isinstance(images, Image):
|
|
images = images.resize((width,height), resampling)
|
|
images = pil2tensor(images)
|
|
else:
|
|
raise Exception("Unsupported images format")
|
|
|
|
return (images, )
|
|
|
|
class Hy3D21LoadImageWithTransparency:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
|
files = folder_paths.filter_files_content_types(files, ["image"])
|
|
return {"required":
|
|
{"image": (sorted(files), {"image_upload": True})},
|
|
}
|
|
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", )
|
|
RETURN_NAMES = ("image", "mask", "image_with_alpha")
|
|
FUNCTION = "load_image"
|
|
def load_image(self, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
|
|
img = node_helpers.pillow(Image.open, image_path)
|
|
|
|
output_images = []
|
|
output_masks = []
|
|
output_images_ori = []
|
|
w, h = None, None
|
|
|
|
excluded_formats = ['MPO']
|
|
|
|
for i in ImageSequence.Iterator(img):
|
|
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
|
|
|
output_images_ori.append(pil2tensor(i))
|
|
|
|
if i.mode == 'I':
|
|
i = i.point(lambda i: i * (1 / 255))
|
|
image = i.convert("RGB")
|
|
|
|
if len(output_images) == 0:
|
|
w = image.size[0]
|
|
h = image.size[1]
|
|
|
|
if image.size[0] != w or image.size[1] != h:
|
|
continue
|
|
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
if 'A' in i.getbands():
|
|
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
elif i.mode == 'P' and 'transparency' in i.info:
|
|
mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
|
|
output_images.append(image)
|
|
output_masks.append(mask.unsqueeze(0))
|
|
|
|
if len(output_images) > 1 and img.format not in excluded_formats:
|
|
output_image = torch.cat(output_images, dim=0)
|
|
output_mask = torch.cat(output_masks, dim=0)
|
|
output_image_ori = torch.cat(output_images_ori, dim=0)
|
|
else:
|
|
output_image = output_images[0]
|
|
output_mask = output_masks[0]
|
|
output_image_ori = output_images_ori[0]
|
|
|
|
return (output_image, output_mask, output_image_ori)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
m = hashlib.sha256()
|
|
with open(image_path, 'rb') as f:
|
|
m.update(f.read())
|
|
return m.digest().hex()
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(s, image):
|
|
if not folder_paths.exists_annotated_filepath(image):
|
|
return "Invalid image file: {}".format(image)
|
|
|
|
return True
|
|
|
|
class Hy3D21PostprocessMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"remove_floaters": ("BOOLEAN", {"default": True}),
|
|
"remove_degenerate_faces": ("BOOLEAN", {"default": True}),
|
|
"reduce_faces": ("BOOLEAN", {"default": True}),
|
|
"max_facenum": ("INT", {"default": 40000, "min": 1, "max": 10000000, "step": 1}),
|
|
"smooth_normals": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, trimesh, remove_floaters, remove_degenerate_faces, reduce_faces, max_facenum, smooth_normals):
|
|
new_mesh = trimesh.copy()
|
|
if remove_floaters:
|
|
new_mesh = FloaterRemover()(new_mesh)
|
|
print(f"Removed floaters, resulting in {new_mesh.vertices.shape[0]} vertices and {new_mesh.faces.shape[0]} faces")
|
|
if remove_degenerate_faces:
|
|
new_mesh = DegenerateFaceRemover()(new_mesh)
|
|
print(f"Removed degenerate faces, resulting in {new_mesh.vertices.shape[0]} vertices and {new_mesh.faces.shape[0]} faces")
|
|
if reduce_faces:
|
|
new_mesh = FaceReducer()(new_mesh, max_facenum=max_facenum)
|
|
print(f"Reduced faces, resulting in {new_mesh.vertices.shape[0]} vertices and {new_mesh.faces.shape[0]} faces")
|
|
if smooth_normals:
|
|
new_mesh.vertex_normals = Trimesh.smoothing.get_vertices_normals(new_mesh)
|
|
|
|
return (new_mesh, )
|
|
|
|
class Hy3D21ExportMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"filename_prefix": ("STRING", {"default": "3D/Trellis2"}),
|
|
"file_format": (["glb", "obj", "ply", "stl", "3mf", "dae"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING","STRING","FILE_3D",)
|
|
RETURN_NAMES = ("glb_path","relative_path","model_3d",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, trimesh, filename_prefix, file_format):
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory())
|
|
output_glb_path = Path(full_output_folder, f'{filename}_{counter:05}_.{file_format}')
|
|
output_glb_path.parent.mkdir(exist_ok=True)
|
|
|
|
if file_format=='obj':
|
|
materialName = f"{filename}_{counter:05}_.mtl"
|
|
if hasattr(trimesh, 'visual') and hasattr(trimesh.visual, 'material') and trimesh.visual.material is not None:
|
|
trimesh.visual.material.name = f"{filename}_{counter:05}"
|
|
|
|
trimesh.export(output_glb_path, file_type=file_format, mtl_name=materialName)
|
|
else:
|
|
trimesh.export(output_glb_path, file_type=file_format)
|
|
|
|
relative_path = Path(subfolder) / f'{filename}_{counter:05}_.{file_format}'
|
|
|
|
from comfy_api.latest import Types
|
|
file_3d = Types.File3D(str(output_glb_path))
|
|
|
|
saved_name = f'{filename}_{counter:05}_.{file_format}'
|
|
return {
|
|
"ui": {"3d": [{"filename": saved_name, "subfolder": subfolder, "type": "output"}]},
|
|
"result": (str(output_glb_path), str(relative_path), file_3d,),
|
|
}
|
|
|
|
class Hy3D21MeshUVWrap:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH", )
|
|
RETURN_NAMES = ("trimesh", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, trimesh):
|
|
trimesh = mesh_uv_wrap(trimesh)
|
|
|
|
return (trimesh,)
|
|
|
|
class Hy3D21LoadMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"glb_path": ("STRING", {"default": "", "tooltip": "The glb path with mesh to load."}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
OUTPUT_TOOLTIPS = ("The glb model with mesh to texturize.",)
|
|
|
|
FUNCTION = "load"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
DESCRIPTION = "Loads a glb model from the given path."
|
|
|
|
def load(self, glb_path):
|
|
|
|
if not os.path.exists(glb_path):
|
|
glb_path = os.path.join(folder_paths.get_input_directory(), glb_path)
|
|
|
|
trimesh = Trimesh.load(glb_path, force="mesh")
|
|
|
|
return (trimesh,)
|
|
|
|
class Hy3D21IMRemesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"merge_vertices": ("BOOLEAN", {"default": True}),
|
|
"vertex_count": ("INT", {"default": 10000, "min": 100, "max": 10000000, "step": 1}),
|
|
"smooth_iter": ("INT", {"default": 8, "min": 0, "max": 100, "step": 1}),
|
|
"align_to_boundaries": ("BOOLEAN", {"default": True}),
|
|
"triangulate_result": ("BOOLEAN", {"default": True}),
|
|
"max_facenum": ("INT", {"default": 40000, "min": 1, "max": 10000000, "step": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "remesh"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
DESCRIPTION = "Remeshes the mesh using instant-meshes: https://github.com/wjakob/instant-meshes, Note: this will remove all vertex colors and textures."
|
|
|
|
def remesh(self, trimesh, merge_vertices, vertex_count, smooth_iter, align_to_boundaries, triangulate_result, max_facenum):
|
|
try:
|
|
import pynanoinstantmeshes as PyNIM
|
|
except ImportError:
|
|
raise ImportError("pynanoinstantmeshes not found. Please install it using 'pip install pynanoinstantmeshes'")
|
|
new_mesh = trimesh.copy()
|
|
if merge_vertices:
|
|
trimesh.merge_vertices(new_mesh)
|
|
|
|
new_verts, new_faces = PyNIM.remesh(
|
|
np.array(trimesh.vertices, dtype=np.float32),
|
|
np.array(trimesh.faces, dtype=np.uint32),
|
|
vertex_count,
|
|
align_to_boundaries=align_to_boundaries,
|
|
smooth_iter=smooth_iter
|
|
)
|
|
if new_verts.shape[0] - 1 != new_faces.max():
|
|
# Skip test as the meshing failed
|
|
raise ValueError("Instant-meshes failed to remesh the mesh")
|
|
new_verts = new_verts.astype(np.float32)
|
|
if triangulate_result:
|
|
new_faces = Trimesh.geometry.triangulate_quads(new_faces)
|
|
|
|
if len(new_mesh.faces) > max_facenum:
|
|
new_mesh = FaceReducer()(new_mesh, max_facenum=max_facenum)
|
|
|
|
return (new_mesh, )
|
|
|
|
class Hy3D21MeshlibDecimate:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"subdivideParts": ("INT",{"default":16, "min":1,"max":64,"step":1, "tooltip":"Should be the number of CPU/Core"}),
|
|
},
|
|
"optional":{
|
|
"target_face_num": ("INT",{"min":0,"max":10000000} ),
|
|
"target_face_ratio": ("FLOAT", {"min":0.000,"max":0.999}),
|
|
"strategy": (["None","MinimizeError","ShortestEdgeFirst"],{"default":"None"}),
|
|
"maxError": ("FLOAT",{"min":0.0,"max":1.0}),
|
|
"maxEdgeLen": ("FLOAT",),
|
|
"maxBdShift": ("FLOAT",),
|
|
"maxTriangleAspectRatio": ("FLOAT",),
|
|
"criticalTriAspectRatio": ("FLOAT",),
|
|
"tinyEdgeLength": ("FLOAT",),
|
|
"stabilizer": ("FLOAT",),
|
|
"angleWeightedDistToPlane": ("BOOLEAN",),
|
|
"optimizeVertexPos": ("BOOLEAN",),
|
|
"collapseNearNotFlippable": ("BOOLEAN",),
|
|
"touchNearBdEdges": ("BOOLEAN",),
|
|
"maxAngleChange": ("FLOAT",),
|
|
"decimateBetweenParts": ("BOOLEAN",),
|
|
"minFacesInPart": ("INT",)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "decimate"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
DESCRIPTION = "Decimate the mesh using meshlib: https://meshlib.io/"
|
|
|
|
def decimate(self, trimesh, subdivideParts, target_face_num=0,target_face_ratio=0.0,strategy="None",maxError=0.0,maxEdgeLen=0.0,maxBdShift=0.0,maxTriangleAspectRatio=0.0,criticalTriAspectRatio=0.0,tinyEdgeLength=0.0,stabilizer=0.0,angleWeightedDistToPlane=False,optimizeVertexPos=False,collapseNearNotFlippable=False,touchNearBdEdges=False,maxAngleChange=0.0,decimateBetweenParts=False,minFacesInPart=0):
|
|
try:
|
|
import meshlib.mrmeshpy as mrmeshpy
|
|
except ImportError:
|
|
raise ImportError("meshlib not found. Please install it using 'pip install meshlib'")
|
|
|
|
if target_face_num == 0 and target_face_ratio == 0.0:
|
|
raise ValueError('target_face_num or target_face_ratio must be set')
|
|
|
|
current_faces_num = trimesh.faces.shape[0]
|
|
print(f'Current Faces Number: {current_faces_num}')
|
|
|
|
settings = mrmeshpy.DecimateSettings()
|
|
if target_face_num > 0:
|
|
faces_to_delete = current_faces_num - target_face_num
|
|
if faces_to_delete<=0:
|
|
return trimesh
|
|
|
|
settings.maxDeletedFaces = faces_to_delete
|
|
elif target_face_ratio > 0.0:
|
|
target_faces = int(current_faces_num * target_face_ratio)
|
|
faces_to_delete = current_faces_num - target_faces
|
|
settings.maxDeletedFaces = faces_to_delete
|
|
else:
|
|
raise ValueError('target_face_num or target_face_ratio must be set')
|
|
|
|
if strategy == "MinimizeError":
|
|
settings.strategy = mrmeshpy.DecimateStrategy.MinimizeError
|
|
elif strategy == "ShortestEdgeFirst":
|
|
settings.strategy = mrmeshpy.DecimateStrategy.ShortestEdgeFirst
|
|
|
|
if maxError > 0.0:
|
|
settings.maxError = maxError
|
|
if maxEdgeLen > 0.0:
|
|
settings.maxEdgeLen = maxEdgeLen
|
|
if maxBdShift > 0.0:
|
|
settings.maxBdShift = maxBdShift
|
|
if maxTriangleAspectRatio > 0.0:
|
|
settings.maxTriangleAspectRatio = maxTriangleAspectRatio
|
|
if criticalTriAspectRatio > 0.0:
|
|
settings.criticalTriAspectRatio = criticalTriAspectRatio
|
|
if tinyEdgeLength > 0.0:
|
|
settings.tinyEdgeLength = tinyEdgeLength
|
|
if stabilizer > 0.0:
|
|
settings.stabilizer = stabilizer
|
|
if angleWeightedDistToPlane == True:
|
|
settings.angleWeightedDistToPlane = angleWeightedDistToPlane
|
|
if optimizeVertexPos == True:
|
|
settings.optimizeVertexPos = optimizeVertexPos
|
|
if collapseNearNotFlippable == True:
|
|
settings.collapseNearNotFlippable = collapseNearNotFlippable
|
|
if touchNearBdEdges == True:
|
|
settings.touchNearBdEdges = touchNearBdEdges
|
|
if maxAngleChange > 0.0:
|
|
settings.maxAngleChange = maxAngleChange
|
|
if decimateBetweenParts == True:
|
|
settings.decimateBetweenParts = decimateBetweenParts
|
|
if minFacesInPart > 0:
|
|
settings.minFacesInPart = minFacesInPart
|
|
|
|
settings.packMesh = True
|
|
settings.subdivideParts = subdivideParts
|
|
|
|
new_mesh = postprocessmesh(trimesh.vertices, trimesh.faces, settings)
|
|
|
|
return (new_mesh, )
|
|
|
|
class Hy3D21SimpleMeshlibDecimate:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"subdivideParts": ("INT",{"default":16, "min":1,"max":64,"step":1, "tooltip":"Should be the number of CPU/Core"}),
|
|
},
|
|
"optional":{
|
|
"target_face_num": ("INT",{"min":0,"max":10000000} ),
|
|
"target_face_ratio": ("FLOAT", {"min":0.000,"max":0.999}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRIMESH",)
|
|
RETURN_NAMES = ("trimesh",)
|
|
FUNCTION = "decimate"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
DESCRIPTION = "Decimate the mesh using meshlib: https://meshlib.io/"
|
|
|
|
def decimate(self, trimesh, subdivideParts, target_face_num=0,target_face_ratio=0.0):
|
|
try:
|
|
import meshlib.mrmeshpy as mrmeshpy
|
|
except ImportError:
|
|
raise ImportError("meshlib not found. Please install it using 'pip install meshlib'")
|
|
|
|
if target_face_num == 0 and target_face_ratio == 0.0:
|
|
raise ValueError('target_face_num or target_face_ratio must be set')
|
|
|
|
current_faces_num = trimesh.faces.shape[0]
|
|
print(f'Current Faces Number: {current_faces_num}')
|
|
|
|
settings = mrmeshpy.DecimateSettings()
|
|
if target_face_num > 0:
|
|
faces_to_delete = current_faces_num - target_face_num
|
|
settings.maxDeletedFaces = faces_to_delete
|
|
elif target_face_ratio > 0.0:
|
|
target_faces = int(current_faces_num * target_face_ratio)
|
|
faces_to_delete = current_faces_num - target_faces
|
|
settings.maxDeletedFaces = faces_to_delete
|
|
else:
|
|
raise ValueError('target_face_num or target_face_ratio must be set')
|
|
|
|
if faces_to_delete<=0:
|
|
return trimesh
|
|
|
|
settings.packMesh = True
|
|
settings.subdivideParts = subdivideParts
|
|
|
|
new_mesh = postprocessmesh(trimesh.vertices, trimesh.faces, settings)
|
|
|
|
return (new_mesh, )
|
|
|
|
class Hy3D21MeshGenerationBatch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"input_folder": ("STRING",),
|
|
"output_folder": ("STRING",),
|
|
"vae_model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}),
|
|
"dit_model_name": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder"}),
|
|
"steps": ("INT", {"default": 50, "min": 1, "max": 100, "step": 1, "tooltip": "Number of diffusion steps"}),
|
|
"guidance_scale": ("FLOAT", {"default": 5.0, "min": 1, "max": 30, "step": 0.1, "tooltip": "Guidance scale"}),
|
|
"attention_mode": (["sdpa", "sageattn"], {"default": "sdpa"}),
|
|
"box_v": ("FLOAT", {"default": 1.01, "min": -10.0, "max": 10.0, "step": 0.001}),
|
|
"octree_resolution": ("INT", {"default": 384, "min": 8, "max": 4096, "step": 8}),
|
|
"num_chunks": ("INT", {"default": 8000, "min": 1, "max": 10000000, "step": 1, "tooltip": "Number of chunks to process at once, higher values use more memory, but make the process faster"}),
|
|
"mc_level": ("FLOAT", {"default": 0, "min": -1.0, "max": 1.0, "step": 0.0001}),
|
|
"mc_algo": (["mc", "dmc"], {"default": "mc"}),
|
|
"simplify": ("BOOLEAN",{"default": True}),
|
|
"target_face_num": ("INT",{"default": 200000,"min":0,"max":10000000} ),
|
|
"seed": ("INT",),
|
|
"generate_random_seed": ("BOOLEAN",{"default":True}),
|
|
"file_format": (["glb", "obj"],),
|
|
"remove_background": ("BOOLEAN",{"default":False}),
|
|
"skip_generated_mesh": ("BOOLEAN", {"default":True}),
|
|
},
|
|
"optional": {
|
|
"enable_flash_vdm": ("BOOLEAN", {"default": True}),
|
|
"force_offload": ("BOOLEAN", {"default": False, "tooltip": "Offloads the model to the offload device once the process is done."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING","STRING","STRING","STRING",)
|
|
RETURN_NAMES = ("input_folder", "output_folder", "processed_input_images", "processed_output_meshes",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
DESCRIPTION = "Process all pictures from a folder"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, input_folder, output_folder, vae_model_name, dit_model_name, steps, guidance_scale, attention_mode, box_v, octree_resolution, num_chunks, mc_level, mc_algo, simplify, target_face_num, seed, generate_random_seed, file_format, remove_background, skip_generated_mesh, enable_flash_vdm, force_offload):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
files = get_picture_files(input_folder)
|
|
nb_pictures = len(files)
|
|
|
|
processed_input_images = []
|
|
processed_output_meshes = []
|
|
|
|
if nb_pictures>0:
|
|
rembg = BackgroundRemover()
|
|
|
|
dit_model_path = folder_paths.get_full_path("diffusion_models", dit_model_name)
|
|
|
|
pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_single_file(
|
|
config_path=os.path.join(script_directory, 'configs', 'dit_config_2_1.yaml'),
|
|
ckpt_path=dit_model_path,
|
|
offload_device=offload_device,
|
|
attention_mode=attention_mode)
|
|
|
|
vae_model_path = folder_paths.get_full_path("vae", vae_model_name)
|
|
vae_sd = load_torch_file(vae_model_path)
|
|
|
|
vae_config = {
|
|
'num_latents': 4096,
|
|
'embed_dim': 64,
|
|
'num_freqs': 8,
|
|
'include_pi': False,
|
|
'heads': 16,
|
|
'width': 1024,
|
|
'num_encoder_layers': 8,
|
|
'num_decoder_layers': 16,
|
|
'qkv_bias': False,
|
|
'qk_norm': True,
|
|
'scale_factor': 1.0039506158752403,
|
|
'geo_decoder_mlp_expand_ratio': 4,
|
|
'geo_decoder_downsample_ratio': 1,
|
|
'geo_decoder_ln_post': True,
|
|
'point_feats': 4,
|
|
'pc_size': 81920,
|
|
'pc_sharpedge_size': 0
|
|
}
|
|
|
|
vae = ShapeVAE(**vae_config)
|
|
vae.load_state_dict(vae_sd)
|
|
vae.eval().to(torch.float16)
|
|
vae.to(device)
|
|
|
|
vae.enable_flashvdm_decoder(enabled=enable_flash_vdm, mc_algo=mc_algo)
|
|
|
|
pbar = ProgressBar(nb_pictures)
|
|
for file in files:
|
|
output_file_name = get_filename_without_extension_os_path(file)
|
|
output_glb_path = Path(output_folder, f'{output_file_name}.{file_format}')
|
|
|
|
processImage = True
|
|
|
|
if skip_generated_mesh:
|
|
if os.path.exists(output_glb_path):
|
|
processImage = False
|
|
|
|
if processImage == True:
|
|
print(f'Processing {file} ...')
|
|
if generate_random_seed:
|
|
seed = int.from_bytes(os.urandom(4), 'big')
|
|
|
|
image = Image.open(file)
|
|
|
|
if remove_background:
|
|
print('Removing background ...')
|
|
image = rembg(image)
|
|
|
|
latents = pipeline(
|
|
image=image,
|
|
num_inference_steps=steps,
|
|
guidance_scale=guidance_scale,
|
|
generator=torch.manual_seed(seed)
|
|
)
|
|
|
|
latents = vae.decode(latents)
|
|
outputs = vae.latents2mesh(
|
|
latents,
|
|
output_type='trimesh',
|
|
bounds=box_v,
|
|
mc_level=mc_level,
|
|
num_chunks=num_chunks,
|
|
octree_resolution=octree_resolution,
|
|
mc_algo=mc_algo,
|
|
enable_pbar=True
|
|
)[0]
|
|
|
|
if force_offload==True:
|
|
vae.to(offload_device)
|
|
|
|
outputs.mesh_f = outputs.mesh_f[:, ::-1]
|
|
|
|
mesh_output = Trimesh.Trimesh(outputs.mesh_v, outputs.mesh_f)
|
|
mesh_output = FloaterRemover()(mesh_output)
|
|
mesh_output = DegenerateFaceRemover()(mesh_output)
|
|
|
|
if simplify==True and target_face_num>0:
|
|
try:
|
|
import meshlib.mrmeshpy as mrmeshpy
|
|
except ImportError:
|
|
raise ImportError("meshlib not found. Please install it using 'pip install meshlib'")
|
|
|
|
if target_face_num == 0 and target_face_ratio == 0.0:
|
|
raise ValueError('target_face_num or target_face_ratio must be set')
|
|
|
|
current_faces_num = len(mesh_output.faces)
|
|
print(f'Current Faces Number: {current_faces_num}')
|
|
|
|
settings = mrmeshpy.DecimateSettings()
|
|
faces_to_delete = current_faces_num - target_face_num
|
|
settings.maxDeletedFaces = faces_to_delete
|
|
settings.packMesh = True
|
|
|
|
print('Decimating ...')
|
|
mesh_output = postprocessmesh(mesh_output.vertices, mesh_output.faces, settings)
|
|
|
|
output_glb_path.parent.mkdir(exist_ok=True)
|
|
|
|
processed_input_images.append(file)
|
|
processed_output_meshes.append(output_glb_path)
|
|
|
|
mesh_output.export(output_glb_path, file_type=file_format)
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
else:
|
|
print(f'Skipping file {file}')
|
|
|
|
pbar.update(1)
|
|
|
|
del pipeline
|
|
del vae
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return (input_folder, output_folder, processed_input_images, processed_output_meshes, )
|
|
|
|
class Hy3D21GenerateMultiViewsBatch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"output_folder": ("STRING",),
|
|
"camera_config": ("HY3D21CAMERA",),
|
|
"view_size": ("INT", {"default": 512, "min": 512, "max":1024, "step":256}),
|
|
"steps": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1, "tooltip": "Number of steps"}),
|
|
"guidance_scale": ("FLOAT", {"default": 3.0, "min": 1, "max": 10, "step": 0.1, "tooltip": "Guidance scale"}),
|
|
"texture_size": ("INT", {"default":1024,"min":512,"max":4096,"step":512}),
|
|
"unwrap_mesh": ("BOOLEAN", {"default":True}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}),
|
|
"generate_random_seed": ("BOOLEAN",{"default":True}),
|
|
"remove_background": ("BOOLEAN",{"default":False}),
|
|
"skip_generated_mesh": ("BOOLEAN",{"default":True}),
|
|
"upscale_multiviews": (["None","CustomModel"],{"default":"None"}),
|
|
"upscale_model_name": (folder_paths.get_filename_list("upscale_models"), ),
|
|
"export_multiviews": ("BOOLEAN",{"default":True, "tooltip":"Multiviews can be used to apply texture to a low poly mesh"}),
|
|
"export_metadata": ("BOOLEAN",{"default":True,"tooltip":"Exporta json file with camera config and multiviews"}),
|
|
},
|
|
"optional": {
|
|
"input_images_folder": ("STRING",),
|
|
"input_meshes_folder": ("STRING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("processed_meshes",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
DESCRIPTION = "Process all meshes from a folder"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, output_folder, camera_config, view_size, steps, guidance_scale, texture_size, unwrap_mesh, seed, generate_random_seed, remove_background, skip_generated_mesh, upscale_multiviews, upscale_model_name, export_multiviews, export_metadata, input_images_folder = None, input_meshes_folder = None):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
rembg = BackgroundRemover()
|
|
processed_meshes = []
|
|
|
|
vertex_inpaint = True
|
|
method = "NS"
|
|
|
|
if input_images_folder != None and input_meshes_folder != None:
|
|
files = get_picture_files(input_images_folder)
|
|
nb_pictures = len(files)
|
|
|
|
if nb_pictures>0:
|
|
conf = Hunyuan3DPaintConfig(view_size, camera_config["selected_camera_azims"], camera_config["selected_camera_elevs"], camera_config["selected_view_weights"], camera_config["ortho_scale"], texture_size)
|
|
|
|
temp_folder_path = os.path.join(comfy_path, "temp")
|
|
os.makedirs(temp_folder_path, exist_ok=True)
|
|
temp_output_path = os.path.join(temp_folder_path, "textured_mesh.obj")
|
|
|
|
pbar = ProgressBar(nb_pictures)
|
|
for file in files:
|
|
image_name = get_filename_without_extension_os_path(file)
|
|
input_meshes = get_mesh_files(input_meshes_folder, image_name)
|
|
if len(input_meshes)>0:
|
|
if len(input_meshes)>1:
|
|
print(f'Warning: Multiple meshes found for input_image {image_name} -> Taking the first one')
|
|
|
|
output_file_name = get_filename_without_extension_os_path(file)
|
|
output_mesh_folder = os.path.join(output_folder, output_file_name)
|
|
output_glb_path = Path(output_mesh_folder, f'{output_file_name}.glb')
|
|
|
|
processMesh = True
|
|
|
|
if skip_generated_mesh and os.path.exists(output_glb_path):
|
|
processMesh = False
|
|
|
|
if processMesh:
|
|
os.makedirs(output_mesh_folder, exist_ok=True)
|
|
|
|
print(f'Processing {file} with {input_meshes[0]} ...')
|
|
metaData = MetaData()
|
|
metaData.camera_config = camera_config
|
|
image = Image.open(file)
|
|
if remove_background:
|
|
print('Removing background ...')
|
|
image = rembg(image)
|
|
|
|
if generate_random_seed:
|
|
seed = int.from_bytes(os.urandom(4), 'big')
|
|
|
|
trimesh = Trimesh.load(input_meshes[0])
|
|
|
|
paint_pipeline = Hunyuan3DPaintPipeline(conf)
|
|
albedo, mr, normal_maps, position_maps = paint_pipeline(mesh=trimesh, image_path=image, output_mesh_path=temp_output_path, num_steps=steps, guidance_scale=guidance_scale, unwrap=unwrap_mesh, seed=seed)
|
|
|
|
if export_multiviews:
|
|
metaData.albedos = []
|
|
metaData.mrs = []
|
|
|
|
for index, img in enumerate(albedo):
|
|
image_output_path = os.path.join(output_mesh_folder, f'Albedo_{index}.png')
|
|
img.save(image_output_path)
|
|
metaData.albedos.append(f'Albedo_{index}.png')
|
|
|
|
for index, img in enumerate(mr):
|
|
image_output_path = os.path.join(output_mesh_folder, f'MR_{index}.png')
|
|
img.save(image_output_path)
|
|
metaData.mrs.append(f'MR_{index}.png')
|
|
|
|
|
|
if upscale_multiviews == "CustomModel":
|
|
model_path = folder_paths.get_full_path_or_raise("upscale_models", upscale_model_name)
|
|
sd = comfy.utils.load_torch_file(model_path, safe_load=True)
|
|
if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
|
|
sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""})
|
|
upscale_model = ModelLoader().load_from_state_dict(sd).eval()
|
|
|
|
if not isinstance(upscale_model, ImageModelDescriptor):
|
|
print("Cannot Upscale: Upscale model must be a single-image model.")
|
|
del upscale_model
|
|
upscale_model = None
|
|
else:
|
|
upscale_model.to(device)
|
|
|
|
if upscale_model != None:
|
|
print('Upscaling Albedo ...')
|
|
albedo_tensors = hy3dpaintimages_to_tensor(albedo)
|
|
in_img = albedo_tensors.movedim(-1,-3).to(device)
|
|
|
|
tile = 512
|
|
overlap = 32
|
|
|
|
oom = True
|
|
while oom:
|
|
try:
|
|
steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
|
|
oom = False
|
|
except mm.OOM_EXCEPTION as e:
|
|
tile //= 2
|
|
if tile < 128:
|
|
raise e
|
|
|
|
#upscale_model.to("cpu")
|
|
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
|
|
|
|
albedo = convert_tensor_images_to_pil(s)
|
|
|
|
if export_multiviews:
|
|
metaData.albedos_upscaled = []
|
|
for index, img in enumerate(albedo):
|
|
image_output_path = os.path.join(output_mesh_folder, f'Albedo_Upscaled_{index}.png')
|
|
img.save(image_output_path)
|
|
metaData.albedos_upscaled.append(f'Albedo_Upscaled_{index}.png')
|
|
|
|
print('Upscaling MR ...')
|
|
mr_tensors = hy3dpaintimages_to_tensor(mr)
|
|
in_img = mr_tensors.movedim(-1,-3).to(device)
|
|
|
|
tile = 512
|
|
overlap = 32
|
|
|
|
oom = True
|
|
while oom:
|
|
try:
|
|
steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
|
|
oom = False
|
|
except mm.OOM_EXCEPTION as e:
|
|
tile //= 2
|
|
if tile < 128:
|
|
raise e
|
|
|
|
#upscale_model.to("cpu")
|
|
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
|
|
|
|
mr = convert_tensor_images_to_pil(s)
|
|
|
|
if export_multiviews:
|
|
metaData.mrs_upscaled = []
|
|
for index, img in enumerate(mr):
|
|
image_output_path = os.path.join(output_mesh_folder, f'MR_Upscaled_{index}.png')
|
|
img.save(image_output_path)
|
|
metaData.mrs_upscaled.append(f'MR_Upscaled_{index}.png')
|
|
|
|
del upscale_model
|
|
|
|
print('Baking MultiViews ...')
|
|
texture, mask, texture_mr, mask_mr = paint_pipeline.bake_from_multiview(albedo,mr,camera_config["selected_camera_elevs"], camera_config["selected_camera_azims"], camera_config["selected_view_weights"])
|
|
|
|
albedo, mr = paint_pipeline.inpaint(texture, mask, texture_mr, mask_mr, vertex_inpaint, method)
|
|
paint_pipeline.set_texture_albedo(albedo)
|
|
paint_pipeline.set_texture_mr(mr)
|
|
|
|
output_mesh_path = os.path.join(comfy_path, "temp", f"{output_file_name}.obj")
|
|
output_temp_path = paint_pipeline.save_mesh(output_mesh_path)
|
|
shutil.copyfile(output_temp_path, output_glb_path)
|
|
metaData.mesh_file = f'{output_file_name}.glb'
|
|
|
|
if export_metadata:
|
|
output_metadata_path = os.path.join(output_mesh_folder,'meta_data.json')
|
|
with open(output_metadata_path,'w') as fw:
|
|
json.dump(metaData.__dict__, indent="\t", fp=fw)
|
|
|
|
processed_meshes.append(output_glb_path)
|
|
|
|
paint_pipeline.clean_memory()
|
|
del paint_pipeline
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
else:
|
|
print(f'Skipping {file}')
|
|
else:
|
|
print(f'Error: No mesh found for input image {image_name}')
|
|
|
|
pbar.update(1)
|
|
|
|
else:
|
|
print('No image found in input_images_folder')
|
|
else:
|
|
print('Nothing to process')
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return (processed_meshes, )
|
|
|
|
class Hy3D21UseMultiViews:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"camera_config": ("HY3D21CAMERA",),
|
|
"albedo": ("IMAGE",),
|
|
"mr": ("IMAGE",),
|
|
"view_size": ("INT",{"default":512}),
|
|
"texture_size": ("INT",{"default":1024}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DPIPELINE", "IMAGE","IMAGE","HY3D21CAMERA",)
|
|
RETURN_NAMES = ("pipeline", "albedo","mr","camera_config",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, trimesh, camera_config, albedo, mr, view_size, texture_size):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
conf = Hunyuan3DPaintConfig(view_size, camera_config["selected_camera_azims"], camera_config["selected_camera_elevs"], camera_config["selected_view_weights"], camera_config["ortho_scale"], texture_size)
|
|
paint_pipeline = Hunyuan3DPaintPipeline(conf)
|
|
paint_pipeline.load_mesh(trimesh)
|
|
|
|
return (paint_pipeline, albedo, mr, camera_config)
|
|
|
|
class Hy3D21UseMultiViewsFromMetaData:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"metadata_file": ("STRING",),
|
|
"view_size": ("INT",{"default":512}),
|
|
"texture_size": ("INT",{"default":1024}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DPIPELINE", "IMAGE","IMAGE","HY3D21CAMERA",)
|
|
RETURN_NAMES = ("pipeline", "albedo","mr","camera_config",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, trimesh, metadata_file, view_size, texture_size):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
with open(metadata_file, 'r') as fr:
|
|
loaded_data = json.load(fr)
|
|
loaded_metaData = MetaData()
|
|
for key, value in loaded_data.items():
|
|
setattr(loaded_metaData, key, value)
|
|
|
|
conf = Hunyuan3DPaintConfig(view_size, loaded_metaData.camera_config["selected_camera_azims"], loaded_metaData.camera_config["selected_camera_elevs"], loaded_metaData.camera_config["selected_view_weights"], loaded_metaData.camera_config["ortho_scale"], texture_size)
|
|
paint_pipeline = Hunyuan3DPaintPipeline(conf)
|
|
|
|
paint_pipeline.load_mesh(trimesh)
|
|
|
|
dir_name = os.path.dirname(metadata_file)
|
|
|
|
albedos = []
|
|
mrs = []
|
|
|
|
if loaded_metaData.albedos_upscaled != None:
|
|
print('Using upscaled pictures ...')
|
|
for file in loaded_metaData.albedos_upscaled:
|
|
albedo_file = os.path.join(dir_name,file)
|
|
albedo = Image.open(albedo_file)
|
|
albedos.append(albedo)
|
|
|
|
for file in loaded_metaData.mrs_upscaled:
|
|
mr_file = os.path.join(dir_name,file)
|
|
mr = Image.open(mr_file)
|
|
mrs.append(mr)
|
|
else:
|
|
print('Using non-upscaled pictures ...')
|
|
for file in loaded_metaData.albedos:
|
|
albedo_file = os.path.join(dir_name,file)
|
|
albedo = Image.open(albedo_file)
|
|
albedos.append(albedo)
|
|
|
|
for file in loaded_metaData.mrs:
|
|
mr_file = os.path.join(dir_name,file)
|
|
mr = Image.open(mr_file)
|
|
mrs.append(mr)
|
|
|
|
albedos_tensor = convert_pil_images_to_tensor(albedos)
|
|
mrs_tensor = convert_pil_images_to_tensor(mrs)
|
|
|
|
return (paint_pipeline, albedos_tensor, mrs_tensor, loaded_metaData.camera_config)
|
|
|
|
class Hy3D21MultiViewsGeneratorWithMetaData:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"trimesh": ("TRIMESH",),
|
|
"camera_config": ("HY3D21CAMERA",),
|
|
"view_size": ("INT", {"default": 512, "min": 512, "max":1024, "step":256}),
|
|
"image": ("IMAGE", {"tooltip": "Image to generate mesh from"}),
|
|
"steps": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1, "tooltip": "Number of steps"}),
|
|
"guidance_scale": ("FLOAT", {"default": 3.0, "min": 1, "max": 10, "step": 0.1, "tooltip": "Guidance scale"}),
|
|
"texture_size": ("INT", {"default":1024,"min":512,"max":4096,"step":512}),
|
|
"unwrap_mesh": ("BOOLEAN", {"default":True}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"output_name":("STRING",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DPIPELINE", "IMAGE","IMAGE","HY3D21METADATA","IMAGE","IMAGE",)
|
|
RETURN_NAMES = ("pipeline", "albedo","mr","metadata","positions","normals",)
|
|
FUNCTION = "genmultiviews"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def genmultiviews(self, trimesh, camera_config, view_size, image, steps, guidance_scale, texture_size, unwrap_mesh, seed, output_name):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
seed = seed % (2**32)
|
|
|
|
conf = Hunyuan3DPaintConfig(view_size, camera_config["selected_camera_azims"], camera_config["selected_camera_elevs"], camera_config["selected_view_weights"], camera_config["ortho_scale"], texture_size)
|
|
paint_pipeline = Hunyuan3DPaintPipeline(conf)
|
|
|
|
image = tensor2pil(image)
|
|
|
|
temp_folder_path = os.path.join(comfy_path, "temp")
|
|
os.makedirs(temp_folder_path, exist_ok=True)
|
|
temp_output_path = os.path.join(temp_folder_path, "textured_mesh.obj")
|
|
|
|
albedo, mr, normal_maps, position_maps = paint_pipeline(mesh=trimesh, image_path=image, output_mesh_path=temp_output_path, num_steps=steps, guidance_scale=guidance_scale, unwrap=unwrap_mesh, seed=seed)
|
|
|
|
albedo_tensor = hy3dpaintimages_to_tensor(albedo)
|
|
mr_tensor = hy3dpaintimages_to_tensor(mr)
|
|
normals_tensor = hy3dpaintimages_to_tensor(normal_maps)
|
|
positions_tensor = hy3dpaintimages_to_tensor(position_maps)
|
|
|
|
output_dir_path = os.path.join(comfy_path, "output", "3D", output_name)
|
|
os.makedirs(output_dir_path, exist_ok=True)
|
|
|
|
metadata = MetaData()
|
|
metadata.mesh_file = output_name
|
|
metadata.camera_config = camera_config
|
|
metadata.albedos = []
|
|
metadata.mrs = []
|
|
|
|
print('Saving Albedo and MR views ...')
|
|
for index, img in enumerate(albedo_tensor):
|
|
output_file_path = os.path.join(output_dir_path,f'Albedo_{index}.png')
|
|
pil_image = tensor2pil(img)
|
|
pil_image.save(output_file_path)
|
|
metadata.albedos.append(f'Albedo_{index}.png')
|
|
|
|
for index, img in enumerate(mr_tensor):
|
|
output_file_path = os.path.join(output_dir_path,f'MR_{index}.png')
|
|
pil_image = tensor2pil(img)
|
|
pil_image.save(output_file_path)
|
|
metadata.mrs.append(f'MR_{index}.png')
|
|
|
|
return (paint_pipeline, albedo_tensor, mr_tensor, metadata, positions_tensor, normals_tensor,)
|
|
|
|
class Hy3DBakeMultiViewsWithMetaData:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HY3DPIPELINE", ),
|
|
"albedo": ("IMAGE", ),
|
|
"mr": ("IMAGE", ),
|
|
"metadata": ("HY3D21METADATA",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","IMAGE","TRIMESH", "STRING", )
|
|
RETURN_NAMES = ("albedo", "mr", "trimesh", "output_glb_path", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def process(self, pipeline, albedo, mr, metadata):
|
|
vertex_inpaint = True
|
|
method = "NS"
|
|
|
|
albedo = convert_tensor_images_to_pil(albedo)
|
|
mr = convert_tensor_images_to_pil(mr)
|
|
|
|
output_mesh_name = metadata.mesh_file
|
|
output_dir_path = os.path.join(comfy_path, "output", "3D", output_mesh_name)
|
|
|
|
#detect if images have been upscaled
|
|
albedo1 = albedo[0]
|
|
width, height = albedo1.size
|
|
if width>pipeline.config.resolution:
|
|
print('Upscaled images detected. Saving Upscaled images ...')
|
|
metadata.albedos_upscaled = []
|
|
metadata.mrs_upscaled = []
|
|
|
|
for index, img in enumerate(albedo):
|
|
output_file_path = os.path.join(output_dir_path,f'Albedo_Upscaled_{index}.png')
|
|
img.save(output_file_path)
|
|
metadata.albedos_upscaled.append(f'Albedo_Upscaled_{index}.png')
|
|
|
|
for index, img in enumerate(mr):
|
|
output_file_path = os.path.join(output_dir_path,f'MR_Upscaled_{index}.png')
|
|
img.save(output_file_path)
|
|
metadata.mrs_upscaled.append(f'MR_Upscaled_{index}.png')
|
|
|
|
camera_config = metadata.camera_config
|
|
texture, mask, texture_mr, mask_mr = pipeline.bake_from_multiview(albedo,mr,camera_config["selected_camera_elevs"], camera_config["selected_camera_azims"], camera_config["selected_view_weights"])
|
|
|
|
albedo, mr = pipeline.inpaint(texture, mask, texture_mr, mask_mr, vertex_inpaint, method)
|
|
|
|
pipeline.set_texture_albedo(albedo)
|
|
pipeline.set_texture_mr(mr)
|
|
|
|
output_glb_path = os.path.join(output_dir_path,f'{output_mesh_name}.obj')
|
|
|
|
pipeline.save_mesh(output_glb_path)
|
|
|
|
output_glb_path = os.path.join(output_dir_path,f'{output_mesh_name}.glb')
|
|
|
|
trimesh = Trimesh.load(output_glb_path)
|
|
|
|
texture_pil = convert_ndarray_to_pil(albedo)
|
|
texture_mr_pil = convert_ndarray_to_pil(mr)
|
|
texture_tensor = pil2tensor(texture_pil)
|
|
texture_mr_tensor = pil2tensor(texture_mr_pil)
|
|
|
|
pipeline.clean_memory()
|
|
|
|
metadata.mesh_file = f'{output_mesh_name}.glb'
|
|
|
|
output_metadata_path = os.path.join(output_dir_path,'meta_data.json')
|
|
with open(output_metadata_path,'w') as fw:
|
|
json.dump(metadata.__dict__, indent="\t", fp=fw)
|
|
|
|
del pipeline
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return (texture_tensor, texture_mr_tensor, trimesh, output_glb_path)
|
|
|
|
class Hy3DHighPolyToLowPolyBakeMultiViewsWithMetaData:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"metadata_file": ("STRING",),
|
|
"view_size": ("INT",{"default":512}),
|
|
"texture_size": ("INT",{"default":1024}),
|
|
"target_face_nums": ("STRING",{"default":"20000,10000,5000"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", )
|
|
RETURN_NAMES = ("output_lowpoly_path", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, metadata_file, view_size, texture_size, target_face_nums):
|
|
try:
|
|
import meshlib.mrmeshpy as mrmeshpy
|
|
except ImportError:
|
|
raise ImportError("meshlib not found. Please install it using 'pip install meshlib'")
|
|
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
output_lowpoly_path = ""
|
|
|
|
vertex_inpaint = True
|
|
method = "NS"
|
|
|
|
with open(metadata_file, 'r') as fr:
|
|
loaded_data = json.load(fr)
|
|
loaded_metaData = MetaData()
|
|
for key, value in loaded_data.items():
|
|
setattr(loaded_metaData, key, value)
|
|
|
|
list_of_faces = parse_string_to_int_list(target_face_nums)
|
|
if len(list_of_faces)>0:
|
|
input_dir = os.path.dirname(metadata_file)
|
|
mesh_name = loaded_metaData.mesh_file.replace(".glb","").replace(".obj","")
|
|
mesh_file_path = os.path.join(input_dir, loaded_metaData.mesh_file)
|
|
|
|
if os.path.exists(mesh_file_path):
|
|
conf = Hunyuan3DPaintConfig(view_size, loaded_metaData.camera_config["selected_camera_azims"], loaded_metaData.camera_config["selected_camera_elevs"], loaded_metaData.camera_config["selected_view_weights"], loaded_metaData.camera_config["ortho_scale"], texture_size)
|
|
|
|
highpoly_mesh = Trimesh.load(mesh_file_path, force="mesh")
|
|
highpoly_mesh = Trimesh.Trimesh(vertices=highpoly_mesh.vertices, faces=highpoly_mesh.faces) # Remove texture coordinates
|
|
highpoly_faces_num = highpoly_mesh.faces.shape[0]
|
|
|
|
albedos = []
|
|
mrs = []
|
|
|
|
if loaded_metaData.albedos_upscaled != None:
|
|
print('Using upscaled pictures ...')
|
|
for file in loaded_metaData.albedos_upscaled:
|
|
albedo_file = os.path.join(input_dir,file)
|
|
albedo = Image.open(albedo_file)
|
|
albedos.append(albedo)
|
|
|
|
for file in loaded_metaData.mrs_upscaled:
|
|
mr_file = os.path.join(input_dir,file)
|
|
mr = Image.open(mr_file)
|
|
mrs.append(mr)
|
|
else:
|
|
print('Using non-upscaled pictures ...')
|
|
for file in loaded_metaData.albedos:
|
|
albedo_file = os.path.join(input_dir,file)
|
|
albedo = Image.open(albedo_file)
|
|
albedos.append(albedo)
|
|
|
|
for file in loaded_metaData.mrs:
|
|
mr_file = os.path.join(dir_name,file)
|
|
mr = Image.open(mr_file)
|
|
mrs.append(mr)
|
|
|
|
output_lowpoly_path = os.path.join(input_dir, "LowPoly")
|
|
|
|
for target_face_num in list_of_faces:
|
|
print('Processing {target_face_num} faces ...')
|
|
pipeline = Hunyuan3DPaintPipeline(conf)
|
|
output_dir_path = os.path.join(input_dir, "LowPoly", f"{target_face_num}")
|
|
os.makedirs(output_dir_path, exist_ok=True)
|
|
|
|
settings = mrmeshpy.DecimateSettings()
|
|
faces_to_delete = highpoly_faces_num - target_face_num
|
|
settings.maxDeletedFaces = faces_to_delete
|
|
settings.subdivideParts = 16
|
|
settings.packMesh = True
|
|
|
|
print(f'Decimating to {target_face_num} faces ...')
|
|
lowpoly_mesh = postprocessmesh(highpoly_mesh.vertices, highpoly_mesh.faces, settings)
|
|
|
|
print('UV Unwrapping ...')
|
|
lowpoly_mesh = mesh_uv_wrap(lowpoly_mesh)
|
|
|
|
pipeline.load_mesh(lowpoly_mesh)
|
|
|
|
camera_config = loaded_metaData.camera_config
|
|
texture, mask, texture_mr, mask_mr = pipeline.bake_from_multiview(albedos,mrs,camera_config["selected_camera_elevs"], camera_config["selected_camera_azims"], camera_config["selected_view_weights"])
|
|
|
|
albedo, mr = pipeline.inpaint(texture, mask, texture_mr, mask_mr, vertex_inpaint, method)
|
|
|
|
pipeline.set_texture_albedo(albedo)
|
|
pipeline.set_texture_mr(mr)
|
|
|
|
output_glb_path = os.path.join(output_dir_path,f'{mesh_name}_{target_face_num}.obj')
|
|
|
|
pipeline.save_mesh(output_glb_path)
|
|
|
|
pipeline.clean_memory()
|
|
|
|
else:
|
|
print(f'Mesh file does not exist: {mesh_file_path}')
|
|
else:
|
|
print('target_face_nums is empty')
|
|
|
|
return (output_lowpoly_path,)
|
|
|
|
class Hy3D21ModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
},
|
|
"optional":{
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "HY3DVAE",)
|
|
RETURN_NAMES = ("model_path", "vae",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def loadmodel(self):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
model_path = os.path.join(folder_paths.models_dir,"diffusion_models", "hunyuan3D-dit-v2-1-fp16.ckpt")
|
|
if not os.path.exists(model_path):
|
|
print('Downloading vae ...')
|
|
from huggingface_hub import hf_hub_download
|
|
import shutil
|
|
|
|
repo_id = "tencent/Hunyuan3D-2.1"
|
|
target_dir = os.path.join(folder_paths.models_dir, "diffusion_models")
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
|
|
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
|
|
|
|
repo_filename = "hunyuan3d-dit-v2-1/model.fp16.ckpt"
|
|
desired_filename = "hunyuan3D-dit-v2-1-fp16.ckpt"
|
|
|
|
print(f"downloading {repo_filename} ...")
|
|
downloaded_path = hf_hub_download(
|
|
repo_id=repo_id,
|
|
filename=repo_filename,
|
|
local_dir=target_dir,
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
final_path = os.path.join(target_dir, desired_filename)
|
|
shutil.move(downloaded_path, final_path)
|
|
print(f"{repo_filename} -> {final_path}")
|
|
|
|
# clean up the now-empty subfolder created by hf_hub_download
|
|
subfolder = os.path.dirname(downloaded_path)
|
|
try:
|
|
os.rmdir(subfolder)
|
|
except OSError:
|
|
pass # not empty, or already gone
|
|
|
|
vae_path = os.path.join(folder_paths.models_dir,"vae", "Hunyuan3D-vae-v2-1-fp16.ckpt")
|
|
if not os.path.exists(vae_path):
|
|
print('Downloading vae ...')
|
|
from huggingface_hub import hf_hub_download
|
|
import shutil
|
|
|
|
repo_id = "tencent/Hunyuan3D-2.1"
|
|
target_dir = os.path.join(folder_paths.models_dir, "vae")
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
|
|
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
|
|
|
|
repo_filename = "hunyuan3d-vae-v2-1/model.fp16.ckpt"
|
|
desired_filename = "Hunyuan3D-vae-v2-1-fp16.ckpt"
|
|
|
|
print(f"downloading {repo_filename} ...")
|
|
downloaded_path = hf_hub_download(
|
|
repo_id=repo_id,
|
|
filename=repo_filename,
|
|
local_dir=target_dir,
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
final_path = os.path.join(target_dir, desired_filename)
|
|
shutil.move(downloaded_path, final_path)
|
|
print(f"{repo_filename} -> {final_path}")
|
|
|
|
# clean up the now-empty subfolder created by hf_hub_download
|
|
subfolder = os.path.dirname(downloaded_path)
|
|
try:
|
|
os.rmdir(subfolder)
|
|
except OSError:
|
|
pass # not empty, or already gone
|
|
|
|
vae_sd = load_torch_file(vae_path)
|
|
|
|
vae_config = {
|
|
'num_latents': 4096,
|
|
'embed_dim': 64,
|
|
'num_freqs': 8,
|
|
'include_pi': False,
|
|
'heads': 16,
|
|
'width': 1024,
|
|
'num_encoder_layers': 8,
|
|
'num_decoder_layers': 16,
|
|
'qkv_bias': False,
|
|
'qk_norm': True,
|
|
'scale_factor': 1.0039506158752403,
|
|
'geo_decoder_mlp_expand_ratio': 4,
|
|
'geo_decoder_downsample_ratio': 1,
|
|
'geo_decoder_ln_post': True,
|
|
'point_feats': 4,
|
|
'pc_size': 81920,
|
|
'pc_sharpedge_size': 0
|
|
}
|
|
|
|
vae = ShapeVAE(**vae_config)
|
|
vae.load_state_dict(vae_sd)
|
|
vae.eval().to(torch.float16)
|
|
|
|
return (model_path, vae,)
|
|
|
|
class Hy3D21MeshGenerator2:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_path": ("STRING",),
|
|
"image": ("IMAGE", {"tooltip": "Image to generate mesh from"}),
|
|
"steps": ("INT", {"default": 50, "min": 1, "max": 100, "step": 1, "tooltip": "Number of diffusion steps"}),
|
|
"guidance_scale": ("FLOAT", {"default": 5.0, "min": 1, "max": 30, "step": 0.1, "tooltip": "Guidance scale"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"attention_mode": (["sdpa", "sageattn"], {"default": "sdpa"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DLATENT",)
|
|
RETURN_NAMES = ("latents",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Hunyuan3D21Wrapper"
|
|
|
|
def loadmodel(self, model_path, image, steps, guidance_scale, seed, attention_mode):
|
|
device = mm.get_torch_device()
|
|
offload_device=mm.unet_offload_device()
|
|
|
|
seed = seed % (2**32)
|
|
|
|
#from .hy3dshape.hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline
|
|
#from .hy3dshape.hy3dshape.rembg import BackgroundRemover
|
|
#import torchvision.transforms as T
|
|
|
|
pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_single_file(
|
|
config_path=os.path.join(script_directory, 'configs', 'dit_config_2_1.yaml'),
|
|
ckpt_path=model_path,
|
|
offload_device=offload_device,
|
|
attention_mode=attention_mode)
|
|
|
|
# to_pil = T.ToPILImage()
|
|
# image = to_pil(image[0].permute(2, 0, 1))
|
|
|
|
# if image.mode == 'RGB':
|
|
# rembg = BackgroundRemover()
|
|
# image = rembg(image)
|
|
|
|
image = tensor2pil(image)
|
|
|
|
latents = pipeline(
|
|
image=image,
|
|
num_inference_steps=steps,
|
|
guidance_scale=guidance_scale,
|
|
generator=torch.manual_seed(seed)
|
|
)
|
|
|
|
del pipeline
|
|
#del vae
|
|
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return (latents,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Hy3D21MeshGenerator": Hy3D21MeshGenerator,
|
|
"Hy3DMultiViewsGenerator": Hy3DMultiViewsGenerator,
|
|
"Hy3DBakeMultiViews": Hy3DBakeMultiViews,
|
|
"Hy3DInPaint": Hy3DInPaint,
|
|
"Hy3D21CameraConfig": Hy3D21CameraConfig,
|
|
"Hy3D21VAELoader": Hy3D21VAELoader,
|
|
"Hy3D21VAEDecode": Hy3D21VAEDecode,
|
|
"Hy3D21VAEConfig": Hy3D21VAEConfig,
|
|
"Hy3D21ResizeImages": Hy3D21ResizeImages,
|
|
"Hy3D21LoadImageWithTransparency": Hy3D21LoadImageWithTransparency,
|
|
"Hy3D21PostprocessMesh": Hy3D21PostprocessMesh,
|
|
"Hy3D21ExportMesh": Hy3D21ExportMesh,
|
|
"Hy3D21MeshUVWrap": Hy3D21MeshUVWrap,
|
|
"Hy3D21LoadMesh": Hy3D21LoadMesh,
|
|
"Hy3D21IMRemesh": Hy3D21IMRemesh,
|
|
"Hy3D21MeshlibDecimate": Hy3D21MeshlibDecimate,
|
|
"Hy3D21MeshGenerationBatch": Hy3D21MeshGenerationBatch,
|
|
"Hy3D21GenerateMultiViewsBatch": Hy3D21GenerateMultiViewsBatch,
|
|
"Hy3D21UseMultiViews": Hy3D21UseMultiViews,
|
|
"Hy3D21UseMultiViewsFromMetaData": Hy3D21UseMultiViewsFromMetaData,
|
|
"Hy3D21MultiViewsGeneratorWithMetaData": Hy3D21MultiViewsGeneratorWithMetaData,
|
|
"Hy3DBakeMultiViewsWithMetaData": Hy3DBakeMultiViewsWithMetaData,
|
|
"Hy3DHighPolyToLowPolyBakeMultiViewsWithMetaData": Hy3DHighPolyToLowPolyBakeMultiViewsWithMetaData,
|
|
"Hy3D21SimpleMeshlibDecimate": Hy3D21SimpleMeshlibDecimate,
|
|
"Hy3D21ModelLoader": Hy3D21ModelLoader,
|
|
"Hy3D21MeshGenerator2": Hy3D21MeshGenerator2,
|
|
#"Hy3D21MultiViewsMeshGenerator": Hy3D21MultiViewsMeshGenerator,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Hy3D21MeshGenerator": "Hunyuan 3D 2.1 Mesh Generator",
|
|
"Hy3DMultiViewsGenerator": "Hunyuan 3D 2.1 MultiViews Generator",
|
|
"Hy3DBakeMultiViews": "Hunyuan 3D 2.1 Bake MultiViews",
|
|
"Hy3DInPaint": "Hunyuan 3D 2.1 InPaint",
|
|
"Hy3D21CameraConfig": "Hunyuan 3D 2.1 Camera Config",
|
|
"Hy3D21VAELoader": "Hunyuan 3D 2.1 VAE Loader",
|
|
"Hy3D21VAEDecode": "Hunyuan 3D 2.1 VAE Decoder",
|
|
"Hy3D21VAEConfig": "Hunyuan 3D 2.1 VAE Config",
|
|
"Hy3D21ResizeImages": "Hunyuan 3D 2.1 Resize Images",
|
|
"Hy3D21LoadImageWithTransparency": "Hunyuan 3D 2.1 Load Image with Transparency",
|
|
"Hy3D21PostprocessMesh": "Hunyuan 3D 2.1 Post Process Trimesh",
|
|
"Hy3D21ExportMesh": "Hunyuan 3D 2.1 Export Mesh",
|
|
"Hy3D21MeshUVWrap": "Hunyuan 3D 2.1 Mesh UV Wrap",
|
|
"Hy3D21LoadMesh": "Hunyuan 3D 2.1 Load Mesh",
|
|
"Hy3D21IMRemesh": "Hunyuan 3D 2.1 Instant-Meshes Remesh",
|
|
"Hy3D21MeshlibDecimate": "Hunyuan 3D 2.1 Meshlib Decimation",
|
|
"Hy3D21MeshGenerationBatch": "Hunyuan 3D 2.1 Mesh Generator from Folder",
|
|
"Hy3D21GenerateMultiViewsBatch": "Hunyuan 3D 2.1 MultiViews Generator Batch",
|
|
"Hy3D21UseMultiViews": "Hunyuan 3D 2.1 Use MultiViews",
|
|
"Hy3D21UseMultiViewsFromMetaData": "Hunyuan 3D 2.1 Use MultiViews From MetaData",
|
|
"Hy3D21MultiViewsGeneratorWithMetaData": "Hunyuan 3D 2.1 MultiViews Generator With MetaData",
|
|
"Hy3DBakeMultiViewsWithMetaData": "Hunyuan 3D 2.1 Bake MultiViews With MetaData",
|
|
"Hy3DHighPolyToLowPolyBakeMultiViewsWithMetaData": "Hunyuan 3D 2.1 HighPoly to LowPoly Bake MultiViews With MetaData",
|
|
"Hy3D21SimpleMeshlibDecimate": "Hunyuan 3D 2.1 Simple Meshlib Decimation",
|
|
"Hy3D21ModelLoader": "Hunyuan 3D 2.1 - Model Loader",
|
|
"Hy3D21MeshGenerator2": "Hunyuan 3D 2.1 Mesh Generator2",
|
|
#"Hy3D21MultiViewsMeshGenerator": "Hunyuan 3D 2.1 MultiViews Mesh Generator"
|
|
} |