diff --git a/nodes.py b/nodes.py index 10447b9..d6c9d8a 100644 --- a/nodes.py +++ b/nodes.py @@ -291,11 +291,7 @@ class Trellis2SimplifyMesh: elif method=="Meshlib": mesh.simplify_with_meshlib(target = target_face_num) else: - raise Exception("Unknown simplification method") - - mm.soft_empty_cache() - torch.cuda.empty_cache() - gc.collect() + raise Exception("Unknown simplification method") return (mesh,) @@ -425,8 +421,6 @@ class Trellis2PostProcessMesh: mesh.faces = new_faces.to(mesh.device) del cumesh - mm.soft_empty_cache() - torch.cuda.empty_cache() gc.collect() return (mesh,) @@ -443,6 +437,7 @@ class Trellis2UnWrapAndRasterizer: "mesh_cluster_smooth_strength": ("INT",{"default":1}), "texture_size": ("INT",{"default":1024, "min":512, "max":16384}), "texture_alpha_mode": (["OPAQUE","MASK","BLEND"],{"default":"OPAQUE"}), + "double_side_material": ("BOOLEAN",{"default":True}), }, } @@ -452,7 +447,7 @@ class Trellis2UnWrapAndRasterizer: CATEGORY = "Trellis2Wrapper" OUTPUT_NODE = True - def process(self, mesh, mesh_cluster_threshold_cone_half_angle_rad, mesh_cluster_refine_iterations, mesh_cluster_global_iterations, mesh_cluster_smooth_strength, texture_size, texture_alpha_mode): + def process(self, mesh, mesh_cluster_threshold_cone_half_angle_rad, mesh_cluster_refine_iterations, mesh_cluster_global_iterations, mesh_cluster_smooth_strength, texture_size, texture_alpha_mode, double_side_material): aabb = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]] vertices = mesh.vertices @@ -587,7 +582,7 @@ class Trellis2UnWrapAndRasterizer: metallicFactor=1.0, roughnessFactor=1.0, alphaMode=alpha_mode, - #doubleSided=True if not remesh else False, + doubleSided=double_side_material, ) vertices_np = out_vertices.cpu().numpy() @@ -609,8 +604,6 @@ class Trellis2UnWrapAndRasterizer: ) del cumesh - mm.soft_empty_cache() - torch.cuda.empty_cache() gc.collect() return (textured_mesh,) @@ -782,7 +775,6 @@ class Trellis2PostProcessAndUnWrapAndRasterizer: else: resolution = int(dual_contouring_resolution) - print(f"Dual Contouring resolution: {resolution}") # Perform Dual Contouring remeshing (rebuilds topology) cumesh.init(*CuMesh.remeshing.remesh_narrow_band_dc( vertices, faces, @@ -919,8 +911,6 @@ class Trellis2PostProcessAndUnWrapAndRasterizer: ) del cumesh - mm.soft_empty_cache() - torch.cuda.empty_cache() gc.collect() return (textured_mesh,) @@ -1033,8 +1023,6 @@ class Trellis2Remesh: mesh.faces = new_faces.to(mesh.device) del cumesh - mm.soft_empty_cache() - torch.cuda.empty_cache() gc.collect() return (mesh,) diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index 12a50c5..096081e 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -83,11 +83,18 @@ class Trellis2ImageTo3DPipeline(Pipeline): Args: path (str): The path to the model. Can be either local path or a Hugging Face repository. - """ + """ pipeline = super(Trellis2ImageTo3DPipeline, Trellis2ImageTo3DPipeline).from_pretrained(path) new_pipeline = Trellis2ImageTo3DPipeline() new_pipeline.__dict__ = pipeline.__dict__ args = pipeline._pretrained_args + + # if os.name=='nt': + # args['models']['sparse_structure_decoder'] = os.path.join(folder_paths.models_dir,"microsoft","TRELLIS-image-large","ckpts","ss_dec_conv3d_16l8_fp16") + # else: + # args['models']['sparse_structure_decoder'] = os.path.join("models","microsoft","TRELLIS-image-large","ckpts","ss_dec_conv3d_16l8_fp16") + + # print(f"Sparse Structure Decoder: {args['models']['sparse_structure_decoder']}") new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args']) new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params'] @@ -550,11 +557,13 @@ class Trellis2ImageTo3DPipeline(Pipeline): torch.manual_seed(seed) cond_512 = self.get_cond([image], 512) cond_1024 = self.get_cond([image], 1024) if pipeline_type != '512' else None - ss_res = {'512': 32, '1024': 64, '1024_cascade': 32, '1536_cascade': 32}[pipeline_type] + ss_res = {'512': 32, '1024': 32, '1024_cascade': 32, '1536_cascade': 32}[pipeline_type] + coords = self.sample_sparse_structure( cond_512, ss_res, num_samples, sparse_structure_sampler_params - ) + ) + if pipeline_type == '512': shape_slat = self.sample_shape_slat( cond_512, self.models['shape_slat_flow_model_512'],