313 lines
11 KiB
Python
313 lines
11 KiB
Python
"""
|
|
ComfyUI nodes for TRELLIS.2 image-to-3D on Apple Silicon.
|
|
|
|
Three nodes that mirror the proven path in trellis-mac/generate.py:
|
|
|
|
LoadImage → TrellisMacLoader → TrellisMacGenerate → TrellisMacExport → Preview3D
|
|
|
|
No CUDA-only dependencies. Device is MPS. Metal o_voxel bake is used when
|
|
available, with a KDTree CPU fallback. The macOS GPU watchdog kill is
|
|
detected and surfaced as an actionable error message.
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from . import _pipeline
|
|
from ._bake import export_mesh
|
|
from ._watchdog import is_watchdog_error, watchdog_help_message
|
|
|
|
CATEGORY = "TrellisMac"
|
|
|
|
# Custom ComfyUI types. Declared here so __init__.py's _NODE_TYPES can
|
|
# register them with the validation layer.
|
|
TYPE_PIPELINE = "TRELLIS_PIPELINE"
|
|
TYPE_MESH = "MESHWITHVOXEL"
|
|
|
|
PIPELINE_TYPE_CHOICES = [
|
|
"512 - fast",
|
|
"1024 - higher quality",
|
|
"1024_cascade - best quality",
|
|
]
|
|
|
|
PIPELINE_TYPE_MAP = {
|
|
"512 - fast": "512",
|
|
"1024 - higher quality": "1024",
|
|
"1024_cascade - best quality": "1024_cascade",
|
|
}
|
|
|
|
|
|
def _tensor_to_pil(image_tensor: torch.Tensor):
|
|
"""
|
|
Convert a ComfyUI IMAGE tensor to a PIL.Image.
|
|
|
|
ComfyUI delivers images as float32 [B,H,W,3] or [H,W,3] in [0,1].
|
|
pipeline.run() expects a PIL.Image. (generate.py opens a PIL.Image
|
|
directly from a path; we adapt the in-memory tensor instead.)
|
|
"""
|
|
from PIL import Image as PILImage
|
|
|
|
if image_tensor.dim() == 4:
|
|
image_tensor = image_tensor[0]
|
|
arr = (image_tensor.cpu().float().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
|
return PILImage.fromarray(arr)
|
|
|
|
|
|
class TrellisMacLoader:
|
|
"""Loads the TRELLIS.2 pipeline onto MPS. Cached for the session."""
|
|
|
|
DESCRIPTION = (
|
|
"Loads the TRELLIS.2 image-to-3D pipeline onto Apple Silicon (MPS).\n\n"
|
|
"Launch ComfyUI with PYTORCH_ENABLE_MPS_FALLBACK=1 and set "
|
|
"TRELLIS_MAC_PATH if trellis-mac is not auto-detected.\n\n"
|
|
"First run downloads ~16 GB of model weights (cached after that).\n"
|
|
"Requires HuggingFace auth and accepted licenses for:\n"
|
|
" • facebook/dinov3-vitl16-pretrain-lvd1689m\n"
|
|
" • briaai/RMBG-2.0\n\n"
|
|
"This node caches the pipeline for the session — it only loads once.\n"
|
|
"Set reload=True to force a re-load."
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_path": (
|
|
"STRING",
|
|
{"default": "microsoft/TRELLIS.2-4B", "multiline": False},
|
|
),
|
|
"reload": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
# Persist a trellis-mac path override without editing env vars.
|
|
# Setting it writes config.json (one-time); leave blank to use
|
|
# the env var / heuristic / fallback resolution in _pipeline.py.
|
|
"trellis_mac_path": (
|
|
"STRING",
|
|
{"default": "", "multiline": False},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (TYPE_PIPELINE,)
|
|
RETURN_NAMES = ("pipeline",)
|
|
FUNCTION = "load"
|
|
CATEGORY = CATEGORY
|
|
|
|
def load(self, model_path, reload, trellis_mac_path=""):
|
|
if trellis_mac_path:
|
|
_pipeline.save_trellis_mac_path(trellis_mac_path)
|
|
if reload:
|
|
_pipeline.clear_cache()
|
|
|
|
pipeline = _pipeline.get_pipeline(model_path)
|
|
return (pipeline,)
|
|
|
|
|
|
class TrellisMacGenerate:
|
|
"""Run image→3D generation. The core node."""
|
|
|
|
DESCRIPTION = (
|
|
"Generates a 3D mesh from a single image using TRELLIS.2 on MPS.\n\n"
|
|
"• pipeline_type: 512 (fastest) · 1024 · 1024_cascade (best quality)\n"
|
|
"• steps: 0 = use pipeline defaults (usually 12). Higher = more detail, slower.\n"
|
|
"• seed: change for different results from the same image.\n\n"
|
|
"If the mesh comes out empty, see the Loader node's notes about the "
|
|
"macOS GPU watchdog workaround."
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline": (TYPE_PIPELINE,),
|
|
"image": ("IMAGE",),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 2**32 - 1}),
|
|
"pipeline_type": (PIPELINE_TYPE_CHOICES,),
|
|
"steps": (
|
|
"INT",
|
|
{"default": 0, "min": 0, "max": 200,
|
|
"tooltip": "0 = use pipeline.json defaults (usually 12)"},
|
|
),
|
|
},
|
|
"optional": {
|
|
"num_samples": ("INT", {"default": 1, "min": 1, "max": 8}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (TYPE_MESH,)
|
|
RETURN_NAMES = ("mesh",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = CATEGORY
|
|
|
|
def generate(self, pipeline, image, seed, pipeline_type, steps, num_samples=1):
|
|
try:
|
|
steps = int(steps)
|
|
except (TypeError, ValueError):
|
|
raise ValueError(f"steps must be an integer or 0 for defaults; got {steps!r}")
|
|
if steps > 200:
|
|
raise ValueError(
|
|
f"steps={steps} is almost certainly a broken workflow value. "
|
|
"Use steps=0 for TRELLIS.2 defaults, or a small value like 8-20. "
|
|
"A value like 1024 can request tens of GB of MPS memory."
|
|
)
|
|
try:
|
|
num_samples = int(num_samples)
|
|
except (TypeError, ValueError):
|
|
raise ValueError(f"num_samples must be an integer; got {num_samples!r}")
|
|
try:
|
|
pipeline_type = PIPELINE_TYPE_MAP[pipeline_type]
|
|
except KeyError:
|
|
raise ValueError(
|
|
f"pipeline_type must be one of {PIPELINE_TYPE_CHOICES}; "
|
|
f"got {pipeline_type!r}"
|
|
)
|
|
|
|
pil_img = _tensor_to_pil(image)
|
|
|
|
# steps=0 means "use pipeline.json defaults" — pass no override.
|
|
# Mirrors generate.py:96 ("sampler_overrides = {...} if args.steps else {}").
|
|
sampler_overrides = {"steps": steps} if steps and steps > 0 else {}
|
|
print(
|
|
"[TrellisMac] generate "
|
|
f"pipeline_type={pipeline_type!r} seed={seed} "
|
|
f"steps={steps} num_samples={num_samples}",
|
|
flush=True,
|
|
)
|
|
|
|
try:
|
|
outputs = pipeline.run(
|
|
pil_img,
|
|
seed=seed,
|
|
num_samples=num_samples,
|
|
pipeline_type=pipeline_type,
|
|
sparse_structure_sampler_params=sampler_overrides,
|
|
shape_slat_sampler_params=sampler_overrides,
|
|
tex_slat_sampler_params=sampler_overrides,
|
|
)
|
|
except (IndexError, AssertionError) as e:
|
|
# Watchdog-corruption signatures from generate.py:128-142.
|
|
if is_watchdog_error(e):
|
|
raise RuntimeError(watchdog_help_message()) from e
|
|
raise
|
|
|
|
# pipeline.run returns List[MeshWithVoxel]; take the first sample.
|
|
# (generate.py:146 — "outputs[0] if isinstance(outputs, list) else outputs")
|
|
mesh_out = outputs[0] if isinstance(outputs, list) else outputs
|
|
|
|
# Validate the mesh isn't empty — same watchdog guard as
|
|
# generate.py:150-152, in case the corruption didn't raise above.
|
|
verts = mesh_out.vertices
|
|
faces = mesh_out.faces
|
|
v_np = verts.cpu().numpy() if hasattr(verts, "cpu") else np.asarray(verts)
|
|
f_np = faces.cpu().numpy() if hasattr(faces, "cpu") else np.asarray(faces)
|
|
if v_np.shape[0] == 0 or f_np.shape[0] == 0:
|
|
raise RuntimeError(watchdog_help_message())
|
|
|
|
return (mesh_out,)
|
|
|
|
|
|
class TrellisMacExport:
|
|
"""Bake PBR textures and export a GLB. OUTPUT_NODE — feeds Preview3D."""
|
|
|
|
DESCRIPTION = (
|
|
"Bakes PBR textures and exports a 3D file to ComfyUI's output folder.\n\n"
|
|
"Automatically picks the best bake method:\n"
|
|
" 1. Metal GPU bake (fastest, via o_voxel + mtldiffrast)\n"
|
|
" 2. KDTree CPU fallback (xatlas UV unwrap + scipy — same quality, slower)\n"
|
|
" 3. Vertex colors only (no bake, if no voxel data)\n\n"
|
|
"Mesh is pre-simplified to ≤200K faces before baking.\n"
|
|
"The glb_path output plugs into ComfyUI's built-in Preview 3D node."
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mesh": (TYPE_MESH,),
|
|
"filename_prefix": (
|
|
"STRING",
|
|
{"default": "TrellisMac/3D", "multiline": False},
|
|
),
|
|
"file_format": (["glb", "obj", "ply"],),
|
|
"texture_size": (["1024", "512", "2048"],),
|
|
"bake_texture": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"also_obj": ("BOOLEAN", {"default": False,
|
|
"tooltip": "Also write a plain .obj alongside"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("glb_path", "relative_path")
|
|
FUNCTION = "export"
|
|
CATEGORY = CATEGORY
|
|
OUTPUT_NODE = True
|
|
|
|
def export(self, mesh, filename_prefix, file_format, texture_size,
|
|
bake_texture, also_obj=False):
|
|
import folder_paths
|
|
from pathlib import Path
|
|
|
|
try:
|
|
texture_size = int(texture_size)
|
|
except (TypeError, ValueError):
|
|
raise ValueError(f"texture_size must be 512, 1024, or 2048; got {texture_size!r}")
|
|
|
|
# Resolve an output path the same way the reference ExportMesh node
|
|
# does (ComfyUI-Trellis2-ref/nodes.py:904): folder_paths handles the
|
|
# counter + subfolder creation.
|
|
full_output_folder, filename, counter, subfolder, prefix = (
|
|
folder_paths.get_save_image_path(
|
|
filename_prefix, folder_paths.get_output_directory()
|
|
)
|
|
)
|
|
out_dir = Path(full_output_folder)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
ext = file_format
|
|
out_path = out_dir / f"{filename}_{counter:05}_.{ext}"
|
|
|
|
# OBJ: skip baking entirely, write geometry directly.
|
|
if file_format == "obj":
|
|
from ._bake import _export_obj
|
|
_export_obj(mesh, str(out_path))
|
|
elif file_format == "ply":
|
|
import trimesh
|
|
verts = mesh.vertices.cpu().numpy()
|
|
faces = mesh.faces.cpu().numpy()
|
|
trimesh.Trimesh(vertices=verts, faces=faces).export(str(out_path))
|
|
else:
|
|
# GLB: full Metal / KDTree / vertex-color decision tree.
|
|
export_mesh(
|
|
mesh,
|
|
glb_path=str(out_path),
|
|
texture_size=texture_size,
|
|
bake_texture=bake_texture,
|
|
also_export_obj=also_obj,
|
|
)
|
|
|
|
relative_path = Path(subfolder) / f"{filename}_{counter:05}_.{ext}"
|
|
return (str(out_path), str(relative_path))
|
|
|
|
|
|
# ComfyUI node registration. __init__.py imports these names.
|
|
NODE_CLASSES = (
|
|
TrellisMacLoader,
|
|
TrellisMacGenerate,
|
|
TrellisMacExport,
|
|
)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TrellisMacLoader": TrellisMacLoader,
|
|
"TrellisMacGenerate": TrellisMacGenerate,
|
|
"TrellisMacExport": TrellisMacExport,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TrellisMacLoader": "Load TrellisMac Pipeline",
|
|
"TrellisMacGenerate": "TrellisMac Image to 3D",
|
|
"TrellisMacExport": "Export TrellisMac Mesh",
|
|
}
|