499 lines
18 KiB
Python
499 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import json
|
|
import torch
|
|
from PIL import Image
|
|
from typing_extensions import override
|
|
import folder_paths
|
|
from folder_paths import get_output_directory, get_input_directory
|
|
from comfy_api.latest import ComfyExtension, io
|
|
|
|
from .api.utils import Hitem3dAPI
|
|
from comfy_api_nodes.util import (
|
|
download_url_to_file_3d,
|
|
)
|
|
|
|
__version__ = "1.0.7"
|
|
|
|
Hi3D_ak = os.environ.get("Hi3D_ak")
|
|
Hi3D_sk = os.environ.get("Hi3D_sk")
|
|
|
|
if not Hi3D_ak or not Hi3D_sk:
|
|
p = os.path.dirname(os.path.realpath(__file__))
|
|
config_path = os.path.join(p, "config.json")
|
|
if os.path.exists(config_path):
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
if not Hi3D_ak:
|
|
Hi3D_ak = config.get("Hi3D_ak")
|
|
if not Hi3D_sk:
|
|
Hi3D_sk = config.get("Hi3D_sk")
|
|
|
|
# ── Constants ─────────────────────────────────────────────────────────
|
|
|
|
IMAGE_TO_3D_SCENE_MODELS = {
|
|
"general": ["hitem3dv2.1", "hitem3dv1.5", "hitem3dv2.0"],
|
|
"portrait": ["scene-portraitv2.1", "scene-portraitv1.5", "scene-portraitv2.0"],
|
|
}
|
|
|
|
IMAGE_TO_3D_MODEL_RESOLUTIONS = {
|
|
"hitem3dv1.5": ["512", "1024", "1536", "1536pro"],
|
|
"hitem3dv2.0": ["512", "1024", "1536", "1536pro"],
|
|
"hitem3dv2.1": ["1536fast", "1536pro"],
|
|
"scene-portraitv1.5": ["1536"],
|
|
"scene-portraitv2.0": ["1536pro"],
|
|
"scene-portraitv2.1": ["1536profast", "1536pro"],
|
|
}
|
|
|
|
ALL_IMAGE_TO_3D_MODELS = [
|
|
"hitem3dv2.1",
|
|
"hitem3dv1.5",
|
|
"hitem3dv2.0",
|
|
"scene-portraitv1.5",
|
|
"scene-portraitv2.0",
|
|
"scene-portraitv2.1",
|
|
]
|
|
|
|
ALL_IMAGE_TO_3D_RESOLUTIONS = [
|
|
"512", "1024", "1536", "1536fast", "1536pro", "1536profast"
|
|
]
|
|
|
|
PBR_SUPPORTED_MODELS = {"hitem3dv2.0", "hitem3dv2.1"}
|
|
|
|
|
|
|
|
ALL_FORMATS = ["glb", "stl", "fbx", "usdz"]
|
|
|
|
Hi3DModelTask = io.Custom("HI3D_MODEL_TASK")
|
|
|
|
# ── Helpers ────────────────────────────────────────────────────────────
|
|
|
|
|
|
def getHitem3dAPI(ak: str, sk: str):
|
|
ak = Hi3D_ak if Hi3D_ak else ak
|
|
sk = Hi3D_sk if Hi3D_sk else sk
|
|
if not ak:
|
|
raise RuntimeError("ak is required")
|
|
if not sk:
|
|
raise RuntimeError("sk is required")
|
|
return Hitem3dAPI(ak, sk), ak, sk
|
|
|
|
|
|
def get_default_model_for_scene(scene):
|
|
return IMAGE_TO_3D_SCENE_MODELS[scene][0]
|
|
|
|
|
|
def infer_scene_from_model(model):
|
|
for scene, models in IMAGE_TO_3D_SCENE_MODELS.items():
|
|
if model in models:
|
|
return scene
|
|
return None
|
|
|
|
|
|
def validate_image_to_3d_options(scene, model, resolution, texture, pbr):
|
|
allowed_models = IMAGE_TO_3D_SCENE_MODELS.get(scene)
|
|
if allowed_models is None:
|
|
valid_scenes = ", ".join(IMAGE_TO_3D_SCENE_MODELS.keys())
|
|
raise RuntimeError(f"Invalid scene '{scene}'. Valid scenes: {valid_scenes}")
|
|
|
|
if model not in allowed_models:
|
|
allowed_models_text = ", ".join(allowed_models)
|
|
raise RuntimeError(
|
|
f"Model '{model}' is not available for scene '{scene}'. "
|
|
f"Allowed models: {allowed_models_text}"
|
|
)
|
|
|
|
allowed_resolutions = IMAGE_TO_3D_MODEL_RESOLUTIONS[model]
|
|
if resolution not in allowed_resolutions:
|
|
allowed_resolutions_text = ", ".join(allowed_resolutions)
|
|
raise RuntimeError(
|
|
f"Resolution '{resolution}' is not available for model '{model}'. "
|
|
f"Allowed resolutions: {allowed_resolutions_text}"
|
|
)
|
|
|
|
pbr_enabled = bool(pbr)
|
|
if pbr_enabled:
|
|
if not texture:
|
|
raise RuntimeError("PBR requires texture to be enabled")
|
|
if model not in PBR_SUPPORTED_MODELS:
|
|
allowed_models_text = ", ".join(sorted(PBR_SUPPORTED_MODELS))
|
|
raise RuntimeError(
|
|
f"PBR is only supported for models: {allowed_models_text}"
|
|
)
|
|
|
|
|
|
|
|
|
|
def save_tensor(image_tensor, filename):
|
|
if image_tensor.dim() > 3:
|
|
image_tensor = image_tensor[0]
|
|
if image_tensor.dtype == torch.float32:
|
|
image_tensor = (image_tensor * 255).byte()
|
|
if image_tensor.dim() == 2:
|
|
image_tensor = image_tensor.unsqueeze(0)
|
|
if image_tensor.dim() == 3 and image_tensor.size(0) == 3:
|
|
image_tensor = image_tensor.permute(1, 2, 0)
|
|
if image_tensor.is_cuda:
|
|
image_tensor = image_tensor.cpu()
|
|
image_np = image_tensor.numpy()
|
|
image_pil = Image.fromarray(image_np)
|
|
if image_np.shape[2] == 4:
|
|
name = filename + ".png"
|
|
image_pil.save(name, "PNG")
|
|
else:
|
|
name = filename + ".jpg"
|
|
image_pil.save(name, "JPEG")
|
|
return name
|
|
|
|
|
|
def toImagePath(image, filename):
|
|
if image is None:
|
|
return None
|
|
if not isinstance(image, (str, bytes, os.PathLike)):
|
|
return save_tensor(image, os.path.join(get_input_directory(), filename))
|
|
return image
|
|
|
|
|
|
def build_multi_view_payload(image_path, image_back=None, image_left=None, image_right=None):
|
|
multi_images = [image_path]
|
|
image_dict = {
|
|
"image_back": image_back,
|
|
"image_left": image_left,
|
|
"image_right": image_right,
|
|
}
|
|
multi_images_bit = "1"
|
|
for image_name in ["image_back", "image_left", "image_right"]:
|
|
img = image_dict[image_name]
|
|
if img is not None:
|
|
image_filename = toImagePath(img, image_name)
|
|
multi_images.append(image_filename)
|
|
multi_images_bit += "1"
|
|
else:
|
|
multi_images_bit += "0"
|
|
multi_images.append(None)
|
|
return multi_images, multi_images_bit
|
|
|
|
|
|
async def run_image_to_3d_task(
|
|
ak, sk, texture, scene, pbr,
|
|
image=None, image_back=None, image_left=None, image_right=None,
|
|
model=None, face=None, resolution=None,
|
|
):
|
|
if image is None:
|
|
raise RuntimeError("image is required")
|
|
|
|
pbr_enabled = bool(pbr)
|
|
if pbr_enabled:
|
|
texture = True
|
|
|
|
if not scene:
|
|
scene = infer_scene_from_model(model) or "general"
|
|
if not model:
|
|
model = get_default_model_for_scene(scene)
|
|
if not resolution:
|
|
default_resolution = "1024"
|
|
allowed_resolutions = IMAGE_TO_3D_MODEL_RESOLUTIONS[model]
|
|
resolution = (
|
|
default_resolution
|
|
if default_resolution in allowed_resolutions
|
|
else allowed_resolutions[0]
|
|
)
|
|
|
|
validate_image_to_3d_options(scene, model, resolution, texture, pbr)
|
|
request_type = 3 if texture else 1
|
|
image_path = toImagePath(image, "image")
|
|
api, ak, sk = getHitem3dAPI(ak, sk)
|
|
|
|
if image_back is None and image_left is None and image_right is None:
|
|
result = await api.image_to_3d(
|
|
image_path, request_type, face, model, resolution, pbr
|
|
)
|
|
else:
|
|
multi_images, multi_images_bit = build_multi_view_payload(
|
|
image_path, image_back=image_back, image_left=image_left, image_right=image_right
|
|
)
|
|
result = await api.multi_view_to_3d(
|
|
multi_images, multi_images_bit, request_type, face, model, resolution, pbr
|
|
)
|
|
|
|
if result["status"] != "success":
|
|
raise RuntimeError(f"Failed to generate mesh: {result['message']}")
|
|
task_id = result["task_id"]
|
|
model_url = result["model_url"]
|
|
glb = await download_url_to_file_3d(model_url, "glb", task_id=task_id)
|
|
|
|
return glb, {"task_id": task_id, "model_url": model_url, "ak": ak, "sk": sk}
|
|
|
|
|
|
# ── Node definitions ─────────────────────────────────────────────────────
|
|
|
|
class ImageTo3DNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ImageTo3DNode",
|
|
display_name="Hi3D:ImageTo3D",
|
|
category="Hi3D",
|
|
description=(
|
|
"Generates a 3D model from a single or multi-view image. "
|
|
"Use `scene` to switch between general object models and portrait/scene models. "
|
|
"Model and resolution options are linked together in the node UI."
|
|
),
|
|
inputs=[
|
|
io.String.Input(
|
|
"ak", default="",
|
|
tooltip="Hi3D API Access Key. Leave empty if configured in config.json or environment variables.",
|
|
),
|
|
io.String.Input(
|
|
"sk", default="",
|
|
tooltip="Hi3D API Secret Key. Leave empty if configured in config.json or environment variables.",
|
|
),
|
|
io.Image.Input(
|
|
"image", optional=True,
|
|
tooltip="Main front image (Required) for 3D model generation.",
|
|
),
|
|
io.Image.Input(
|
|
"image_back", optional=True,
|
|
tooltip="Back reference image (Optional) for multi-view 3D reconstruction.",
|
|
),
|
|
io.Image.Input(
|
|
"image_left", optional=True,
|
|
tooltip="Left reference image (Optional) for multi-view 3D reconstruction.",
|
|
),
|
|
io.Image.Input(
|
|
"image_right", optional=True,
|
|
tooltip="Right reference image (Optional) for multi-view 3D reconstruction.",
|
|
),
|
|
io.Boolean.Input(
|
|
"texture", default=False, optional=True,
|
|
tooltip="Whether to generate textures. If enabled, the model will include vertex colors/textures.",
|
|
),
|
|
io.Boolean.Input(
|
|
"pbr", default=True, optional=True,
|
|
tooltip="PBR开关（默认开启）。PBR需要texture开启，并且仅支持hitem3dv2.0/hitem3dv2.1。",
|
|
),
|
|
io.Combo.Input(
|
|
"scene", options=["general", "portrait"], default="general", optional=True,
|
|
tooltip="Scene preset. `general` is for common objects, `portrait` is for portrait/scene generation.",
|
|
),
|
|
io.Combo.Input(
|
|
"model", options=ALL_IMAGE_TO_3D_MODELS, default="hitem3dv2.1", optional=True,
|
|
tooltip="3D generation model version. The node UI filters valid versions according to the selected scene.",
|
|
),
|
|
io.Combo.Input(
|
|
"resolution", options=ALL_IMAGE_TO_3D_RESOLUTIONS, default="1024", optional=True,
|
|
tooltip="Generation resolution. The node UI filters valid resolutions according to the selected model.",
|
|
),
|
|
io.Int.Input(
|
|
"face", min=100000, max=2000000, default=500000, optional=True,
|
|
tooltip="Output poly count. Higher count means finer detail but larger file size.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.File3DGLB.Output(
|
|
display_name="GLB",
|
|
),
|
|
Hi3DModelTask.Output(
|
|
display_name="model_task",
|
|
tooltip="Task information to be passed to the Texture node for re-texturing",
|
|
),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls, ak, sk,
|
|
image=None, image_back=None, image_left=None, image_right=None,
|
|
texture=False, pbr=True, scene="general", model=None, resolution=None,
|
|
face=500000,
|
|
):
|
|
glb, task_info = await run_image_to_3d_task(
|
|
ak, sk, texture, scene, pbr,
|
|
image=image, image_back=image_back,
|
|
image_left=image_left, image_right=image_right,
|
|
model=model, face=face, resolution=resolution,
|
|
)
|
|
return io.NodeOutput(glb, task_info)
|
|
|
|
|
|
class TextureNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="TextureNode",
|
|
display_name="Hi3D:Texture",
|
|
category="Hi3D",
|
|
description=(
|
|
"Regenerates texture maps for existing 3D models. "
|
|
"Provide a reference image and the original model (via path/URL or upstream node) "
|
|
"to generate new textures."
|
|
),
|
|
inputs=[
|
|
io.String.Input(
|
|
"ak", default="",
|
|
tooltip="Hi3D API Access Key. Leave empty if configured in config.json or environment variables.",
|
|
),
|
|
io.String.Input(
|
|
"sk", default="",
|
|
tooltip="Hi3D API Secret Key. Leave empty if configured in config.json or environment variables.",
|
|
),
|
|
io.MultiType.Input(
|
|
"GLB",
|
|
types=[io.File3DGLB, io.File3DAny],
|
|
optional=True,
|
|
tooltip="Input 3D model in GLB format. Either this or model_task is required.",
|
|
),
|
|
io.Image.Input(
|
|
"image", optional=True,
|
|
tooltip="Texture reference image (Required). The model will generate a new texture map based on this image.",
|
|
),
|
|
Hi3DModelTask.Input(
|
|
"model_task", optional=True,
|
|
tooltip="Task output from ImageTo3D nodes. Automatically uses that model's URL if connected.",
|
|
),
|
|
io.Combo.Input(
|
|
"model", options=["hitem3dv2.1", "hitem3dv1.5", "hitem3dv2.0", "scene-portraitv1.5"],
|
|
default="hitem3dv2.1", optional=True,
|
|
tooltip="Model version used for texture generation.",
|
|
),
|
|
io.Boolean.Input(
|
|
"pbr", default=True, optional=True,
|
|
tooltip="PBR开关（默认开启）。仅hitem3dv2.0/hitem3dv2.1支持。",
|
|
),
|
|
|
|
],
|
|
outputs=[
|
|
io.File3DGLB.Output(
|
|
display_name="GLB",
|
|
),
|
|
Hi3DModelTask.Output(
|
|
display_name="model_task",
|
|
tooltip="Task information to be passed to subsequent nodes",
|
|
),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls, ak, sk,
|
|
GLB=None, image=None, model_task=None, model=None, pbr=True,
|
|
):
|
|
if image is None:
|
|
raise RuntimeError("image is required")
|
|
if bool(pbr) and model not in PBR_SUPPORTED_MODELS:
|
|
allowed_models_text = ", ".join(sorted(PBR_SUPPORTED_MODELS))
|
|
raise RuntimeError(
|
|
f"PBR is only supported for models: {allowed_models_text}"
|
|
)
|
|
|
|
if model_task is not None:
|
|
if not ak and "ak" in model_task:
|
|
ak = model_task["ak"]
|
|
if not sk and "sk" in model_task:
|
|
sk = model_task["sk"]
|
|
|
|
api, ak, sk = getHitem3dAPI(ak, sk)
|
|
|
|
if model_task is not None:
|
|
mesh_url = model_task["model_url"]
|
|
elif GLB is not None:
|
|
if GLB.is_disk_backed:
|
|
mesh_path = GLB.get_source()
|
|
else:
|
|
tmp_path = os.path.join(get_input_directory(), f"_tmp_texture_input.{GLB.format or 'glb'}")
|
|
GLB.save_to(tmp_path)
|
|
mesh_path = tmp_path
|
|
mesh_url = await api.upload_file(mesh_path)
|
|
else:
|
|
raise RuntimeError("GLB or model_task is required")
|
|
image_path = toImagePath(image, "image")
|
|
result = await api.texture(image_path, mesh_url, model, pbr)
|
|
|
|
if result["status"] != "success":
|
|
raise RuntimeError(f"Failed to generate mesh: {result['message']}")
|
|
|
|
task_id = result["task_id"]
|
|
model_url = result["model_url"]
|
|
glb = await download_url_to_file_3d(model_url, "glb", task_id=task_id)
|
|
return io.NodeOutput(
|
|
glb,
|
|
{"task_id": task_id, "model_url": model_url, "ak": ak, "sk": sk},
|
|
)
|
|
|
|
|
|
class LoadGLBNode:
|
|
SUPPORTED_EXTENSIONS = {'.glb'}
|
|
|
|
@classmethod
|
|
def _scan_files(cls):
|
|
from pathlib import Path
|
|
results = []
|
|
for base_dir in [get_input_directory(), get_output_directory()]:
|
|
base = Path(base_dir)
|
|
if not base.exists():
|
|
continue
|
|
for f in base.rglob("*"):
|
|
if f.suffix.lower() in cls.SUPPORTED_EXTENSIONS:
|
|
results.append(str(f.relative_to(base)).replace("\\", "/"))
|
|
return sorted(set(results))
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
files = cls._scan_files()
|
|
return {
|
|
"required": {
|
|
"model_file": (sorted(files) if files else ["(no files found)"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FILE_3D",)
|
|
RETURN_NAMES = ("model_3d",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "Hi3D"
|
|
DESCRIPTION = "Load a GLB 3D model file from disk."
|
|
|
|
def execute(self, model_file):
|
|
from comfy_api.latest._util.geometry_types import File3D
|
|
|
|
filepath = folder_paths.get_annotated_filepath(model_file)
|
|
if not os.path.isfile(filepath):
|
|
raise RuntimeError(f"3D model file not found: {filepath}")
|
|
return (File3D(filepath),)
|
|
|
|
|
|
# ── Extension registration ───────────────────────────────────────────────
|
|
|
|
|
|
class Hitem3DExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
ImageTo3DNode,
|
|
TextureNode,
|
|
]
|
|
|
|
|
|
NODES_LIST: list[type[io.ComfyNode]] = [
|
|
ImageTo3DNode,
|
|
TextureNode,
|
|
]
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageTo3DNode": ImageTo3DNode,
|
|
"TextureNode": TextureNode,
|
|
"LoadGLBNode": LoadGLBNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageTo3DNode": "Hi3D:ImageTo3D",
|
|
"TextureNode": "Hi3D:Texture",
|
|
"LoadGLBNode": "Hi3D:Load3DModel",
|
|
}
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
|
|
async def comfy_entrypoint() -> Hitem3DExtension:
|
|
return Hitem3DExtension()
|