nodes/visualbruno-ComfyUI-Trellis2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Experimental resolutions for sparse structure up to 128 (only for cascade pipelines)

Browse Source 
Increased DinoV3 Image resolution to 2048
This commit is contained in:
Bruno Fargnoli
2026-03-04 15:24:37 +01:00
parent 7ce26deae4
commit 76531be650
6 changed files with 114 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -14,6 +14,7 @@
| Date | Description |
| --- | --- |
| **2026-03-04** | Sparse Structure Resolution supported up to 128<br>Experimental for "cascade" pipelines only<br>Can increase the details |
| **2026-02-27** | Added the Wheels for Windows Python 3.13, Torch 2.10.0, CUDA 13.1 |
| **2026-02-26** | Added FP8 models<br>Added "sdpa" and "flash_attn_3" for the backend |
| **2026-02-21** | Fixed "Vertical lines" bug |
nodes.py
+66 -7
View File
@@ -422,7 +422,7 @@ class Trellis2MeshWithVoxelGenerator:
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
"max_num_tokens": ("INT",{"default":49152,"min":0,"max":999999}),
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
"sparse_structure_resolution": ("INT", {"default":32,"min":8,"max":128,"step":8}),
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
"generate_texture_slat": ("BOOLEAN", {"default":True}),
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
},
@@ -1224,7 +1224,7 @@ class Trellis2MeshWithVoxelAdvancedGenerator:
"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}),
"max_views": ("INT", {"default": 4, "min": 1, "max": 16}),
"sparse_structure_resolution": ("INT", {"default":32,"min":8,"max":128,"step":8}),
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
"generate_texture_slat": ("BOOLEAN", {"default":True}),
"sparse_structure_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
"sparse_structure_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
@@ -1325,7 +1325,7 @@ class Trellis2MeshWithVoxelMultiViewGenerator:
"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}),
"sparse_structure_resolution": ("INT", {"default":32,"min":8,"max":128,"step":8}),
"sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}),
"generate_texture_slat": ("BOOLEAN", {"default":True}),
"sparse_structure_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
"sparse_structure_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
@@ -2029,7 +2029,7 @@ class Trellis2ReconstructMeshWithQuad:
"remesh_band": ("FLOAT",{"default":1.0}),
"resolution": ([128,256,512,1024,2048],{"default":512}),
"remove_floaters": ("BOOLEAN",{"default":True}),
"remove_inner_faces": ("BOOLEAN",{"default":False}),
"remove_inner_faces": ("BOOLEAN",{"default":False}),
}
}
@@ -2429,6 +2429,10 @@ class Trellis2PostProcess2:
"remove_duplicate_faces": ("BOOLEAN",{"default":True}),
"weld_vertices": ("BOOLEAN",{"default":True}),
"weld_vertices_digits": ("INT",{"default":4,"min":1,"max":8}),
"smooth": ("BOOLEAN",{"default":False}),
"smooth_iterations": ("INT",{"default":10,"min":1,"max":99,"step":1}),
"subdivide": ("BOOLEAN",{"default":False}),
"subdivide_iterations": ("INT",{"default":1,"min":1,"max":10}),
},
}
@@ -2438,7 +2442,7 @@ class Trellis2PostProcess2:
CATEGORY = "Trellis2Wrapper"
OUTPUT_NODE = True
def process(self, mesh, fill_holes, fix_normals, fix_face_orientation, remove_duplicate_faces, weld_vertices, weld_vertices_digits,):
def process(self, mesh, fill_holes, fix_normals, fix_face_orientation, remove_duplicate_faces, weld_vertices, weld_vertices_digits,smooth,smooth_iterations,subdivide,subdivide_iterations):
mesh_copy = copy.deepcopy(mesh)
vertices_np = mesh_copy.vertices.cpu().numpy()
@@ -2480,8 +2484,18 @@ class Trellis2PostProcess2:
nb_faces_removed = faces_count - new_faces_count
print(f"Weld Vertices: Removed {nb_vertices_removed} vertices / {nb_faces_removed} faces")
if smooth:
print('Smoothing ...')
Trimesh.smoothing.filter_taubin(trimesh, lamb=0.5, nu=-0.53, iterations=smooth_iterations)
if subdivide:
print('Subdividing ...')
trimesh = trimesh.subdivide_loop(iterations=subdivide_iterations)
new_vertices = torch.from_numpy(trimesh.vertices).float()
new_faces = torch.from_numpy(trimesh.faces).int()
new_faces = torch.from_numpy(trimesh.faces).int()
print(f"After postprocessing: {len(new_faces)} faces")
mesh_copy.vertices = new_vertices.to(mesh_copy.device)
mesh_copy.faces = new_faces.to(mesh_copy.device)
@@ -2736,7 +2750,7 @@ class Trellis2RemeshWithQuad:
"remesh_project": ("FLOAT",{"default":0.0}),
"dual_contouring_resolution": (["Auto","128","256","512","1024","2048"],{"default":"Auto"}),
"remove_floaters": ("BOOLEAN",{"default":True}),
"remove_inner_faces": ("BOOLEAN",{"default":True}),
"remove_inner_faces": ("BOOLEAN",{"default":True})
}
}
@@ -3055,6 +3069,49 @@ class Trellis2FillHolesWithCuMesh:
mesh_copy.fill_holes(max_hole_perimeter = max_permieters)
return (mesh_copy,)
class Trellis2LaplacianSmoothingWithOpen3d:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mesh": ("MESHWITHVOXEL",),
"iterations": ("INT",{"default":10, "min":1, "max":100}),
"method": (["Laplacian", "Taubin"],{"default":"Laplacian"}),
},
}
RETURN_TYPES = ("MESHWITHVOXEL", )
RETURN_NAMES = ("mesh", )
FUNCTION = "process"
CATEGORY = "Trellis2Wrapper"
OUTPUT_NODE = True
def process(self, mesh, iterations, method):
import open3d
mesh_copy = copy.deepcopy(mesh)
vertices = mesh_copy.vertices.cpu().numpy()
faces = mesh_copy.faces.cpu().numpy().astype(np.int32)
open3d_mesh = open3d.geometry.TriangleMesh()
open3d_mesh.vertices = open3d.utility.Vector3dVector(vertices)
open3d_mesh.triangles = open3d.utility.Vector3iVector(faces)
if method == "Laplacian":
open3d_mesh = open3d_mesh.filter_smooth_laplacian(number_of_iterations=iterations)
elif method == "Taubin":
open3d_mesh = open3d_mesh.filter_smooth_taubin(number_of_iterations=iterations)
open3d_mesh.compute_vertex_normals()
new_vertices = np.asarray(open3d_mesh.vertices)
new_faces = np.asarray(open3d_mesh.triangles)
mesh_copy.vertices = torch.from_numpy(new_vertices).float().to(mesh_copy.device)
mesh_copy.faces = torch.from_numpy(new_faces).int().to(mesh_copy.device)
return (mesh_copy,)
NODE_CLASS_MAPPINGS = {
"Trellis2LoadModel": Trellis2LoadModel,
@@ -3090,6 +3147,7 @@ NODE_CLASS_MAPPINGS = {
"Trellis2ReconstructMeshWithQuad": Trellis2ReconstructMeshWithQuad,
"Trellis2StringSelector": Trellis2StringSelector,
"Trellis2FillHolesWithCuMesh": Trellis2FillHolesWithCuMesh,
"Trellis2LaplacianSmoothingWithOpen3d": Trellis2LaplacianSmoothingWithOpen3d,
}
@@ -3127,4 +3185,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"Trellis2ReconstructMeshWithQuad": "Trellis2 - Reconstruct Mesh With Quad",
"Trellis2StringSelector": "Trellis2 - String Selector",
"Trellis2FillHolesWithCuMesh": "Trellis2 - Fill Holes with CuMesh",
"Trellis2LaplacianSmoothingWithOpen3d": "Trellis2 - Laplacian Smoothing (using open3d)",
}
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "trellis2"
description = "ComfyUI Wrapper for Microsoft Trellis.2 - Native and Compact Structured Latents for 3D Generation"
version = "1.0.12"
version = "1.0.13"
license = {file = "LICENSE"}
# classifiers = [
# # For OS-independent nodes (works on all operating systems)
trellis2/models/sparse_structure_vae.py
-24
View File
@@ -99,18 +99,6 @@ class UpsampleBlock3d(nn.Module):
class SparseStructureEncoder(nn.Module):
"""
Encoder for Sparse Structure (\mathcal{E}_S in the paper Sec. 3.3).
Args:
in_channels (int): Channels of the input.
latent_channels (int): Channels of the latent representation.
num_res_blocks (int): Number of residual blocks at each resolution.
channels (List[int]): Channels of the encoder blocks.
num_res_blocks_middle (int): Number of residual blocks in the middle.
norm_type (Literal["group", "layer"]): Type of normalization layer.
use_fp16 (bool): Whether to use FP16.
"""
def __init__(
self,
in_channels: int,
@@ -222,18 +210,6 @@ class SparseStructureEncoder(nn.Module):
class SparseStructureDecoder(nn.Module):
"""
Decoder for Sparse Structure (\mathcal{D}_S in the paper Sec. 3.3).
Args:
out_channels (int): Channels of the output.
latent_channels (int): Channels of the latent representation.
num_res_blocks (int): Number of residual blocks at each resolution.
channels (List[int]): Channels of the decoder blocks.
num_res_blocks_middle (int): Number of residual blocks in the middle.
norm_type (Literal["group", "layer"]): Type of normalization layer.
use_fp16 (bool): Whether to use FP16.
"""
def __init__(
self,
out_channels: int,
trellis2/modules/image_feature_extractor.py
+8 -2
View File
@@ -60,7 +60,7 @@ class DinoV3FeatureExtractor:
"""
Feature extractor for DINOv3 models.
"""
def __init__(self, model_name: str, image_size=512):
def __init__(self, model_name: str, image_size=2048):
self.model_name = model_name
self.model = DINOv3ViTModel.from_pretrained(model_name)
self.model.eval()
@@ -106,7 +106,13 @@ class DinoV3FeatureExtractor:
assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
elif isinstance(image, list):
assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
# We resize the images only if they are bigger than self.image_size
image = [
i.resize((self.image_size, self.image_size), Image.LANCZOS)
if max(i.size) > self.image_size else i
for i in image
]
#image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
image = torch.stack(image).cuda()
trellis2/pipelines/trellis2_image_to_3d.py
+38 -11
View File
@@ -497,9 +497,15 @@ class Trellis2ImageTo3DPipeline(Pipeline):
decoded = decoder(z_s)>0
if self.low_vram:
decoder.cpu()
# if resolution != decoded.shape[2]:
# ratio = decoded.shape[2] // resolution
# decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
if resolution != decoded.shape[2]:
ratio = decoded.shape[2] // resolution
decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
if resolution < decoded.shape[2]:
ratio = decoded.shape[2] // resolution
decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
else:
decoded = torch.nn.functional.interpolate(decoded.float(), size=(resolution, resolution, resolution), mode='nearest') > 0.5
coords = torch.argwhere(decoded)[:, [0, 2, 3, 4]].int()
coords = coords.cpu()
@@ -571,6 +577,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
coords: torch.Tensor,
sampler_params: dict = {},
max_num_tokens: int = 49152,
sparse_structure_resolution: int = 32,
) -> SparseTensor:
"""
Sample structured latent with the given conditioning.
@@ -628,23 +635,28 @@ class Trellis2ImageTo3DPipeline(Pipeline):
if not self.keep_models_loaded:
self.unload_shape_slat_decoder()
ratio = (sparse_structure_resolution / 32)
while True:
quant_coords = torch.cat([
hr_coords[:, :1],
((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
((hr_coords[:, 1:] + 0.5) / (lr_resolution * ratio) * (hr_resolution // 16)).int(),
], dim=1)
coords = quant_coords.unique(dim=0)
num_tokens = coords.shape[0]
if num_tokens < max_num_tokens:
if hr_resolution != resolution:
print(f"Due to the limited number of tokens, the resolution is reduced to {hr_resolution}.")
print(f"Num Tokens: {num_tokens}")
break
hr_resolution -= 128
if hr_resolution < 1024 and resolution >= 1024:
print(f"Num Tokens: {num_tokens}")
hr_resolution = 1024
break
if hr_resolution < 512:
print(f"Num Tokens: {num_tokens}")
hr_resolution = 512
break
@@ -1015,7 +1027,8 @@ class Trellis2ImageTo3DPipeline(Pipeline):
self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
512, 1024,
coords, shape_slat_sampler_params,
max_num_tokens
max_num_tokens,
sparse_structure_resolution
)
if pbar is not None:
@@ -1046,7 +1059,8 @@ class Trellis2ImageTo3DPipeline(Pipeline):
self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
512, 2048,
coords, shape_slat_sampler_params,
max_num_tokens
max_num_tokens,
sparse_structure_resolution
)
if pbar is not None:
@@ -1077,7 +1091,8 @@ class Trellis2ImageTo3DPipeline(Pipeline):
self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
512, 4096,
coords, shape_slat_sampler_params,
max_num_tokens
max_num_tokens,
sparse_structure_resolution
)
if pbar is not None:
@@ -1108,7 +1123,8 @@ class Trellis2ImageTo3DPipeline(Pipeline):
self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
512, 1536,
coords, shape_slat_sampler_params,
max_num_tokens
max_num_tokens,
sparse_structure_resolution
)
if pbar is not None:
@@ -1255,6 +1271,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
max_num_tokens,
front_axis=front_axis,
blend_temperature=blend_temperature,
sparse_structure_resolution = sparse_structure_resolution
)
res = 1024
@@ -1273,6 +1290,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
max_num_tokens,
front_axis=front_axis,
blend_temperature=blend_temperature,
sparse_structure_resolution = sparse_structure_resolution
)
res = 1536
@@ -1417,9 +1435,15 @@ class Trellis2ImageTo3DPipeline(Pipeline):
decoder.cpu()
self._cleanup_cuda()
# if resolution != decoded.shape[2]:
# ratio = decoded.shape[2] // resolution
# decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
if resolution != decoded.shape[2]:
ratio = decoded.shape[2] // resolution
decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
if resolution < decoded.shape[2]:
ratio = decoded.shape[2] // resolution
decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
else:
decoded = torch.nn.functional.interpolate(decoded.float(), size=(resolution, resolution, resolution), mode='nearest') > 0.5
# Extract coordinates (N, 4) -> (b, d, h, w)
# argwhere returns (b, c, d, h, w), so we want [0, 2, 3, 4]
@@ -1507,6 +1531,7 @@ class Trellis2ImageTo3DPipeline(Pipeline):
max_num_tokens: int = 49152,
front_axis: str = 'z',
blend_temperature: float = 2.0,
sparse_structure_resolution: int = 32,
) -> SparseTensor:
# LR
if self.low_vram:
@@ -1561,12 +1586,14 @@ class Trellis2ImageTo3DPipeline(Pipeline):
if self.low_vram:
self.models['shape_slat_decoder'].cpu()
self.models['shape_slat_decoder'].low_vram = False
ratio = sparse_structure_resolution / 32
hr_resolution = resolution
while True:
quant_coords = torch.cat([
hr_coords[:, :1],
((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
((hr_coords[:, 1:] + 0.5) / (lr_resolution * ratio) * (hr_resolution // 16)).int(),
], dim=1)
coords = quant_coords.unique(dim=0)
num_tokens = coords.shape[0]