nodes/cuzelac-ComfyUI-Trellis2-MultiViewRefiner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
cuzelac-ComfyUI-Trellis2-Mu…/nodes_multiview_refiner.py
T
Chris UzelacandClaude Opus 4.7 f8c676ae94 Initial commit: multi-view mesh refiner add-on for ComfyUI-Trellis2 
Standalone ComfyUI custom node that adds the Trellis2MeshRefinerMultiView
node. Monkey-patches sample_mesh_slat_multiview and refine_mesh_multiview
onto Trellis2ImageTo3DPipeline lazily (on first node execution) so the
addon is agnostic to sibling-package load order under ComfyUI's loader.

Method bodies are ports of the tested implementation from
https://github.com/visualbruno/ComfyUI-Trellis2/pull/125 — test matrix
covered front-only, front+back, and all-4-view inputs at 512/1024/1536.

Credit to bribblej-web for suggesting the drop-in packaging pattern on
the PR thread.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-04-18 19:59:14 -07:00

523 lines
19 KiB
Python
Raw Permalink Blame History

"""
Multi-view mesh refiner for ComfyUI-Trellis2.
Adds one node: Trellis2MeshRefinerMultiView.
This package sits next to ComfyUI-Trellis2 in custom_nodes/ and does not
import from it at load time — ComfyUI registers custom-node modules under
path-derived sys.modules keys, and load order across sibling packages is
not guaranteed. Instead, this node resolves everything lazily from the
pipeline instance it receives at runtime (via type(pipeline).__module__),
then monkey-patches two methods onto the pipeline class on first use:
- sample_mesh_slat_multiview
- refine_mesh_multiview
The method bodies are ports of the tested implementation from
https://github.com/visualbruno/ComfyUI-Trellis2/pull/125 — tested across
front-only / front+back / all-4-view inputs at resolutions 512, 1024, 1536.
"""
import gc
import logging
import sys
import numpy as np
import torch
from PIL import Image
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Helpers inlined from ComfyUI-Trellis2's nodes.py (kept self-contained so we
# don't have to chase cross-package imports at load time).
# ---------------------------------------------------------------------------
def reset_cuda():
torch.cuda.synchronize()
gc.collect()
torch.cuda.empty_cache()
def tensor2pil(image):
if isinstance(image, torch.Tensor):
t = image.detach().cpu()
if t.ndim == 4:
if t.shape[0] != 1:
raise ValueError(f"tensor2pil expects batch of 1, got batch={t.shape[0]}")
t = t[0]
elif t.ndim != 3:
raise ValueError(f"tensor2pil expects (H,W,C) or (1,H,W,C), got shape={tuple(t.shape)}")
arr = (t.numpy() * 255.0).clip(0, 255).astype(np.uint8)
return Image.fromarray(arr)
raise TypeError(f"tensor2pil expected torch.Tensor, got {type(image)}")
# ---------------------------------------------------------------------------
# Pipeline method: sample_mesh_slat_multiview
# Resolves SparseTensor and samplers lazily from the pipeline's own module,
# so this function is agnostic to how the parent package was registered.
# ---------------------------------------------------------------------------
def _sample_mesh_slat_multiview(
self,
mesh_slat,
conds,
views,
flow_model,
resolution,
sampler_params=None,
max_num_tokens=49152,
downsampling=16,
front_axis='z',
blend_temperature=2.0,
):
if sampler_params is None:
sampler_params = {}
pipeline_mod = sys.modules[type(self).__module__]
SparseTensor = pipeline_mod.SparseTensor
samplers = pipeline_mod.samplers
self.load_shape_slat_decoder()
logger.info('Decoding mesh slat (MultiView) ...')
if self.low_vram:
self.models['shape_slat_decoder'].to(self.device)
self.models['shape_slat_decoder'].low_vram = True
hr_coords = self.models['shape_slat_decoder'].upsample(mesh_slat, upsample_times=4)
if self.low_vram:
self.models['shape_slat_decoder'].cpu()
self.models['shape_slat_decoder'].low_vram = False
hr_resolution = resolution
if not self.keep_models_loaded:
self.unload_shape_slat_decoder()
lr_resolution = resolution
while True:
quant_coords = torch.cat([
hr_coords[:, :1],
((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // downsampling)).int(),
], dim=1)
coords = quant_coords.unique(dim=0)
num_tokens = coords.shape[0]
if num_tokens < max_num_tokens:
if hr_resolution != resolution:
logger.info(f"Due to the limited number of tokens, the resolution is reduced to {hr_resolution}.")
break
hr_resolution -= 128
if hr_resolution < 1024 and resolution >= 1024:
hr_resolution = 1024
break
if hr_resolution < 512:
hr_resolution = 512
break
if self.low_vram:
for v in conds:
conds[v] = self._cond_to(conds[v], self.device)
coords_dev = coords.to(self.device)
noise = SparseTensor(
feats=torch.randn(coords.shape[0], flow_model.in_channels, device=self.device),
coords=coords_dev,
)
sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
sampler_class = getattr(
samplers,
f"Flow{self._sampler_prefix}MultiViewGuidanceIntervalSampler",
samplers.FlowEulerMultiViewGuidanceIntervalSampler,
)
sampler = sampler_class(
sigma_min=1e-5,
resolution=flow_model.resolution if hasattr(flow_model, 'resolution') else flow_model[0].resolution,
)
if self.low_vram:
flow_model.to(self.device)
slat = sampler.sample(
flow_model,
noise,
conds=conds,
views=views,
front_axis=front_axis,
blend_temperature=blend_temperature,
**sampler_params,
verbose=True,
tqdm_desc="Sampling mesh shape SLat (MultiView)",
).samples
if self.low_vram:
flow_model.cpu()
self._cleanup_cuda()
std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
slat = slat * std + mean
del coords_dev
if self.low_vram:
for v in conds:
conds[v] = self._cond_cpu(conds[v])
self._cleanup_cuda()
return slat, hr_resolution
# ---------------------------------------------------------------------------
# Pipeline method: refine_mesh_multiview
# ---------------------------------------------------------------------------
@torch.no_grad()
def _refine_mesh_multiview(
self,
mesh,
front,
back=None,
left=None,
right=None,
seed=42,
shape_slat_sampler_params=None,
tex_slat_sampler_params=None,
resolution=1024,
max_num_tokens=50000,
generate_texture_slat=True,
return_latent=False,
downsampling=16,
use_tiled=True,
front_axis='z',
blend_temperature=2.0,
sampler='euler',
):
if shape_slat_sampler_params is None:
shape_slat_sampler_params = {}
if tex_slat_sampler_params is None:
tex_slat_sampler_params = {}
pipeline_mod = sys.modules[type(self).__module__]
seed_all = pipeline_mod.seed_all
self.switch_samplers(sampler)
mesh = self.preprocess_mesh(mesh)
seed_all(seed)
views_dict = {'front': front}
if back is not None:
views_dict['back'] = back
if left is not None:
views_dict['left'] = left
if right is not None:
views_dict['right'] = right
views_list = list(views_dict.keys())
self.load_image_cond_model()
conds = {}
cond_resolution = 512 if resolution == 512 else 1024
for v, img in views_dict.items():
conds[v] = self.get_cond([img], cond_resolution)
if not self.keep_models_loaded:
self.unload_image_cond_model()
mesh_slat = self.encode_shape_slat(mesh, resolution)
if resolution == 512:
self.unload_shape_slat_flow_model_1024()
self.load_shape_slat_flow_model_512()
shape_slat, res = self.sample_mesh_slat_multiview(
mesh_slat, conds, views_list,
self.models['shape_slat_flow_model_512'],
512, shape_slat_sampler_params,
max_num_tokens, downsampling,
front_axis, blend_temperature,
)
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_512()
if generate_texture_slat:
self.unload_tex_slat_flow_model_1024()
self.load_tex_slat_flow_model_512()
tex_slat = self.sample_tex_slat_multiview(
conds, views_list,
shape_slat=shape_slat,
flow_model=self.models['tex_slat_flow_model_512'],
sampler_params=tex_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
if not self.keep_models_loaded:
self.unload_tex_slat_flow_model_512()
elif resolution == 1024:
self.unload_shape_slat_flow_model_512()
self.load_shape_slat_flow_model_1024()
shape_slat, res = self.sample_mesh_slat_multiview(
mesh_slat, conds, views_list,
self.models['shape_slat_flow_model_1024'],
1024, shape_slat_sampler_params,
max_num_tokens, downsampling,
front_axis, blend_temperature,
)
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_1024()
if generate_texture_slat:
self.unload_tex_slat_flow_model_512()
self.load_tex_slat_flow_model_1024()
tex_slat = self.sample_tex_slat_multiview(
conds, views_list,
shape_slat=shape_slat,
flow_model=self.models['tex_slat_flow_model_1024'],
sampler_params=tex_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
if not self.keep_models_loaded:
self.unload_tex_slat_flow_model_1024()
elif resolution == 1536:
self.unload_shape_slat_flow_model_512()
self.load_shape_slat_flow_model_1024()
shape_slat, res = self.sample_mesh_slat_multiview(
mesh_slat, conds, views_list,
self.models['shape_slat_flow_model_1024'],
1536, shape_slat_sampler_params,
max_num_tokens, downsampling,
front_axis, blend_temperature,
)
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_1024()
if generate_texture_slat:
self.unload_tex_slat_flow_model_512()
self.load_tex_slat_flow_model_1024()
tex_slat = self.sample_tex_slat_multiview(
conds, views_list,
shape_slat=shape_slat,
flow_model=self.models['tex_slat_flow_model_1024'],
sampler_params=tex_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
if not self.keep_models_loaded:
self.unload_tex_slat_flow_model_1024()
torch.cuda.synchronize()
torch.cuda.empty_cache()
if generate_texture_slat:
out_mesh = self.decode_latent(shape_slat, tex_slat, res, use_tiled=use_tiled)
else:
out_mesh = self.decode_latent(shape_slat, None, res, use_tiled=use_tiled)
torch.cuda.synchronize()
torch.cuda.empty_cache()
if return_latent:
if generate_texture_slat:
return out_mesh, (shape_slat, tex_slat, res)
else:
return out_mesh, (shape_slat, None, res)
else:
return out_mesh
# ---------------------------------------------------------------------------
# Deferred pipeline patching. Runs once on first process() call.
# ---------------------------------------------------------------------------
_REQUIRED_PIPELINE_ATTRS = [
'load_shape_slat_decoder', 'unload_shape_slat_decoder',
'load_shape_slat_flow_model_512', 'unload_shape_slat_flow_model_512',
'load_shape_slat_flow_model_1024', 'unload_shape_slat_flow_model_1024',
'load_tex_slat_flow_model_512', 'unload_tex_slat_flow_model_512',
'load_tex_slat_flow_model_1024', 'unload_tex_slat_flow_model_1024',
'load_image_cond_model', 'unload_image_cond_model',
'get_cond', 'encode_shape_slat', 'preprocess_mesh', 'switch_samplers',
'sample_tex_slat_multiview', 'decode_latent',
'_cond_to', '_cond_cpu', '_cleanup_cuda',
]
_REQUIRED_MODULE_ATTRS = ['SparseTensor', 'samplers', 'seed_all']
_patched_classes = set()
def _ensure_patched(pipeline_cls):
if pipeline_cls in _patched_classes:
return
missing_cls = [a for a in _REQUIRED_PIPELINE_ATTRS if not hasattr(pipeline_cls, a)]
if missing_cls:
raise RuntimeError(
"ComfyUI-Trellis2-MultiViewRefiner is out of date: "
f"{pipeline_cls.__name__} is missing expected helpers {missing_cls}. "
"Upstream ComfyUI-Trellis2 has changed. "
"Check for an updated version at "
"https://github.com/cuzelac/ComfyUI-Trellis2-MultiViewRefiner"
)
pipeline_mod = sys.modules[pipeline_cls.__module__]
missing_mod = [a for a in _REQUIRED_MODULE_ATTRS if not hasattr(pipeline_mod, a)]
if missing_mod:
raise RuntimeError(
"ComfyUI-Trellis2-MultiViewRefiner is out of date: "
f"pipeline module {pipeline_cls.__module__} is missing {missing_mod}. "
"Upstream ComfyUI-Trellis2 has changed. "
"Check for an updated version at "
"https://github.com/cuzelac/ComfyUI-Trellis2-MultiViewRefiner"
)
samplers = pipeline_mod.samplers
if not hasattr(samplers, 'FlowEulerMultiViewGuidanceIntervalSampler'):
raise RuntimeError(
"ComfyUI-Trellis2-MultiViewRefiner is out of date: "
"'FlowEulerMultiViewGuidanceIntervalSampler' not found in "
"ComfyUI-Trellis2's samplers module."
)
if not hasattr(pipeline_cls, 'sample_mesh_slat_multiview'):
pipeline_cls.sample_mesh_slat_multiview = _sample_mesh_slat_multiview
if not hasattr(pipeline_cls, 'refine_mesh_multiview'):
pipeline_cls.refine_mesh_multiview = _refine_mesh_multiview
_patched_classes.add(pipeline_cls)
logger.info(
"ComfyUI-Trellis2-MultiViewRefiner: patched %s with "
"sample_mesh_slat_multiview / refine_mesh_multiview",
pipeline_cls.__name__,
)
# ---------------------------------------------------------------------------
# Node
# ---------------------------------------------------------------------------
class Trellis2MeshRefinerMultiView:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pipeline": ("TRELLIS2PIPELINE",),
"trimesh": ("TRIMESH",),
"front_image": ("IMAGE",),
"seed": ("INT", {"default": 12345, "min": 0, "max": 0x7fffffff}),
"resolution": ([512, 1024, 1536], {"default": 1024}),
"shape_steps": ("INT", {"default": 12, "min": 1, "max": 100}),
"shape_guidance_strength": ("FLOAT", {"default": 6.50, "min": 0.00, "max": 99.99, "step": 0.01}),
"shape_guidance_rescale": ("FLOAT", {"default": 0.05, "min": 0.00, "max": 1.00, "step": 0.01}),
"shape_rescale_t": ("FLOAT", {"default": 4.00, "min": 0.00, "max": 9.99, "step": 0.01}),
"texture_steps": ("INT", {"default": 12, "min": 1, "max": 100}),
"texture_guidance_strength": ("FLOAT", {"default": 3.00, "min": 0.00, "max": 99.99, "step": 0.01}),
"texture_guidance_rescale": ("FLOAT", {"default": 0.20, "min": 0.00, "max": 1.00, "step": 0.01}),
"texture_rescale_t": ("FLOAT", {"default": 3.00, "min": 0.00, "max": 9.99, "step": 0.01}),
"max_num_tokens": ("INT", {"default": 999999, "min": 0, "max": 999999}),
"generate_texture_slat": ("BOOLEAN", {"default": True}),
"downsampling": ([16, 32, 64], {"default": 16}),
"shape_guidance_interval_start": ("FLOAT", {"default": 0.10, "min": 0.00, "max": 1.00, "step": 0.01}),
"shape_guidance_interval_end": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 1.00, "step": 0.01}),
"texture_guidance_interval_start": ("FLOAT", {"default": 0.00, "min": 0.00, "max": 1.00, "step": 0.01}),
"texture_guidance_interval_end": ("FLOAT", {"default": 0.90, "min": 0.00, "max": 1.00, "step": 0.01}),
"use_tiled_decoder": ("BOOLEAN", {"default": True}),
"front_axis": (["z", "x"], {"default": "z"}),
"blend_temperature": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
"sampler": (["euler", "heun", "rk4", "rk5"], {"default": "euler"}),
},
"optional": {
"back_image": ("IMAGE",),
"left_image": ("IMAGE",),
"right_image": ("IMAGE",),
},
}
RETURN_TYPES = ("MESHWITHVOXEL", "BVH",)
RETURN_NAMES = ("mesh", "bvh",)
FUNCTION = "process"
CATEGORY = "Trellis2Wrapper"
OUTPUT_NODE = True
def process(
self,
pipeline, trimesh, front_image, seed, resolution,
shape_steps, shape_guidance_strength, shape_guidance_rescale, shape_rescale_t,
texture_steps, texture_guidance_strength, texture_guidance_rescale, texture_rescale_t,
max_num_tokens, generate_texture_slat, downsampling,
shape_guidance_interval_start, shape_guidance_interval_end,
texture_guidance_interval_start, texture_guidance_interval_end,
use_tiled_decoder, front_axis, blend_temperature, sampler,
back_image=None, left_image=None, right_image=None,
):
_ensure_patched(type(pipeline))
reset_cuda()
front_pil = tensor2pil(front_image)
back_pil = tensor2pil(back_image) if back_image is not None else None
left_pil = tensor2pil(left_image) if left_image is not None else None
right_pil = tensor2pil(right_image) if right_image is not None else None
shape_slat_sampler_params = {
"steps": shape_steps,
"guidance_strength": shape_guidance_strength,
"guidance_rescale": shape_guidance_rescale,
"guidance_interval": [shape_guidance_interval_start, shape_guidance_interval_end],
"rescale_t": shape_rescale_t,
}
tex_slat_sampler_params = {
"steps": texture_steps,
"guidance_strength": texture_guidance_strength,
"guidance_rescale": texture_guidance_rescale,
"guidance_interval": [texture_guidance_interval_start, texture_guidance_interval_end],
"rescale_t": texture_rescale_t,
}
mesh = pipeline.refine_mesh_multiview(
mesh=trimesh,
front=front_pil,
back=back_pil,
left=left_pil,
right=right_pil,
seed=seed,
shape_slat_sampler_params=shape_slat_sampler_params,
tex_slat_sampler_params=tex_slat_sampler_params,
resolution=resolution,
max_num_tokens=max_num_tokens,
generate_texture_slat=generate_texture_slat,
downsampling=downsampling,
use_tiled=use_tiled_decoder,
front_axis=front_axis,
blend_temperature=blend_temperature,
sampler=sampler,
)[0]
vertices = mesh.vertices.cuda()
faces = mesh.faces.cuda()
if generate_texture_slat:
import cumesh as CuMesh
logger.info("Building BVH for current mesh...")
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
bvh.vertices = vertices.detach().clone()
bvh.faces = faces.detach().clone()
else:
logger.info("Not building BVH, only used for texturing")
bvh = None
return (mesh, bvh,)
# ---------------------------------------------------------------------------
# ComfyUI registration
# ---------------------------------------------------------------------------
NODE_CLASS_MAPPINGS = {
"Trellis2MeshRefinerMultiView": Trellis2MeshRefinerMultiView,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Trellis2MeshRefinerMultiView": "Trellis2 - Mesh Refiner Multi-View",
}
Reference in New Issue View Git Blame Copy Permalink