diff --git a/README.md b/README.md index 7aec3da..57e8eea 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ | Date | Description | | --- | --- | +| **2026-03-04** | Sparse Structure Resolution supported up to 128
Experimental for "cascade" pipelines only
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
Added "sdpa" and "flash_attn_3" for the backend | | **2026-02-21** | Fixed "Vertical lines" bug | diff --git a/nodes.py b/nodes.py index a5ac339..20df00f 100644 --- a/nodes.py +++ b/nodes.py @@ -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)", } diff --git a/pyproject.toml b/pyproject.toml index 6a8f6cd..a7fa8d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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) diff --git a/trellis2/models/sparse_structure_vae.py b/trellis2/models/sparse_structure_vae.py index 36714a9..5157882 100644 --- a/trellis2/models/sparse_structure_vae.py +++ b/trellis2/models/sparse_structure_vae.py @@ -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, diff --git a/trellis2/modules/image_feature_extractor.py b/trellis2/modules/image_feature_extractor.py index c3cb515..bf152c8 100644 --- a/trellis2/modules/image_feature_extractor.py +++ b/trellis2/modules/image_feature_extractor.py @@ -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() diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index 0a36b84..1586108 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -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]