mirror of
https://git.datalinker.icu/kijai/ComfyUI-Hunyuan3DWrapper.git
synced 2025-12-08 20:34:28 +08:00
454 lines
16 KiB
Python
454 lines
16 KiB
Python
import os
|
|
import torch
|
|
import gc
|
|
from pathlib import Path
|
|
import numpy as np
|
|
|
|
from .hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline, FaceReducer, FloaterRemover, DegenerateFaceRemover
|
|
from .hy3dgen.texgen import Hunyuan3DPaintPipeline
|
|
from .hy3dgen.texgen.utils.dehighlight_utils import Light_Shadow_Remover
|
|
|
|
from accelerate import init_empty_weights
|
|
from accelerate.utils import set_module_tensor_to_device
|
|
|
|
import folder_paths
|
|
|
|
import comfy.model_management as mm
|
|
from comfy.utils import load_torch_file
|
|
|
|
script_directory = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
import logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
log = logging.getLogger(__name__)
|
|
|
|
#region Model loading
|
|
class Hy3DModelLoader:
|
|
@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",}),
|
|
|
|
"base_precision": (["fp32", "bf16"], {"default": "bf16"}),
|
|
"load_device": (["main_device", "offload_device"], {"default": "main_device"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DMODEL",)
|
|
RETURN_NAMES = ("model", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def loadmodel(self, model, base_precision, load_device):
|
|
device = mm.get_torch_device()
|
|
|
|
config_path = os.path.join(script_directory, "configs", "dit_config.yaml")
|
|
model_path = folder_paths.get_full_path("diffusion_models", model)
|
|
pipe = Hunyuan3DDiTFlowMatchingPipeline.from_single_file(ckpt_path=model_path, config_path=config_path, use_safetensors=True, device=device)
|
|
|
|
return (pipe,)
|
|
|
|
class DownloadAndLoadHy3DDelightModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (["hunyuan3d-delight-v2-0"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("DELIGHTMODEL",)
|
|
RETURN_NAMES = ("delight_model", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def loadmodel(self, model):
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
download_path = os.path.join(folder_paths.models_dir,"diffusers")
|
|
model_path = os.path.join(download_path, model)
|
|
|
|
if not os.path.exists(model_path):
|
|
log.info(f"Downloading model to: {model_path}")
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(
|
|
repo_id="tencent/Hunyuan3D-2",
|
|
allow_patterns=["*hunyuan3d-delight-v2-0*"],
|
|
local_dir=download_path,
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
|
|
|
|
delight_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
|
|
model_path,
|
|
torch_dtype=torch.float16,
|
|
safety_checker=None,
|
|
)
|
|
delight_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(delight_pipe.scheduler.config)
|
|
delight_pipe = delight_pipe.to(device, torch.float16)
|
|
|
|
return (delight_pipe,)
|
|
|
|
class Hy3DDelightImage:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"delight_pipe": ("DELIGHTMODEL",),
|
|
"image": ("IMAGE", ),
|
|
"steps": ("INT", {"default": 50, "min": 1}),
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 16}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 16}),
|
|
"cfg_image": ("FLOAT", {"default": 1.5, "min": 0.0, "max": 100.0, "step": 0.01}),
|
|
"cfg_text": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def process(self, delight_pipe, image, width, height, cfg_image, cfg_text, steps, seed):
|
|
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
image = image.permute(0, 3, 1, 2).to(device)
|
|
|
|
delight_pipe = delight_pipe.to(device)
|
|
|
|
image = delight_pipe(
|
|
prompt="",
|
|
image=image,
|
|
generator=torch.manual_seed(seed),
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=steps,
|
|
image_guidance_scale=cfg_image,
|
|
guidance_scale=cfg_text,
|
|
output_type="pt"
|
|
).images[0]
|
|
|
|
delight_pipe = delight_pipe.to(offload_device)
|
|
|
|
out_tensor = image.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
|
|
|
return (out_tensor, )
|
|
|
|
class DownloadAndLoadHy3DPaintModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (["hunyuan3d-paint-v2-0"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DPAINTMODEL",)
|
|
RETURN_NAMES = ("multiview_model", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def loadmodel(self, model):
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
download_path = os.path.join(folder_paths.models_dir,"diffusers")
|
|
model_path = os.path.join(download_path, model)
|
|
|
|
if not os.path.exists(model_path):
|
|
log.info(f"Downloading model to: {model_path}")
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(
|
|
repo_id="tencent/Hunyuan3D-2",
|
|
allow_patterns=[f"*{model}*"],
|
|
local_dir=download_path,
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
|
|
custom_pipeline_path = os.path.join(script_directory, 'hy3dgen', 'texgen', 'hunyuanpaint')
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
model_path,
|
|
custom_pipeline=custom_pipeline_path,
|
|
torch_dtype=torch.float16)
|
|
|
|
pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config, timestep_spacing='trailing')
|
|
|
|
return (pipeline,)
|
|
|
|
class Hy3DRenderMultiView:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HY3DPAINTMODEL",),
|
|
"mesh": ("HY3DMESH",),
|
|
"image": ("IMAGE", ),
|
|
"view_size": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 16}),
|
|
"render_size": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 16}),
|
|
"texture_size": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 16}),
|
|
"steps": ("INT", {"default": 30, "min": 1}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MESHRENDER")
|
|
RETURN_NAMES = ("image", "renderer")
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def process(self, pipeline, image, mesh, view_size, render_size, texture_size, seed, steps):
|
|
device = mm.get_torch_device()
|
|
mm.soft_empty_cache()
|
|
torch.manual_seed(seed)
|
|
generator=torch.Generator(device=pipeline.device).manual_seed(seed)
|
|
|
|
from .hy3dgen.texgen.differentiable_renderer.mesh_render import MeshRender
|
|
|
|
self.render = MeshRender(
|
|
default_resolution=render_size,
|
|
texture_size=texture_size)
|
|
|
|
input_image = image.permute(0, 3, 1, 2).unsqueeze(0).to(device)
|
|
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
from .hy3dgen.texgen.utils.uv_warp_utils import mesh_uv_wrap
|
|
|
|
mesh = mesh_uv_wrap(mesh)
|
|
|
|
self.render.load_mesh(mesh)
|
|
|
|
selected_camera_azims = [0, 90, 180, 270, 0, 180]
|
|
selected_camera_elevs = [0, 0, 0, 0, 90, -90]
|
|
selected_view_weights = [1, 0.1, 0.5, 0.1, 0.05, 0.05]
|
|
|
|
normal_maps = self.render_normal_multiview(
|
|
selected_camera_elevs, selected_camera_azims, use_abs_coor=True)
|
|
position_maps = self.render_position_multiview(
|
|
selected_camera_elevs, selected_camera_azims)
|
|
|
|
camera_info = [(((azim // 30) + 9) % 12) // {-20: 1, 0: 1, 20: 1, -90: 3, 90: 3}[
|
|
elev] + {-20: 0, 0: 12, 20: 24, -90: 36, 90: 40}[elev] for azim, elev in
|
|
zip(selected_camera_azims, selected_camera_elevs)]
|
|
|
|
control_images = normal_maps + position_maps
|
|
|
|
for i in range(len(control_images)):
|
|
control_images[i] = control_images[i].resize((view_size, view_size))
|
|
if control_images[i].mode == 'L':
|
|
control_images[i] = control_images[i].point(lambda x: 255 if x > 1 else 0, mode='1')
|
|
|
|
num_view = len(control_images) // 2
|
|
normal_image = [[control_images[i] for i in range(num_view)]]
|
|
position_image = [[control_images[i + num_view] for i in range(num_view)]]
|
|
|
|
pipeline = pipeline.to(device)
|
|
|
|
multiview_images = pipeline(
|
|
input_image,
|
|
width=view_size,
|
|
height=view_size,
|
|
generator=generator,
|
|
num_in_batch = num_view,
|
|
camera_info_gen = [camera_info],
|
|
camera_info_ref = [[0]],
|
|
normal_imgs = normal_image,
|
|
position_imgs = position_image,
|
|
num_inference_steps=steps,
|
|
output_type="pt",
|
|
).images
|
|
|
|
pipeline = pipeline.to(offload_device)
|
|
|
|
out_tensors = multiview_images.permute(0, 2, 3, 1).cpu().float()
|
|
|
|
return (out_tensors, self.render)
|
|
|
|
def render_normal_multiview(self, camera_elevs, camera_azims, use_abs_coor=True):
|
|
normal_maps = []
|
|
for elev, azim in zip(camera_elevs, camera_azims):
|
|
normal_map = self.render.render_normal(
|
|
elev, azim, use_abs_coor=use_abs_coor, return_type='pl')
|
|
normal_maps.append(normal_map)
|
|
|
|
return normal_maps
|
|
|
|
def render_position_multiview(self, camera_elevs, camera_azims):
|
|
position_maps = []
|
|
for elev, azim in zip(camera_elevs, camera_azims):
|
|
position_map = self.render.render_position(
|
|
elev, azim, return_type='pl')
|
|
position_maps.append(position_map)
|
|
|
|
return position_maps
|
|
|
|
class Hy3DBakeFromMultiview:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", ),
|
|
"renderer": ("MESHRENDER",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DMESH",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def process(self, images, renderer):
|
|
device = mm.get_torch_device()
|
|
self.render = renderer
|
|
|
|
multiviews = images.permute(0, 3, 1, 2).to(device)
|
|
|
|
device = mm.get_torch_device()
|
|
|
|
selected_camera_azims = [0, 90, 180, 270, 0, 180]
|
|
selected_camera_elevs = [0, 0, 0, 0, 90, -90]
|
|
selected_view_weights = [1, 0.1, 0.5, 0.1, 0.05, 0.05]
|
|
merge_method = 'fast'
|
|
|
|
texture, mask = self.bake_from_multiview(multiviews,
|
|
selected_camera_elevs, selected_camera_azims, selected_view_weights,
|
|
method=merge_method)
|
|
|
|
mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
texture_np = self.render.uv_inpaint(texture, mask_np)
|
|
texture = torch.tensor(texture_np / 255).float().to(texture.device)
|
|
|
|
self.render.set_texture(texture)
|
|
textured_mesh = self.render.save_mesh()
|
|
|
|
return (textured_mesh,)
|
|
|
|
def bake_from_multiview(self, views, camera_elevs,
|
|
camera_azims, view_weights, method='graphcut'):
|
|
project_textures, project_weighted_cos_maps = [], []
|
|
project_boundary_maps = []
|
|
for view, camera_elev, camera_azim, weight in zip(
|
|
views, camera_elevs, camera_azims, view_weights):
|
|
project_texture, project_cos_map, project_boundary_map = self.render.back_project(
|
|
view, camera_elev, camera_azim)
|
|
project_cos_map = weight * (project_cos_map ** self.config.bake_exp)
|
|
project_textures.append(project_texture)
|
|
project_weighted_cos_maps.append(project_cos_map)
|
|
project_boundary_maps.append(project_boundary_map)
|
|
|
|
if method == 'fast':
|
|
texture, ori_trust_map = self.render.fast_bake_texture(
|
|
project_textures, project_weighted_cos_maps)
|
|
else:
|
|
raise f'no method {method}'
|
|
return texture, ori_trust_map > 1E-8
|
|
|
|
class Hy3DGenerateMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HY3DMODEL",),
|
|
"image": ("IMAGE", ),
|
|
"octree_resolution": ("INT", {"default": 256, "min": 64, "max": 4096, "step": 16}),
|
|
"guidance_scale": ("FLOAT", {"default": 5.5, "min": 0.0, "max": 100.0, "step": 0.01}),
|
|
"steps": ("INT", {"default": 30, "min": 1}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"remove_floaters": ("BOOLEAN", {"default": True}),
|
|
"remove_degenerate_faces": ("BOOLEAN", {"default": True}),
|
|
"reduce_faces": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK", ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HY3DMESH",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def process(self, pipeline, image, steps, guidance_scale, octree_resolution, seed, remove_floaters, remove_degenerate_faces, reduce_faces,
|
|
mask=None):
|
|
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
image = image.permute(0, 3, 1, 2).to(device)
|
|
|
|
pipeline.to(device)
|
|
|
|
mesh = pipeline(
|
|
image=image,
|
|
mask=mask,
|
|
num_inference_steps=steps,
|
|
mc_algo='mc',
|
|
guidance_scale=guidance_scale,
|
|
octree_resolution=octree_resolution,
|
|
generator=torch.manual_seed(seed))[0]
|
|
|
|
if remove_floaters:
|
|
mesh = FloaterRemover()(mesh)
|
|
if remove_degenerate_faces:
|
|
mesh = DegenerateFaceRemover()(mesh)
|
|
if reduce_faces:
|
|
mesh = FaceReducer()(mesh)
|
|
|
|
pipeline.to(offload_device)
|
|
|
|
return (mesh, )
|
|
|
|
class Hy3DExportMesh:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("HY3DMESH",),
|
|
"filename_prefix": ("STRING", {"default": "3D/Hy3D"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("glb_path",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Hunyuan3DWrapper"
|
|
|
|
def process(self, mesh, filename_prefix):
|
|
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}_.glb')
|
|
output_glb_path.parent.mkdir(exist_ok=True)
|
|
mesh.export(output_glb_path)
|
|
|
|
return (str(output_glb_path), )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Hy3DModelLoader": Hy3DModelLoader,
|
|
"Hy3DGenerateMesh": Hy3DGenerateMesh,
|
|
"Hy3DExportMesh": Hy3DExportMesh,
|
|
"DownloadAndLoadHy3DDelightModel": DownloadAndLoadHy3DDelightModel,
|
|
"DownloadAndLoadHy3DPaintModel": DownloadAndLoadHy3DPaintModel,
|
|
"Hy3DDelightImage": Hy3DDelightImage,
|
|
"Hy3DRenderMultiView": Hy3DRenderMultiView,
|
|
"Hy3DBakeFromMultiview": Hy3DBakeFromMultiview
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Hy3DModelLoader": "Hy3DModelLoader",
|
|
"Hy3DGenerateMesh": "Hy3DGenerateMesh",
|
|
"Hy3DExportMesh": "Hy3DExportMesh",
|
|
"DownloadAndLoadHy3DDelightModel": "(Down)Load Hy3D DelightModel",
|
|
"DownloadAndLoadHy3DPaintModel": "(Down)Load Hy3D PaintModel",
|
|
"Hy3DDelightImage": "Hy3DDelightImage",
|
|
"Hy3DRenderMultiView": "Hy3D Render MultiView",
|
|
"Hy3DBakeFromMultiview": "Hy3D Bake From Multiview"
|
|
}
|