diff --git a/README.md b/README.md index 128bb72..42d9f94 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,16 @@ --- +**REQUIREMENTS** + +You need to have access to facebook dinov3 models in order to use Trellis.2 [https://huggingface.co/facebook/dinov3-vitl16-pretrain-lvd1689m](https://huggingface.co/facebook/dinov3-vitl16-pretrain-lvd1689m) + +Save the model in ComfyUI models folder under "facebook/dinov3-vitl16-pretrain-lvd1689m" + +So in ComfyUI/models/facebook/dinov3-vitl16-pretrain-lvd1689m + +--- + ## ⚙️ Installation Guide > Tested on **Windows 11** with **Python 3.11** with **Torch = 2.7.0 + cu128**. diff --git a/nodes.py b/nodes.py index 413eb53..d5bbf44 100644 --- a/nodes.py +++ b/nodes.py @@ -9,6 +9,8 @@ import trimesh as Trimesh from tqdm import tqdm import folder_paths +import node_helpers +import hashlib import comfy.model_management as mm from comfy.utils import load_torch_file, ProgressBar, common_upscale @@ -22,7 +24,8 @@ comfy_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) def pil2tensor(image): return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,] -tensor2pil = transforms.ToPILImage() +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) def convert_tensor_images_to_pil(images): pil_array = [] @@ -37,7 +40,8 @@ class Trellis2LoadModel: def INPUT_TYPES(s): return { "required": { - "modelname": ("STRING",["microsoft/TRELLIS.2-4B"]), + "modelname": (["TRELLIS.2-4B"],), + "backend": (["flash_attn","xformers"],{"default":"xformers"}), }, } @@ -45,18 +49,251 @@ class Trellis2LoadModel: RETURN_NAMES = ("pipeline", ) FUNCTION = "process" CATEGORY = "Trellis2Wrapper" + OUTPUT_NODE = True - def process(self, modelname): - pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B") + def process(self, modelname, backend): + os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # Can save GPU memory + os.environ['ATTN_BACKEND'] = backend + + download_path = os.path.join(folder_paths.models_dir,"microsoft") + model_path = os.path.join(download_path, modelname) + + hf_model_name = f"microsoft/{modelname}" + + if not os.path.exists(model_path): + print(f"Downloading model to: {model_path}") + from huggingface_hub import snapshot_download + snapshot_download( + repo_id=hf_model_name, + local_dir=model_path, + local_dir_use_symlinks=False, + ) + + dinov3_model_path = os.path.join(folder_paths.models_dir,"facebook","dinov3-vitl16-pretrain-lvd1689m","model.safetensors") + if not os.path.exists(dinov3_model_path): + raise Exception("Facebook Dinov3 model not found in models/facebook/dinov3-vitl16-pretrain-lvd1689m folder") + + pipeline = Trellis2ImageTo3DPipeline.from_pretrained(model_path) pipeline.cuda() return (pipeline,) + +class Trellis2MeshWithVoxelGenerator: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pipeline": ("TRELLIS2PIPELINE",), + "image": ("IMAGE",), + "seed": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}), + "pipeline_type": (["512","1024","1024_cascade","1536_cascade"],{"default":"1024_cascade"}), + "sparse_structure_steps": ("INT",{"default":12, "min":1, "max":100},), + "shape_steps": ("INT",{"default":12, "min":1, "max":100},), + "texture_steps": ("INT",{"default":12, "min":1, "max":100},), + "max_num_tokens": ("INT",{"default":49152,"min":0,"max":999999}), + }, + } + + RETURN_TYPES = ("MESHWITHVOXEL", ) + RETURN_NAMES = ("mesh", ) + FUNCTION = "process" + CATEGORY = "Trellis2Wrapper" + OUTPUT_NODE = True + + def process(self, pipeline, image, seed, pipeline_type, sparse_structure_steps, shape_steps, texture_steps, max_num_tokens): + image = tensor2pil(image) + + sparse_structure_sampler_params = {"steps":sparse_structure_steps} + shape_slat_sampler_params = {"steps":shape_steps} + tex_slat_sampler_params = {"steps":texture_steps} + + mesh = pipeline.run(image=image, seed=seed, pipeline_type=pipeline_type, sparse_structure_sampler_params = sparse_structure_sampler_params, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, max_num_tokens = max_num_tokens)[0] + + return (mesh,) + +class Trellis2LoadImageWithTransparency: + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + files = folder_paths.filter_files_content_types(files, ["image"]) + return {"required": + {"image": (sorted(files), {"image_upload": True})}, + } + + CATEGORY = "Trellis2Wrapper" + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", ) + RETURN_NAMES = ("image", "mask", "image_with_alpha") + FUNCTION = "load_image" + def load_image(self, image): + image_path = folder_paths.get_annotated_filepath(image) + + img = node_helpers.pillow(Image.open, image_path) + + output_images = [] + output_masks = [] + output_images_ori = [] + w, h = None, None + + excluded_formats = ['MPO'] + + for i in ImageSequence.Iterator(img): + i = node_helpers.pillow(ImageOps.exif_transpose, i) + + output_images_ori.append(pil2tensor(i)) + + if i.mode == 'I': + i = i.point(lambda i: i * (1 / 255)) + image = i.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + elif i.mode == 'P' and 'transparency' in i.info: + mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") + output_images.append(image) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and img.format not in excluded_formats: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) + output_image_ori = torch.cat(output_images_ori, dim=0) + else: + output_image = output_images[0] + output_mask = output_masks[0] + output_image_ori = output_images_ori[0] + + return (output_image, output_mask, output_image_ori) + + @classmethod + def IS_CHANGED(s, image): + image_path = folder_paths.get_annotated_filepath(image) + m = hashlib.sha256() + with open(image_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, image): + if not folder_paths.exists_annotated_filepath(image): + return "Invalid image file: {}".format(image) + + return True + +class Trellis2SimplifyMesh: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mesh": ("MESHWITHVOXEL",), + "target_face_num": ("INT",{"default":1000000,"min":1,"max":16000000}), + "method": (["Cumesh","Meshlib"],{"default":"Meshlib"}), + }, + } + + RETURN_TYPES = ("MESHWITHVOXEL", ) + RETURN_NAMES = ("mesh", ) + FUNCTION = "process" + CATEGORY = "Trellis2Wrapper" + OUTPUT_NODE = True + + def process(self, mesh, target_face_num, method): + if method=="Cumesh": + mesh.simplify_with_cumesh(target = target_face_num) + elif method=="Meshlib": + mesh.simplify_with_meshlib(target = target_face_num) + else: + raise Exception("Unknown simplification method") + + return (mesh,) + +class Trellis2MeshWithVoxelToTrimesh: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mesh": ("MESHWITHVOXEL",), + }, + } + + RETURN_TYPES = ("TRIMESH", ) + RETURN_NAMES = ("trimesh", ) + FUNCTION = "process" + CATEGORY = "Trellis2Wrapper" + OUTPUT_NODE = True + + def process(self, mesh): + trimesh = Trimesh.Trimesh( + vertices=mesh.vertices.cpu().numpy(), + faces=mesh.faces.cpu().numpy(), + process=False + ) + + return (trimesh,) + +class Trellis2ExportMesh: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "trimesh": ("TRIMESH",), + "filename_prefix": ("STRING", {"default": "3D/Hy3D"}), + "file_format": (["glb", "obj", "ply", "stl", "3mf", "dae"],), + }, + "optional": { + "save_file": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("glb_path",) + FUNCTION = "process" + CATEGORY = "Trellis2Wrapper" + OUTPUT_NODE = True + + def process(self, trimesh, filename_prefix, file_format, save_file=True): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory()) + output_glb_path = Path(full_output_folder, f'{filename}_{counter:05}_.{file_format}') + output_glb_path.parent.mkdir(exist_ok=True) + if save_file: + trimesh.export(output_glb_path, file_type=file_format) + relative_path = Path(subfolder) / f'{filename}_{counter:05}_.{file_format}' + else: + temp_file = Path(full_output_folder, f'hy3dtemp_.{file_format}') + trimesh.export(temp_file, file_type=file_format) + relative_path = Path(subfolder) / f'hy3dtemp_.{file_format}' + + return (str(relative_path), ) NODE_CLASS_MAPPINGS = { "Trellis2LoadModel": Trellis2LoadModel, + "Trellis2MeshWithVoxelGenerator": Trellis2MeshWithVoxelGenerator, + "Trellis2LoadImageWithTransparency": Trellis2LoadImageWithTransparency, + "Trellis2SimplifyMesh": Trellis2SimplifyMesh, + "Trellis2MeshWithVoxelToTrimesh": Trellis2MeshWithVoxelToTrimesh, + "Trellis2ExportMesh": Trellis2ExportMesh, } NODE_DISPLAY_NAME_MAPPINGS = { "Trellis2LoadModel": "Trellis2 - LoadModel", + "Trellis2MeshWithVoxelGenerator": "Trellis2 - Mesh With Voxel Generator", + "Trellis2LoadImageWithTransparency": "Trellis2 - Load Image with Transparency", + "Trellis2SimplifyMesh": "Trellis2 - Simplify Mesh", + "Trellis2MeshWithVoxelToTrimesh": "Trellis2 - Mesh With Voxel To Trimesh", + "Trellis2ExportMesh": "Trellis2 - Export Mesh", } diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index 8d7afd5..115983c 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -9,6 +9,9 @@ from ..modules.sparse import SparseTensor from ..modules import image_feature_extractor from ..representations import Mesh, MeshWithVoxel +import folder_paths +import os + class Trellis2ImageTo3DPipeline(Pipeline): """ @@ -92,8 +95,14 @@ class Trellis2ImageTo3DPipeline(Pipeline): new_pipeline.shape_slat_normalization = args['shape_slat_normalization'] new_pipeline.tex_slat_normalization = args['tex_slat_normalization'] + facebook_model_path = os.path.join(folder_paths.models_dir,"facebook","dinov3-vitl16-pretrain-lvd1689m") + args['image_cond_model']['args']['model_name'] = facebook_model_path + new_pipeline.image_cond_model = getattr(image_feature_extractor, args['image_cond_model']['name'])(**args['image_cond_model']['args']) - new_pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args']) + + #new_pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args']) + + new_pipeline.rembg_model = None new_pipeline.low_vram = args.get('low_vram', True) new_pipeline.default_pipeline_type = args.get('default_pipeline_type', '1024_cascade') @@ -487,7 +496,7 @@ class Trellis2ImageTo3DPipeline(Pipeline): sparse_structure_sampler_params: dict = {}, shape_slat_sampler_params: dict = {}, tex_slat_sampler_params: dict = {}, - preprocess_image: bool = True, + preprocess_image: bool = False, return_latent: bool = False, pipeline_type: Optional[str] = None, max_num_tokens: int = 49152, diff --git a/trellis2/representations/mesh/base.py b/trellis2/representations/mesh/base.py index b70e4cc..c0c217b 100644 --- a/trellis2/representations/mesh/base.py +++ b/trellis2/representations/mesh/base.py @@ -4,6 +4,10 @@ from ..voxel import Voxel import cumesh from flex_gemm.ops.grid_sample import grid_sample_3d +import numpy as np +import meshlib.mrmeshnumpy as mrmeshnumpy +import meshlib.mrmeshpy as mrmeshpy + class Mesh: def __init__(self, @@ -68,7 +72,13 @@ class Mesh: self.vertices = new_vertices.to(self.device) self.faces = new_faces.to(self.device) - def simplify(self, target=1000000, verbose: bool=False, options: dict={}): + def simplify_with_cumesh(self, target=1000000, verbose: bool=False, options: dict={}): + current_faces_num = len(self.faces) + print(f'Current Faces Number: {current_faces_num}') + + if current_faces_num