From 3bbf2642e081888d57caafbf1d4a658b9935e7d9 Mon Sep 17 00:00:00 2001 From: Chris Uzelac Date: Wed, 11 Mar 2026 15:40:13 -0700 Subject: [PATCH] Add multi-view Mesh Refiner node with spatial blending Add a new "Mesh Refiner Multi-View" node that refines meshes using per-view spatial blending rather than the existing single-view concat fusion approach. Pipeline changes (trellis2_image_to_3d.py): - sample_mesh_slat_multiview(): combines mesh encoding/upsampling from sample_mesh_slat() with multi-view spatial blending samplers, building per-view conditioning and using FlowMultiViewGuidanceInterval samplers for position-aware blending across views - refine_mesh_multiview(): orchestrates the full multi-view refinement pipeline with named view inputs (front/back/left/right), per-view conditioning, and multi-view sampling for both shape and texture Node changes (nodes.py): - Trellis2MeshRefinerMultiView: ComfyUI node with front_image (required) plus optional back_image, left_image, right_image inputs, with front_axis and blend_temperature controls for spatial blending Co-Authored-By: Claude Opus 4.6 --- nodes.py | 120 +++++++++- trellis2/pipelines/trellis2_image_to_3d.py | 253 ++++++++++++++++++++- 2 files changed, 369 insertions(+), 4 deletions(-) diff --git a/nodes.py b/nodes.py index 1fda815..6a408f4 100644 --- a/nodes.py +++ b/nodes.py @@ -1,4 +1,5 @@ import os +import logging import torch import torchvision.transforms as transforms from PIL import Image, ImageSequence, ImageOps @@ -2438,7 +2439,122 @@ def process(self, pipeline, trimesh, image, seed, resolution, print('Not building BVH, only used for texturing') bvh = None - return (mesh, bvh,) + return (mesh, bvh,) + + +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): + + 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_guidance_interval = [shape_guidance_interval_start, shape_guidance_interval_end] + texture_guidance_interval = [texture_guidance_interval_start, texture_guidance_interval_end] + + shape_slat_sampler_params = {"steps": shape_steps, "guidance_strength": shape_guidance_strength, "guidance_rescale": shape_guidance_rescale, "guidance_interval": shape_guidance_interval, "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, "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: + logging.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: + logging.info('Not building BVH, only used for texturing') + bvh = None + + return (mesh, bvh,) + class Trellis2PostProcess2: @classmethod @@ -3837,6 +3953,7 @@ def process(self, trimesh, target_face_num, thresh, lambda_edge_length, lambda_s "Trellis2LoadMesh": Trellis2LoadMesh, "Trellis2PreProcessImage": Trellis2PreProcessImage, "Trellis2MeshRefiner": Trellis2MeshRefiner, + "Trellis2MeshRefinerMultiView": Trellis2MeshRefinerMultiView, "Trellis2PostProcess2": Trellis2PostProcess2, "Trellis2OvoxelExportToGLB": Trellis2OvoxelExportToGLB, "Trellis2TrimeshToMeshWithVoxel": Trellis2TrimeshToMeshWithVoxel, @@ -3885,6 +4002,7 @@ def process(self, trimesh, target_face_num, thresh, lambda_edge_length, lambda_s "Trellis2LoadMesh": "Trellis2 - Load Mesh", "Trellis2PreProcessImage": "Trellis2 - PreProcess Image", "Trellis2MeshRefiner": "Trellis2 - Mesh Refiner", + "Trellis2MeshRefinerMultiView": "Trellis2 - Mesh Refiner Multi-View", "Trellis2PostProcess2": "Trellis2 - PostProcess Mesh (using Trimesh)", "Trellis2OvoxelExportToGLB": "Trellis2 - Ovoxel Export to GLB", "Trellis2TrimeshToMeshWithVoxel": "Trellis2 - Trimesh to Mesh with Voxel", diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index 3f69059..c6fdc76 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -2692,8 +2692,106 @@ def sample_mesh_slat( cond = self._cond_cpu(cond) self._cleanup_cuda() - return slat, hr_resolution - + return slat, hr_resolution + + def sample_mesh_slat_multiview( + self, + mesh_slat, + conds: dict, + views: list, + flow_model, + resolution: int, + sampler_params: dict = {}, + max_num_tokens: int = 49152, + downsampling: int = 16, + front_axis: str = 'z', + blend_temperature: float = 2.0, + ) -> tuple: + # Upsample + 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 + @torch.no_grad() def refine_mesh( self, @@ -2825,4 +2923,153 @@ def refine_mesh( else: return out_mesh, (shape_slat, None, res) else: - return out_mesh \ No newline at end of file + return out_mesh + + @torch.no_grad() + def refine_mesh_multiview( + self, + mesh: trimesh.Trimesh, + front: Image.Image, + back: Image.Image = None, + left: Image.Image = None, + right: Image.Image = None, + seed: int = 42, + shape_slat_sampler_params: dict = {}, + tex_slat_sampler_params: dict = {}, + resolution: int = 1024, + max_num_tokens: int = 50000, + generate_texture_slat: bool = True, + return_latent: bool = False, + downsampling: int = 16, + use_tiled: bool = True, + front_axis: str = 'z', + blend_temperature: float = 2.0, + sampler: str = 'euler', + ): + self.switch_samplers(sampler) + + mesh = self.preprocess_mesh(mesh) + seed_all(seed) + + # Build views dict + 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()) + + # Build per-view conditioning + 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 \ No newline at end of file