495 lines
16 KiB
Python
495 lines
16 KiB
Python
import os
|
|
import shutil
|
|
import warnings
|
|
import sys
|
|
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
from scipy.spatial import KDTree
|
|
|
|
warnings.simplefilter("ignore", category=UserWarning)
|
|
warnings.simplefilter("ignore", category=FutureWarning)
|
|
warnings.simplefilter("ignore", category=DeprecationWarning)
|
|
from folder_paths import (
|
|
get_filename_list,
|
|
get_full_path,
|
|
get_save_image_path,
|
|
get_output_directory,
|
|
)
|
|
from huggingface_hub import snapshot_download
|
|
import torch
|
|
import numpy as np
|
|
from rembg import remove, new_session
|
|
import mcubes
|
|
import cv2
|
|
from einops import rearrange
|
|
from PIL import Image, ImageSequence, ImageOps, ImageColor
|
|
import io
|
|
import base64
|
|
import ipywidgets as widgets
|
|
from IPython.display import display
|
|
import trimesh
|
|
import torch.nn.functional as F
|
|
from typing import Tuple
|
|
import folder_paths
|
|
from os import path
|
|
from .infer import Removebg, Image2Views, Views2Mesh
|
|
|
|
# 添加当前目录到 sys.path
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if current_dir not in sys.path:
|
|
sys.path.append(current_dir)
|
|
sys.path.insert(0, path.dirname(__file__))
|
|
|
|
# 设置访问路径
|
|
base64_data_Array = {""}
|
|
|
|
|
|
@PromptServer.instance.routes.get("/get_obj_Base64")
|
|
async def get_obj(request):
|
|
return web.json_response(base64_data_Array)
|
|
|
|
|
|
def barycentric_interpolate(v0, v1, v2, c0, c1, c2, p):
|
|
v0v1 = v1 - v0
|
|
v0v2 = v2 - v0
|
|
v0p = p - v0
|
|
d00 = np.dot(v0v1, v0v1)
|
|
d01 = np.dot(v0v1, v0v2)
|
|
d11 = np.dot(v0v2, v0v2)
|
|
d20 = np.dot(v0p, v0v1)
|
|
d21 = np.dot(v0p, v0v2)
|
|
denom = d00 * d11 - d01 * d01
|
|
if abs(denom) < 1e-8:
|
|
return (c0 + c1 + c2) / 3
|
|
v = (d11 * d20 - d01 * d21) / denom
|
|
w = (d00 * d21 - d01 * d20) / denom
|
|
u = 1.0 - v - w
|
|
u = np.clip(u, 0, 1)
|
|
v = np.clip(v, 0, 1)
|
|
w = np.clip(w, 0, 1)
|
|
interpolate_color = u * c0 + v * c1 + w * c2
|
|
return np.clip(interpolate_color, 0, 255)
|
|
|
|
|
|
def is_point_in_triangle(p, v0, v1, v2):
|
|
def sign(p1, p2, p3):
|
|
return (p1[0] - p3[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[1] - p3[1])
|
|
|
|
d1 = sign(p, v0, v1)
|
|
d2 = sign(p, v1, v2)
|
|
d3 = sign(p, v2, v0)
|
|
|
|
has_neg = (d1 < 0) or (d2 < 0) or (d3 < 0)
|
|
has_pos = (d1 > 0) or (d2 > 0) or (d3 > 0)
|
|
|
|
return not (has_neg and has_pos)
|
|
|
|
|
|
def pil2tensor(image: Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
|
|
def tensor2pil(t_image: torch.Tensor) -> Image:
|
|
return Image.fromarray(
|
|
np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
)
|
|
|
|
|
|
class RemoveBackground:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"session": (
|
|
"STRING",
|
|
{
|
|
"multiline": False,
|
|
"default": "u2net",
|
|
"lazy": True,
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "make_six_views"
|
|
CATEGORY = "EX"
|
|
|
|
def __init__(self):
|
|
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
self.worker_i2v = None
|
|
|
|
def make_six_views(
|
|
self,
|
|
image,
|
|
session,
|
|
):
|
|
if not isinstance(image, Image.Image):
|
|
if isinstance(image, torch.Tensor):
|
|
print(f"Original tensor shape: {image.shape}")
|
|
if image.dim() == 4 and image.shape[0] == 1:
|
|
image = image.squeeze(0)
|
|
print(f"Squeezed tensor shape: {image.shape}")
|
|
if image.dim() == 3:
|
|
if image.shape[2] in [1, 3, 4]:
|
|
print(f"Image is in [H, W, C] format with shape: {image.shape}")
|
|
pass
|
|
elif image.shape[0] in [1, 3, 4]:
|
|
print(f"Image is in [C, H, W] format with shape: {image.shape}")
|
|
image = image.permute(1, 2, 0)
|
|
print(f"Transposed image shape: {image.shape}")
|
|
else:
|
|
raise ValueError(f"Unsupported image shape: {image.shape}")
|
|
else:
|
|
raise ValueError(f"Unsupported image dimensions: {image.dim()}")
|
|
image = image.cpu().numpy()
|
|
print(
|
|
f"Converted to numpy array, shape: {image.shape}, dtype: {image.dtype}"
|
|
)
|
|
if image.dtype in [np.float32, np.float64]:
|
|
image = np.clip(image * 255, 0, 255).astype(np.uint8)
|
|
print(f"Scaled image to uint8, dtype: {image.dtype}")
|
|
elif image.dtype == np.uint8:
|
|
print(f"Image is already uint8, dtype: {image.dtype}")
|
|
else:
|
|
image = image.astype(np.uint8)
|
|
print(f"Converted image to uint8, dtype: {image.dtype}")
|
|
if image.ndim != 3 or image.shape[2] not in [1, 3, 4]:
|
|
raise ValueError(
|
|
f"Invalid image shape after processing: {image.shape}"
|
|
)
|
|
image = Image.fromarray(image)
|
|
print(f"Converted to PIL Image, size: {image.size}, mode: {image.mode}")
|
|
else:
|
|
raise TypeError(f"Unsupported image type: {type(image)}")
|
|
else:
|
|
print(
|
|
f"Input image is already a PIL Image, size: {image.size}, mode: {image.mode}"
|
|
)
|
|
rgba = remove(image, session=new_session(session))
|
|
rgba_img_tensor = pil2tensor(rgba)
|
|
return (rgba_img_tensor,)
|
|
|
|
|
|
class GenerateSixViews:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"rgba": ("IMAGE",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"step": ("INT", {"default": 50}),
|
|
"use_lite": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"IMAGE",
|
|
)
|
|
RETURN_NAMES = (
|
|
"sixImages",
|
|
"originalImage",
|
|
)
|
|
FUNCTION = "make_six_views"
|
|
CATEGORY = "EX"
|
|
|
|
def __init__(self):
|
|
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
self.worker_i2v = None
|
|
|
|
def make_six_views(self, rgba, seed, step, use_lite):
|
|
rgba = tensor2pil(rgba)
|
|
if rgba.mode != "RGBA":
|
|
raise ValueError("Input image is not in RGBA format.")
|
|
if self.worker_i2v is None:
|
|
self.worker_i2v = Image2Views(use_lite=use_lite, device=self.device)
|
|
res_img, pils = self.worker_i2v(rgba, seed, step)
|
|
res_multi_img_tensor = pil2tensor(res_img[0])
|
|
use_image_tensor = pil2tensor(res_img[1])
|
|
return (
|
|
res_multi_img_tensor,
|
|
use_image_tensor,
|
|
)
|
|
|
|
|
|
class Hunyuan3DNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"sixImages": ("IMAGE",),
|
|
"originalImage": ("IMAGE",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"mesh_size": (
|
|
"INT",
|
|
{
|
|
"default": 500,
|
|
"min": 0,
|
|
"max": 1024,
|
|
"step": 1,
|
|
"display": "number",
|
|
},
|
|
),
|
|
"max_number_of_faces": (
|
|
"INT",
|
|
{
|
|
"default": 90000,
|
|
"min": 1000,
|
|
"max": 1000000,
|
|
"step": 100,
|
|
"display": "number",
|
|
},
|
|
),
|
|
"use_lite": ("BOOLEAN", {"default": False}),
|
|
"is_simplify": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MESH",)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "imgTo3D"
|
|
CATEGORY = "EX"
|
|
|
|
def __init__(self):
|
|
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
self.worker_v23 = None
|
|
|
|
@staticmethod
|
|
def download_model_files(repo_id: str, local_dir: str):
|
|
"""
|
|
从 Huggingface 仓库下载所有文件到指定的本地目录。
|
|
|
|
:param repo_id: Huggingface 仓库的 ID，例如 "tencent/Hunyuan3D-1"
|
|
:param local_dir: 本地目标目录，例如 "weights"
|
|
"""
|
|
print(
|
|
f"模型文件未找到，正在从 Huggingface 仓库 {repo_id} 下载所有文件到 {local_dir} ..."
|
|
)
|
|
try:
|
|
snapshot_download(
|
|
repo_id=repo_id,
|
|
local_dir=local_dir,
|
|
repo_type="model",
|
|
ignore_patterns=["*.gitignore"],
|
|
)
|
|
print("模型文件下载完成。")
|
|
except Exception as e:
|
|
print(f"下载模型文件时出错: {e}")
|
|
raise
|
|
|
|
def imgTo3D(
|
|
self,
|
|
sixImages,
|
|
originalImage,
|
|
seed,
|
|
mesh_size,
|
|
max_number_of_faces,
|
|
use_lite,
|
|
is_simplify,
|
|
):
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
weights_dir = os.path.join(base_dir, "weights")
|
|
svrm_dir = os.path.join(weights_dir, "svrm")
|
|
weights_path = os.path.join(svrm_dir, "svrm.safetensors")
|
|
|
|
# 检查模型权重文件是否存在
|
|
if not os.path.exists(weights_path):
|
|
print(f"未找到模型文件 {weights_path}，开始下载模型文件...")
|
|
repo_id = "tencent/Hunyuan3D-1" # Huggingface 仓库 ID
|
|
self.download_model_files(repo_id=repo_id, local_dir=weights_dir)
|
|
# 检查下载后是否存在模型文件
|
|
if not os.path.exists(weights_path):
|
|
raise FileNotFoundError(
|
|
f"下载后仍未找到模型文件 {weights_path}。请检查仓库 ID 或网络连接。"
|
|
)
|
|
|
|
config_path = os.path.join(base_dir, "svrm/configs/svrm.yaml")
|
|
weights_path = os.path.join(base_dir, "weights/svrm/svrm.safetensors")
|
|
|
|
if self.worker_v23 is None:
|
|
self.worker_v23 = Views2Mesh(
|
|
config_path,
|
|
weights_path,
|
|
use_lite=use_lite,
|
|
device=self.device,
|
|
)
|
|
|
|
# 从视图生成网格
|
|
mesh = self.worker_v23(
|
|
tensor2pil(sixImages),
|
|
tensor2pil(originalImage),
|
|
seed=seed,
|
|
mesh_size=mesh_size,
|
|
)
|
|
mesh.apply_transform(
|
|
np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
|
|
)
|
|
|
|
# 简化模型
|
|
if is_simplify:
|
|
simplified_mesh = mesh.simplify_quadric_decimation(
|
|
face_count=max_number_of_faces
|
|
)
|
|
# 映射颜色：使用KD树找到新顶点的最近原始顶点
|
|
tree = KDTree(mesh.vertices)
|
|
_, closest_indices = tree.query(simplified_mesh.vertices)
|
|
|
|
# 根据最近点的颜色重新赋值
|
|
simplified_colors = mesh.visual.vertex_colors[closest_indices]
|
|
simplified_mesh.visual.vertex_colors = simplified_colors
|
|
|
|
mesh = simplified_mesh
|
|
|
|
# gltf材质贴图
|
|
# texture_image = mesh.visual.to_texture().material.image.copy()
|
|
# texture_image_tensor = pil2tensor(
|
|
# texture_image.resize((1024, 1024), Image.BILINEAR)
|
|
# )
|
|
|
|
return (mesh,)
|
|
|
|
|
|
class TriMeshViewer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("MESH",),
|
|
"saveName": (
|
|
"STRING",
|
|
{
|
|
"multiline": False,
|
|
"default": "mesh",
|
|
},
|
|
),
|
|
"saveToOutputFolder": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"hidden": {"display": "DISPLAY"},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "display"
|
|
CATEGORY = "EX"
|
|
|
|
def display(self, mesh, saveName, saveToOutputFolder):
|
|
saved = list()
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = (
|
|
get_save_image_path(saveName, get_output_directory())
|
|
)
|
|
|
|
filename_with_batch_num = filename.replace("%batch_num%", str(1))
|
|
file = f"{filename_with_batch_num}_{counter:05}_.glb"
|
|
buffer = io.BytesIO()
|
|
mesh.export(buffer, "glb")
|
|
buffer.seek(0)
|
|
binary_data = buffer.getvalue()
|
|
base64_data = base64.b64encode(binary_data).decode("utf-8")
|
|
|
|
# 全局变量修改
|
|
global base64_data_Array
|
|
base64_data_Array = {"base64": {"file": file, "data": base64_data}}
|
|
# 是否保存到输出目录
|
|
if saveToOutputFolder:
|
|
mesh.export(path.join(full_output_folder, file))
|
|
saved.append({"filename": file, "type": "output", "subfolder": subfolder})
|
|
return {"ui": {"mesh": file}}
|
|
|
|
|
|
class SquareImage:
|
|
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"resolution": (
|
|
"INT",
|
|
{"default": 1024, "min": 8, "max": 8096, "step": 16},
|
|
),
|
|
"upscale_method": (cls.upscale_methods,),
|
|
"padding_color": ("COLOR", {"default": (255, 255, 255)}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("Square_Image",)
|
|
FUNCTION = "make_square"
|
|
CATEGORY = "EX"
|
|
|
|
def make_square(
|
|
self, image, resolution, upscale_method, padding_color=(255, 255, 255)
|
|
):
|
|
ret_images = []
|
|
|
|
# 映射用户选择的缩放方法到 PIL 的 resample 方法
|
|
resample_methods = {
|
|
"nearest-exact": Image.NEAREST,
|
|
"bilinear": Image.BILINEAR,
|
|
"area": Image.BOX,
|
|
"bicubic": Image.BICUBIC,
|
|
"lanczos": Image.LANCZOS,
|
|
}
|
|
|
|
resample = resample_methods.get(upscale_method, Image.NEAREST)
|
|
|
|
for img_tensor in image:
|
|
# 将张量转换为 PIL 图像
|
|
pil_image = tensor2pil(img_tensor)
|
|
|
|
# 获取原始尺寸
|
|
original_width, original_height = pil_image.size
|
|
|
|
# 判断是否已经是正方形
|
|
if original_width == original_height:
|
|
square_image = pil_image
|
|
else:
|
|
# 计算填充尺寸
|
|
max_side = max(original_width, original_height)
|
|
delta_w = max_side - original_width
|
|
delta_h = max_side - original_height
|
|
padding = (
|
|
delta_w // 2,
|
|
delta_h // 2,
|
|
delta_w - (delta_w // 2),
|
|
delta_h - (delta_h // 2),
|
|
)
|
|
|
|
# 填充图像
|
|
square_image = ImageOps.expand(pil_image, padding, fill=padding_color)
|
|
|
|
# 缩放到指定分辨率
|
|
square_image = square_image.resize(
|
|
(resolution, resolution), resample=resample
|
|
)
|
|
|
|
# 转换回张量
|
|
tensor_img = pil2tensor(square_image)
|
|
ret_images.append(tensor_img)
|
|
|
|
# 将处理后的图像张量堆叠起来
|
|
return (torch.cat(ret_images, dim=0),)
|
|
|
|
|
|
# Register the node
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Hunyuan3DNode": Hunyuan3DNode,
|
|
"GenerateSixViews": GenerateSixViews,
|
|
"SquareImage": SquareImage,
|
|
"RemoveBackground": RemoveBackground,
|
|
"TriMeshViewer": TriMeshViewer,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Hunyuan3DNode": "EX_Hunyuan3DNode",
|
|
"GenerateSixViews": "EX_GenerateSixViews",
|
|
"SquareImage": "EX_SquareImage",
|
|
"RemoveBackground": "EX_RemoveBackground",
|
|
"TriMeshViewer": "EX_TriMeshViewer",
|
|
}
|