360 lines
14 KiB
Python
360 lines
14 KiB
Python
import os
|
|
import shutil
|
|
import warnings
|
|
import sys
|
|
warnings.simplefilter('ignore', category=UserWarning)
|
|
warnings.simplefilter('ignore', category=FutureWarning)
|
|
warnings.simplefilter('ignore', category=DeprecationWarning)
|
|
from huggingface_hub import snapshot_download
|
|
import torch
|
|
import numpy as np
|
|
from einops import rearrange
|
|
from PIL import Image, ImageSequence, ImageOps, ImageColor
|
|
from io import BytesIO
|
|
import ipywidgets as widgets
|
|
from IPython.display import display
|
|
import tempfile
|
|
import torch.nn.functional as F
|
|
from typing import Tuple
|
|
import folder_paths
|
|
|
|
# 导入您的 infer 模块中的必要类和函数
|
|
from .infer import seed_everything, save_gif
|
|
from .infer import Removebg, Image2Views, Views2Mesh, GifRenderer
|
|
# 添加当前目录到 sys.path
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if current_dir not in sys.path:
|
|
sys.path.append(current_dir)
|
|
|
|
|
|
class Hunyuan3DNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"seed": ("INT", {"default": 0}),
|
|
"step": ("INT", {"default": 50}),
|
|
"max_number_of_faces": ("INT", {"default": 90000}),
|
|
"do_texture_mapping": ("BOOLEAN", {"default": False}),
|
|
"do_render_gif": ("BOOLEAN", {"default": False}),
|
|
"use_lite": ("BOOLEAN", {"default": False}), # 新增参数
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "STRING", "STRING", "STRING", )
|
|
RETURN_NAMES = ("rmbg_image", "multiview_image", "obj_file_path", "glb_file_path", "gif_file_path", )
|
|
FUNCTION = "imgTo3D"
|
|
CATEGORY = "Image/3D"
|
|
|
|
def __init__(self):
|
|
# 移除 super().__init__() 调用
|
|
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
# 将模型初始化为 None，待执行时再加载
|
|
self.worker_xbg = None
|
|
self.worker_i2v = None
|
|
self.worker_v23 = None
|
|
self.worker_gif = 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, image, seed, step, max_number_of_faces, do_texture_mapping, do_render_gif, use_lite):
|
|
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}")
|
|
|
|
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 或网络连接。")
|
|
|
|
if self.worker_xbg is None:
|
|
self.worker_xbg = Removebg()
|
|
if self.worker_i2v is None:
|
|
self.worker_i2v = Image2Views(
|
|
use_lite=use_lite,
|
|
device=self.device
|
|
)
|
|
if use_lite:
|
|
config_path = os.path.join(base_dir, "svrm/configs/svrm.yaml")
|
|
weights_path = os.path.join(base_dir, "weights/svrm/svrm.safetensors")
|
|
else:
|
|
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
|
|
)
|
|
if self.worker_gif is None:
|
|
self.worker_gif = GifRenderer(self.device)
|
|
|
|
save_folder = self.prepare_save_folder()
|
|
|
|
try:
|
|
# 阶段 1：去除背景
|
|
rgba = self.worker_xbg(image)
|
|
rgba.save(os.path.join(save_folder, 'img_nobg.png'))
|
|
|
|
# 阶段 2：生成多视图图像
|
|
res_img, pils = self.worker_i2v(rgba, seed, step)
|
|
save_gif(pils, os.path.join(save_folder, 'views.gif'))
|
|
views_img, cond_img = res_img[0], res_img[1]
|
|
|
|
# 准备显示的多视图图像
|
|
img_array = np.asarray(views_img, dtype=np.uint8)
|
|
show_img = rearrange(img_array, '(n h) (m w) c -> (n m) h w c', n=3, m=2)
|
|
show_img = show_img[self.worker_i2v.order, ...]
|
|
show_img = rearrange(show_img, '(n m) h w c -> (n h) (m w) c', n=2, m=3)
|
|
show_img = Image.fromarray(show_img)
|
|
|
|
# 阶段 3：从视图生成网格
|
|
self.worker_v23(
|
|
views_img, cond_img,
|
|
seed=seed, save_folder=save_folder,
|
|
target_face_count=max_number_of_faces,
|
|
do_texture_mapping=do_texture_mapping
|
|
)
|
|
|
|
# 阶段 4：渲染 GIF
|
|
gif_file_path = None
|
|
if do_render_gif:
|
|
self.worker_gif(
|
|
os.path.join(save_folder, 'mesh.obj'),
|
|
gif_dst_path=os.path.join(save_folder, 'output.gif')
|
|
)
|
|
gif_file_path = os.path.join(save_folder, 'output.gif')
|
|
|
|
|
|
# 准备输出
|
|
rmbg_image = np.array(rgba, dtype=np.uint8) # [H, W, C]
|
|
multiview_image = np.array(show_img, dtype=np.uint8) # [H, W, C]
|
|
obj_file_path = os.path.join(save_folder, 'mesh_with_colors.obj')
|
|
glb_file_path = os.path.join(save_folder, 'mesh.glb')
|
|
|
|
# Convert to torch.Tensor, keep dimensions [H, W, C], normalize to [0, 1]
|
|
rmbg_image_tensor = torch.from_numpy(rmbg_image).float() / 255.0 # [H, W, C]
|
|
multiview_image_tensor = torch.from_numpy(multiview_image).float() / 255.0 # [H, W, C]
|
|
|
|
# Add batch dimension
|
|
rmbg_image_tensor = rmbg_image_tensor.unsqueeze(0) # [1, H, W, C]
|
|
multiview_image_tensor = multiview_image_tensor.unsqueeze(0) # [1, H, W, C]
|
|
|
|
# Return
|
|
outputs = (
|
|
rmbg_image_tensor,
|
|
multiview_image_tensor,
|
|
obj_file_path,
|
|
glb_file_path,
|
|
gif_file_path,
|
|
)
|
|
return outputs
|
|
|
|
except Exception as e:
|
|
print(f"Error in processing: {str(e)}")
|
|
raise
|
|
|
|
def prepare_save_folder(self):
|
|
# 获取当前文件的上级再上一级目录
|
|
base_dir = folder_paths.get_output_directory()
|
|
|
|
# 组合新的保存路径
|
|
output_dir = os.path.join(base_dir, '3D_output')
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 检查已有的子文件夹
|
|
exists = set(int(_) for _ in os.listdir(output_dir) if not _.startswith(".") and _.isdigit())
|
|
if len(exists) == 30:
|
|
shutil.rmtree(os.path.join(output_dir, "0"))
|
|
cur_id = 0
|
|
else:
|
|
cur_id = min(set(range(30)) - exists)
|
|
|
|
# 删除循环中下一个要用的文件夹
|
|
next_folder = os.path.join(output_dir, str((cur_id + 1) % 30))
|
|
if os.path.exists(next_folder):
|
|
shutil.rmtree(next_folder)
|
|
|
|
# 创建当前ID的保存文件夹
|
|
save_folder = os.path.join(output_dir, str(cur_id))
|
|
os.makedirs(save_folder, exist_ok=True)
|
|
return save_folder
|
|
|
|
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 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 = "TTP/Image"
|
|
|
|
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),)
|
|
|
|
class GifImageViewerNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"gif_file_path": ("STRING", {"default": '', "multiline": False, "forceInput": True}),
|
|
},
|
|
}
|
|
|
|
# 明确指定返回类型，这里假设返回的是 UI_IMAGE
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "preview_gif_image"
|
|
CATEGORY = "Image/Visualize"
|
|
|
|
def preview_gif_image(self, gif_file_path):
|
|
|
|
filename, subdirectory = self.get_fn_subf(gif_file_path)
|
|
# 准备返回结果
|
|
results = [
|
|
{
|
|
"filename": filename, # 使用临时文件的路径
|
|
"subfolder": subdirectory, # 如果有子文件夹，填写子文件夹名称，否则为空字符串
|
|
"type": "output"
|
|
}
|
|
]
|
|
|
|
return { "ui": { "images": results } }
|
|
|
|
def get_fn_subf(self, gif_file_path):
|
|
output_directory = folder_paths.get_output_directory()
|
|
filename = os.path.basename(gif_file_path)
|
|
directory_path = os.path.dirname(gif_file_path)
|
|
subdirectory = os.path.relpath(directory_path, output_directory)
|
|
|
|
return filename, subdirectory
|
|
|
|
# Register the node
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Hunyuan3DNode": Hunyuan3DNode,
|
|
"SquareImage" : SquareImage,
|
|
"GifImageViewerNode" : GifImageViewerNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Hunyuan3DNode": "TTP_Hunyuan3DNode",
|
|
"SquareImage": "TTP_SquareImage",
|
|
"GifImageViewerNode": "TTP_GIFViewer",
|
|
}
|