diff --git a/example_workflows/_Example_Inputs_Files/meshes/robot.glb b/example_workflows/_Example_Inputs_Files/meshes/robot.glb new file mode 100644 index 0000000..8ab7f41 Binary files /dev/null and b/example_workflows/_Example_Inputs_Files/meshes/robot.glb differ diff --git a/loadmesh.py b/loadmesh.py new file mode 100644 index 0000000..b1bb3c2 --- /dev/null +++ b/loadmesh.py @@ -0,0 +1,269 @@ +class Preview_3DMesh: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mesh_file_path": ("STRING", {"default": '', "multiline": False}), + }, + } + + OUTPUT_NODE = True + RETURN_TYPES = () + FUNCTION = "preview_mesh" + CATEGORY = "Comfy3D/Visualize" + + def preview_mesh(self, mesh_file_path): + + mesh_folder_path, filename = os.path.split(mesh_file_path) + + if not os.path.isabs(mesh_file_path): + mesh_file_path = os.path.join(comfy_paths.output_directory, mesh_folder_path, filename) + + if not filename.lower().endswith(SUPPORTED_3D_EXTENSIONS): + cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print() + mesh_file_path = "" + + print(f"[Preview_3DMesh] Final mesh path: {mesh_file_path}") + print(f"[Preview_3DMesh] File exists: {os.path.exists(mesh_file_path) if mesh_file_path else False}") + + previews = [ + { + "filepath": mesh_file_path, + } + ] + return {"ui": {"previews": previews}, "result": ()} + +class Load_3D_Mesh: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mesh_file_path": ("STRING", {"default": '', "multiline": False}), + "use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Использовать FastMesh для GLB/GLTF файлов"}), + "resize": ("BOOLEAN", {"default": False},), + "renormal": ("BOOLEAN", {"default": True},), + "retex": ("BOOLEAN", {"default": False},), + "optimizable": ("BOOLEAN", {"default": False},), + "clean": ("BOOLEAN", {"default": False},), + "resize_bound": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1000.0, "step": 0.001}), + }, + } + + RETURN_TYPES = ( + "MESH", + ) + RETURN_NAMES = ( + "mesh", + ) + FUNCTION = "load_mesh" + CATEGORY = "Comfy3D/Import|Export" + + def load_mesh(self, mesh_file_path, use_fastmesh, resize, renormal, retex, optimizable, clean, resize_bound): + import time + mesh = None + + if not os.path.isabs(mesh_file_path): + mesh_file_path = os.path.join(comfy_paths.input_directory, mesh_file_path) + + if os.path.exists(mesh_file_path): + folder, filename = os.path.split(mesh_file_path) + if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS): + with torch.inference_mode(not optimizable): + start_time = time.time() + if use_fastmesh and mesh_file_path.endswith((".glb", ".gltf")): + mesh = FastMesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound) + print(f"[FastMesh] Loading completed in {time.time() - start_time:.3f}s") + else: + mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound) + print(f"[Mesh] Loading completed in {time.time() - start_time:.3f}s") + else: + cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print() + else: + cstr(f"[{self.__class__.__name__}] File {mesh_file_path} does not exist").error.print() + return (mesh, ) + +class Save_3D_Mesh: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mesh": ("MESH",), + "save_path": ("STRING", {"default": 'Mesh_%Y-%m-%d-%M-%S-%f.glb', "multiline": False}), + "use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Использовать FastMesh для быстрого сохранения GLB/PLY/OBJ"}), + }, + } + + OUTPUT_NODE = True + RETURN_TYPES = ( + "STRING", + ) + RETURN_NAMES = ( + "save_path", + ) + FUNCTION = "save_mesh" + CATEGORY = "Comfy3D/Import|Export" + + def save_mesh(self, mesh, save_path, use_fastmesh): + import time + save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3D_EXTENSIONS, self.__class__.__name__) + + if save_path is not None: + start_time = time.time() + + # Используем FastMesh методы для быстрого сохранения + if use_fastmesh and save_path.endswith((".glb", ".ply", ".obj")): + # Если это уже FastMesh - используем напрямую + print(f"[Save_3D_Mesh] FastMesh detected") + if hasattr(mesh, '_write_glb_fast'): + mesh.write(save_path) + save_time = time.time() - start_time + print(f"[FastMesh] Saving completed in {save_time:.3f}s") + else: + # Если обычный Mesh - создаем временный FastMesh для быстрых методов + print(f"[Save_3D_Mesh] REGULAR Mesh detected") + fast_mesh = FastMesh( + v=mesh.v, f=mesh.f, vn=mesh.vn, fn=mesh.fn, + vt=mesh.vt, ft=mesh.ft, vc=mesh.vc, + albedo=mesh.albedo, metallicRoughness=getattr(mesh, 'metallicRoughness', None), + device=mesh.device + ) + fast_mesh.ori_center = getattr(mesh, 'ori_center', 0) + fast_mesh.ori_scale = getattr(mesh, 'ori_scale', 1) + fast_mesh.write(save_path) + save_time = time.time() - start_time + print(f"[FastMesh] Saving completed in {save_time:.3f}s") + else: + # Обычное сохранение через стандартный метод + mesh.write(save_path) + save_time = time.time() - start_time + print(f"[Mesh] Saving completed in {save_time:.3f}s") + + return (save_path, ) + + +class Preview_3D_FastMesh: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mesh_file_path": ("STRING", {"default": '', "multiline": False}), + }, + } + + OUTPUT_NODE = True + RETURN_TYPES = () + FUNCTION = "preview_mesh" + CATEGORY = "Comfy3D/Visualize" + + def preview_mesh(self, mesh_file_path): + + mesh_folder_path, filename = os.path.split(mesh_file_path) + + if not os.path.isabs(mesh_file_path): + mesh_file_path = os.path.join(comfy_paths.output_directory, mesh_folder_path, filename) + + if not filename.lower().endswith(SUPPORTED_3D_EXTENSIONS): + cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print() + mesh_file_path = "" + + print(f"[Preview_3D_FastMesh] Final mesh path: {mesh_file_path}") + print(f"[Preview_3D_FastMesh] File exists: {os.path.exists(mesh_file_path) if mesh_file_path else False}") + + previews = [ + { + "filepath": mesh_file_path, + } + ] + return {"ui": {"previews": previews}, "result": ()} + + +class Load_3D_FastMesh: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mesh_file_path": ("STRING", {"default": '', "multiline": False}), + "use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Использовать FastMesh загрузчик"}), + "resize": ("BOOLEAN", {"default": False},), + "renormal": ("BOOLEAN", {"default": True},), + "retex": ("BOOLEAN", {"default": False},), + "optimizable": ("BOOLEAN", {"default": False},), + "clean": ("BOOLEAN", {"default": False},), + "resize_bound": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1000.0, "step": 0.001}), + }, + } + + RETURN_TYPES = ( + "MESH", + ) + RETURN_NAMES = ( + "mesh", + ) + FUNCTION = "load_mesh" + CATEGORY = "Comfy3D/Import|Export" + + def load_mesh(self, mesh_file_path, use_fastmesh, resize, renormal, retex, optimizable, clean, resize_bound): + import time + mesh = None + + if not os.path.isabs(mesh_file_path): + mesh_file_path = os.path.join(comfy_paths.input_directory, mesh_file_path) + + if os.path.exists(mesh_file_path): + folder, filename = os.path.split(mesh_file_path) + if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS): + with torch.inference_mode(not optimizable): + start_time = time.time() + if use_fastmesh and mesh_file_path.endswith((".glb", ".gltf")): + from mesh_processor.mesh import FastMesh + mesh = FastMesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound) + print(f"[FastMesh] Loading completed in {time.time() - start_time:.3f}s") + else: + mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound) + mesh_type = "FastMesh" if use_fastmesh else "Mesh" + print(f"[{mesh_type}] Loading completed in {time.time() - start_time:.3f}s") + else: + cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print() + else: + cstr(f"[{self.__class__.__name__}] File {mesh_file_path} does not exist").error.print() + return (mesh, ) + + +class Save_3D_FastMesh: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mesh": ("MESH",), + "save_path": ("STRING", {"default": 'FastMesh_%Y-%m-%d-%M-%S-%f.glb', "multiline": False}), + }, + } + + OUTPUT_NODE = True + RETURN_TYPES = ( + "STRING", + ) + RETURN_NAMES = ( + "save_path", + ) + FUNCTION = "save_mesh" + CATEGORY = "Comfy3D/Import|Export" + + def save_mesh(self, mesh, save_path): + import time + save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3D_EXTENSIONS, self.__class__.__name__) + + if save_path is not None: + start_time = time.time() + mesh.write(save_path) + save_time = time.time() - start_time + mesh_type = "FastMesh" if hasattr(mesh, '_create_empty_albedo') else "Mesh" + print(f"[{mesh_type}] Saving completed in {save_time:.3f}s") + + return (save_path, ) \ No newline at end of file diff --git a/mesh_processor/__init__.py b/mesh_processor/__init__.py index 2ba5108..36db804 100644 --- a/mesh_processor/__init__.py +++ b/mesh_processor/__init__.py @@ -34,9 +34,9 @@ from .rigging_ops import ( from .validate import validate_before_save from .mesh import ( Mesh, - FastMesh, PointCloud, ) +from .fastmesh import FastMesh __all__ = [ "GltfDocument", diff --git a/mesh_processor/export_utils.py b/mesh_processor/export_utils.py index b8612ed..9103f4a 100644 --- a/mesh_processor/export_utils.py +++ b/mesh_processor/export_utils.py @@ -5,7 +5,8 @@ Direct functions WITHOUT trimesh dependency import torch import numpy as np -from .mesh import Mesh, FastMesh +from .mesh import Mesh +from .fastmesh import FastMesh def export_to_fastmesh(mesh_output, device=None): diff --git a/mesh_processor/fastmesh.py b/mesh_processor/fastmesh.py new file mode 100644 index 0000000..f9ea58d --- /dev/null +++ b/mesh_processor/fastmesh.py @@ -0,0 +1,629 @@ +import os +import cv2 +from typing import NamedTuple, Optional +import torch +from torch import Tensor +import trimesh +import numpy as np + +from kiui.op import safe_normalize, dot +from kiui.typing import * + +from .io_gltf import load_gltf_or_glb, get_binary_data, save_glb, GltfDocument +from .mesh_ops import get_all_meshes_triangles +from .accessors import access_data, append_accessor_and_bufferview +from .mesh import Mesh + +class FastMesh: + """ + Fast implementation of mesh based on FastGLB without dependency on trimesh. + Uses direct loading of GLB/GLTF through mesh_processor. + """ + + def __init__( + self, + v: Optional[Tensor] = None, + f: Optional[Tensor] = None, + vn: Optional[Tensor] = None, + fn: Optional[Tensor] = None, + vt: Optional[Tensor] = None, + ft: Optional[Tensor] = None, + vc: Optional[Tensor] = None, + albedo: Optional[Tensor] = None, + metallicRoughness: Optional[Tensor] = None, + device: Optional[torch.device] = None, + ): + """Initialize FastMesh with the same parameters as Mesh""" + self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.v = v + self.vn = vn + self.vt = vt + self.f = f + self.fn = fn + self.ft = ft + self.vc = vc + self.albedo = albedo + self.metallicRoughness = metallicRoughness + self.ori_center = 0 + self.ori_scale = 1 + self._lazy_texture_doc = None + + @classmethod + def load(cls, path, resize=True, renormal=True, retex=False, clean=False, bound=0.5, front_dir='+z', lazy_texture=True, **kwargs): + """Loading mesh through FastGLB approach""" + + device = kwargs.get('device') + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if path.endswith(".obj"): + mesh = cls._load_obj_fast(path, device=device, **kwargs) + elif path.endswith((".glb", ".gltf")): + mesh = cls._load_gltf_fast(path, device=device, lazy_texture=lazy_texture, **kwargs) + else: + standard_mesh = Mesh.load(path, resize=False, renormal=False, retex=False, clean=False, **kwargs) + mesh = cls._from_standard_mesh(standard_mesh) + + if mesh is None: + return None + + if clean: + mesh._clean_mesh() + + print(f"[FastMesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}") + + if resize: + mesh.auto_size(bound=bound) + + if renormal or mesh.vn is None: + mesh.auto_normal() + print(f"[FastMesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}") + + print(f"[FastMesh] UV check: retex={retex}, albedo={mesh.albedo is not None}, vt={mesh.vt is not None}") + if retex or (mesh.albedo is not None and mesh.vt is None): + print(f"[FastMesh] Starting UV generation for {mesh.v.shape[0]} vertices...") + mesh.auto_uv(cache_path=path) + print(f"[FastMesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}") + + print(f"[FastMesh] Checking front_dir: {front_dir}") + if front_dir != "+z": + print(f"[FastMesh] Applying front_dir rotation...") + mesh._apply_front_dir_rotation(front_dir) + + print(f"[FastMesh] Loading completed successfully!") + + return mesh + + @classmethod + def _load_gltf_fast(cls, path, device=None, lazy_texture=True, **kwargs): + """Fast loading of GLB/GLTF through mesh_processor""" + + try: + print(f"[FastMesh] Loading GLB/GLTF: {path}") + + # Load GLTF document + + doc = load_gltf_or_glb(path) + print(f"[FastMesh] GLTF document: {len(doc.meshes())} meshes") + + # Extract geometry + vertices, faces, mesh_groups = get_all_meshes_triangles(doc, transform_to_global=True) + + if len(vertices) == 0 or len(faces) == 0: + print(f"[FastMesh] No geometry in {path}") + return None + + mesh = cls(device=device) + mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) + mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) + + # Extract UV, normals, colors from all meshes + uvs, normals, colors = cls._extract_all_mesh_attributes(doc, mesh_groups, device) + if uvs is not None: + mesh.vt = uvs + mesh.ft = mesh.f # Same topology + print(f"[FastMesh] UV: {mesh.vt.shape}") + if normals is not None: + mesh.vn = normals + mesh.fn = mesh.f + print(f"[FastMesh] Normals: {mesh.vn.shape}") + if colors is not None: + mesh.vc = colors + print(f"[FastMesh] Colors: {mesh.vc.shape}") + if not lazy_texture: + cls._extract_gltf_textures(mesh, doc, device) + else: + mesh._lazy_texture_doc = doc + + print(f"[FastMesh] Loaded: {len(vertices)} vertices, {len(faces)} faces") + return mesh + + except Exception as e: + print(f"[FastMesh] Error loading GLB/GLTF {path}: {e}") + return None + + @classmethod + def _extract_all_mesh_attributes(cls, doc, mesh_groups, device): + """Extract UV, normals, colors from all meshes in GLTF""" + + all_uvs = [] + all_normals = [] + all_colors = [] + + for mesh_idx, gltf_mesh in enumerate(doc.meshes()): + for primitive_idx, primitive in enumerate(gltf_mesh.get("primitives", [])): + attributes = primitive.get("attributes", {}) + + # Find corresponding mesh group + mesh_group = None + for group in mesh_groups: + if group.get('mesh_idx') == mesh_idx and group.get('primitive_idx') == primitive_idx: + mesh_group = group + break + + if mesh_group is None: + continue + + vertex_offset = mesh_group['vertex_offset'] + vertex_count = mesh_group['vertex_count'] + + # Extract UV coordinates + if "TEXCOORD_0" in attributes: + try: + uv_accessor_idx = int(attributes["TEXCOORD_0"]) + uv_data = access_data(doc, uv_accessor_idx).astype(np.float32) + uv_data[:, 1] = 1.0 - uv_data[:, 1] # Flip V + all_uvs.append(uv_data) + print(f"[FastMesh] Mesh {mesh_idx}.{primitive_idx} UV: {uv_data.shape}") + except Exception as e: + print(f"[FastMesh] Error extracting UV from mesh {mesh_idx}.{primitive_idx}: {e}") + all_uvs.append(None) + else: + all_uvs.append(None) + + # Extract normals + if "NORMAL" in attributes: + try: + normal_accessor_idx = int(attributes["NORMAL"]) + normal_data = access_data(doc, normal_accessor_idx).astype(np.float32) + all_normals.append(normal_data) + print(f"[FastMesh] Mesh {mesh_idx}.{primitive_idx} Normals: {normal_data.shape}") + except Exception as e: + print(f"[FastMesh] Error extracting normals from mesh {mesh_idx}.{primitive_idx}: {e}") + all_normals.append(None) + else: + all_normals.append(None) + + # Extract vertex colors + if "COLOR_0" in attributes: + try: + color_accessor_idx = int(attributes["COLOR_0"]) + color_data = access_data(doc, color_accessor_idx).astype(np.float32) + if color_data.shape[1] > 3: + color_data = color_data[:, :3] + all_colors.append(color_data) + print(f"[FastMesh] Mesh {mesh_idx}.{primitive_idx} Colors: {color_data.shape}") + except Exception as e: + print(f"[FastMesh] Error extracting colors from mesh {mesh_idx}.{primitive_idx}: {e}") + all_colors.append(None) + else: + all_colors.append(None) + + # Concatenate all data + final_uvs = None + final_normals = None + final_colors = None + + if any(uv is not None for uv in all_uvs): + # Pad missing UV data with zeros + padded_uvs = [] + for i, uv in enumerate(all_uvs): + if uv is not None: + padded_uvs.append(uv) + else: + # Create dummy UV coordinates + vertex_count = mesh_groups[i]['vertex_count'] if i < len(mesh_groups) else 1000 + dummy_uv = np.zeros((vertex_count, 2), dtype=np.float32) + padded_uvs.append(dummy_uv) + + if padded_uvs: + final_uvs = torch.tensor(np.concatenate(padded_uvs, axis=0), dtype=torch.float32, device=device) + + if any(normal is not None for normal in all_normals): + valid_normals = [n for n in all_normals if n is not None] + if valid_normals: + final_normals = torch.tensor(np.concatenate(valid_normals, axis=0), dtype=torch.float32, device=device) + + if any(color is not None for color in all_colors): + valid_colors = [c for c in all_colors if c is not None] + if valid_colors: + final_colors = torch.tensor(np.concatenate(valid_colors, axis=0), dtype=torch.float32, device=device) + + return final_uvs, final_normals, final_colors + + @classmethod + def _extract_gltf_attributes(cls, mesh, doc, device): + """Extract attributes from GLTF (UV, normals, colors)""" + + if not doc.meshes(): + return + + first_mesh = doc.meshes()[0] + for primitive in first_mesh.get("primitives", []): + attributes = primitive.get("attributes", {}) + + # UV coordinates + if "TEXCOORD_0" in attributes and mesh.vt is None: + try: + uv_accessor_idx = int(attributes["TEXCOORD_0"]) + uv_data = access_data(doc, uv_accessor_idx).astype(np.float32) + uv_data[:, 1] = 1.0 - uv_data[:, 1] # Flip V + mesh.vt = torch.tensor(uv_data, dtype=torch.float32, device=device) + mesh.ft = mesh.f + print(f"[FastMesh] UV: {mesh.vt.shape}") + except Exception as e: + print(f"[FastMesh] Error UV: {e}") + + # Normals + if "NORMAL" in attributes and mesh.vn is None: + try: + normal_accessor_idx = int(attributes["NORMAL"]) + normal_data = access_data(doc, normal_accessor_idx).astype(np.float32) + mesh.vn = torch.tensor(normal_data, dtype=torch.float32, device=device) + mesh.fn = mesh.f + print(f"[FastMesh] Normals: {mesh.vn.shape}") + except Exception as e: + print(f"[FastMesh] Error normals: {e}") + + # Vertex colors + if "COLOR_0" in attributes and mesh.vc is None: + try: + color_accessor_idx = int(attributes["COLOR_0"]) + color_data = access_data(doc, color_accessor_idx).astype(np.float32) + if color_data.shape[1] > 3: + color_data = color_data[:, :3] + mesh.vc = torch.tensor(color_data, dtype=torch.float32, device=device) + print(f"[FastMesh] Colors: {mesh.vc.shape}") + except Exception as e: + print(f"[FastMesh] Error colors: {e}") + + @classmethod + def _extract_gltf_textures(cls, mesh, doc, device): + """Extract textures from GLTF""" + + should_create_empty = True + + try: + if not (doc.materials() and doc.textures() and doc.images()): + mesh._create_empty_albedo() + return + + material = doc.materials()[0] + pbr = material.get("pbrMetallicRoughness", {}) + + # Albedo texture + base_color_texture = pbr.get("baseColorTexture") + if base_color_texture: + albedo_texture = cls._extract_texture_by_index(doc, base_color_texture.get("index", 0)) + if albedo_texture is not None: + albedo_float = albedo_texture.astype(np.float32) / 255.0 + mesh.albedo = torch.tensor(albedo_float, dtype=torch.float32, device=device).contiguous() + should_create_empty = False + print(f"[FastMesh] Albedo: {mesh.albedo.shape}") + + # Metallic-Roughness texture + mr_texture_info = pbr.get("metallicRoughnessTexture") + if mr_texture_info and not should_create_empty: + mr_texture = cls._extract_texture_by_index(doc, mr_texture_info.get("index", 0)) + if mr_texture is not None: + mr_float = mr_texture.astype(np.float32) / 255.0 + mesh.metallicRoughness = torch.tensor(mr_float, dtype=torch.float32, device=device).contiguous() + print(f"[FastMesh] MetallicRoughness: {mesh.metallicRoughness.shape}") + + except Exception as e: + print(f"[FastMesh] Error extracting textures: {e}") + + if should_create_empty: + mesh._create_empty_albedo() + + @classmethod + def _extract_texture_by_index(cls, doc, texture_idx): + """Extract texture by index from GLTF""" + + try: + if texture_idx >= len(doc.textures()): + return None + + texture = doc.textures()[texture_idx] + image_idx = texture.get("source", 0) + + if image_idx >= len(doc.images()): + return None + + image_info = doc.images()[image_idx] + buffer_view_idx = image_info.get("bufferView") + + if buffer_view_idx is None: + return None + + buffer_view = doc.bufferViews()[buffer_view_idx] + buffer_idx = buffer_view.get("buffer", 0) + byte_offset = buffer_view.get("byteOffset", 0) + byte_length = buffer_view.get("byteLength", 0) + + binary_data = get_binary_data(doc, buffer_idx) + image_bytes = binary_data[byte_offset:byte_offset + byte_length] + + # Decode image + image_array = np.frombuffer(image_bytes, dtype=np.uint8) + texture_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + + if texture_image is not None: + texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB) + return texture_image + + except Exception as e: + print(f"[FastMesh] Error extracting texture {texture_idx}: {e}") + + return None + + @classmethod + def _load_obj_fast(cls, path, device=None, **kwargs): + """Fast loading of OBJ (fallback to standard method)""" + standard_mesh = Mesh.load_obj(path, device=device, **kwargs) + return cls._from_standard_mesh(standard_mesh) + + @classmethod + def _from_standard_mesh(cls, standard_mesh): + """Convert standard Mesh to FastMesh""" + if standard_mesh is None: + return None + + return cls( + v=standard_mesh.v, + f=standard_mesh.f, + vn=standard_mesh.vn, + fn=standard_mesh.fn, + vt=standard_mesh.vt, + ft=standard_mesh.ft, + vc=standard_mesh.vc, + albedo=standard_mesh.albedo, + metallicRoughness=standard_mesh.metallicRoughness, + device=standard_mesh.device + ) + + def _create_empty_albedo(self): + """Create empty albedo texture""" + texture = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) + self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device) + print(f"[FastMesh] Empty texture: {self.albedo.shape}") + + def _clean_mesh(self): + """Clean mesh (simplified version)""" + try: + from kiui.mesh_utils import clean_mesh + vertices = self.v.detach().cpu().numpy() + triangles = self.f.detach().cpu().numpy() + vertices, triangles = clean_mesh(vertices, triangles, remesh=False) + self.v = torch.from_numpy(vertices).contiguous().float().to(self.device) + self.f = torch.from_numpy(triangles).contiguous().int().to(self.device) + except Exception as e: + print(f"[FastMesh] Error cleaning: {e}") + + def _apply_front_dir_rotation(self, front_dir): + """Apply rotation for aligning front_dir with +z""" + # Use the same logic as in standard Mesh + if front_dir == "+z": + return + + # Define transformation matrix + if "-z" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=self.device, dtype=torch.float32) + elif "+x" in front_dir: + T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32) + elif "-x" in front_dir: + T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32) + elif "+y" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=self.device, dtype=torch.float32) + elif "-y" in front_dir: + T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=self.device, dtype=torch.float32) + else: + T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + + # Additional rotations + if '1' in front_dir: + T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + elif '2' in front_dir: + T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + elif '3' in front_dir: + T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) + + # Apply transformation + if self.v is not None: + self.v @= T + if self.vn is not None: + self.vn @= T + + # Use methods from standard Mesh + def aabb(self): + """AABB mesh""" + return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values + + @torch.no_grad() + def auto_size(self, bound=0.9): + """Automatic size change""" + vmin, vmax = self.aabb() + self.ori_center = (vmax + vmin) / 2 + self.ori_scale = 2 * bound / torch.max(vmax - vmin).item() + self.v = (self.v - self.ori_center) * self.ori_scale + + def auto_normal(self): + """Automatic calculation of normals""" + i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long() + v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :] + + face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) + + vn = torch.zeros_like(self.v) + vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals) + vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals) + vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals) + + vn = torch.where( + dot(vn, vn) > 1e-20, + vn, + torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device), + ) + vn = safe_normalize(vn) + + self.vn = vn + self.fn = self.f + + def auto_uv(self, cache_path=None, vmap=True): + """Automatic calculation of UV (use method from standard Mesh)""" + # Temporarily delegate to standard Mesh + temp_mesh = Mesh( + v=self.v, f=self.f, vn=self.vn, fn=self.fn, + vt=self.vt, ft=self.ft, vc=self.vc, + albedo=self.albedo, metallicRoughness=self.metallicRoughness, + device=self.device + ) + temp_mesh.auto_uv(cache_path=cache_path, vmap=vmap) + + # Copy result back + self.vt = temp_mesh.vt + self.ft = temp_mesh.ft + if vmap: + self.v = temp_mesh.v + self.f = temp_mesh.f + self.vn = temp_mesh.vn + self.fn = temp_mesh.fn + if temp_mesh.vc is not None: + self.vc = temp_mesh.vc + + def to(self, device): + """Move to another device""" + self.device = device + for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo", "vc", "metallicRoughness"]: + tensor = getattr(self, name) + if tensor is not None: + setattr(self, name, tensor.to(device)) + return self + + def write(self, path): + """Fast saving using mesh_processor for GLB/GLTF files""" + if path.endswith(".glb"): + self._write_glb_fast(path) + elif path.endswith(".gltf"): + # For GLTF, fallback to standard method (less common) + self._write_fallback(path) + elif path.endswith(".ply"): + self._write_ply_fast(path) + elif path.endswith(".obj"): + self._write_obj_fast(path) + else: + # Fallback for unknown formats + self._write_fallback(path) + + def _write_glb_fast(self, path): + """Fast GLB writing using mesh_processor""" + try: + # Prepare data + vertices = self.v.detach().cpu().numpy().astype(np.float32) + faces = self.f.detach().cpu().numpy().astype(np.uint32) + + # Create minimal GLTF structure + gltf_json = { + "asset": {"version": "2.0", "generator": "FastMesh"}, + "scenes": [{"nodes": [0]}], + "nodes": [{"mesh": 0}], + "meshes": [{ + "primitives": [{ + "mode": 4 # TRIANGLES + }] + }] + } + + # Create document with empty buffer + doc = GltfDocument(json_data=gltf_json, binary_buffers={0: np.array([], dtype=np.uint8)}) + + # Add vertex data using helper function + pos_accessor, _ = append_accessor_and_bufferview( + doc, vertices, + component_type=5126, # FLOAT + element_type="VEC3", + target=34962 # ARRAY_BUFFER + ) + + # Add face data using helper function + faces_flat = faces.flatten() + face_component_type = 5123 if vertices.shape[0] <= 65535 else 5125 # UNSIGNED_SHORT or UNSIGNED_INT + face_dtype = np.uint16 if face_component_type == 5123 else np.uint32 + + idx_accessor, _ = append_accessor_and_bufferview( + doc, faces_flat.astype(face_dtype), + component_type=face_component_type, + element_type="SCALAR", + target=34963 # ELEMENT_ARRAY_BUFFER + ) + + # Update primitive with accessor indices + primitive = doc.json()["meshes"][0]["primitives"][0] + primitive["attributes"] = {"POSITION": pos_accessor} + primitive["indices"] = idx_accessor + + # Save to file + with open(path, "wb") as f: + save_glb(doc, f) + + print(f"[FastMesh] Fast GLB saved: {path}") + + except Exception as e: + print(f"[FastMesh] Fast GLB save failed: {e}, using fallback") + self._write_fallback(path) + + def _write_ply_fast(self, path): + """Fast PLY writing""" + vertices = self.v.detach().cpu().numpy() + faces = self.f.detach().cpu().numpy() + + with open(path, 'w') as f: + f.write("ply\n") + f.write("format ascii 1.0\n") + f.write(f"element vertex {len(vertices)}\n") + f.write("property float x\n") + f.write("property float y\n") + f.write("property float z\n") + f.write(f"element face {len(faces)}\n") + f.write("property list uchar int vertex_indices\n") + f.write("end_header\n") + + for v in vertices: + f.write(f"{v[0]} {v[1]} {v[2]}\n") + + for face in faces: + f.write(f"3 {face[0]} {face[1]} {face[2]}\n") + + def _write_obj_fast(self, path): + """Fast OBJ writing""" + vertices = self.v.detach().cpu().numpy() + faces = self.f.detach().cpu().numpy() + + with open(path, 'w') as f: + for v in vertices: + f.write(f"v {v[0]} {v[1]} {v[2]}\n") + + for face in faces: + f.write(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n") + + def _write_fallback(self, path): + """Fallback to standard Mesh write method""" + temp_mesh = Mesh( + v=self.v, f=self.f, vn=self.vn, fn=self.fn, + vt=self.vt, ft=self.ft, vc=self.vc, + albedo=self.albedo, metallicRoughness=self.metallicRoughness, + device=self.device + ) + temp_mesh.ori_center = self.ori_center + temp_mesh.ori_scale = self.ori_scale + return temp_mesh.write(path) \ No newline at end of file diff --git a/mesh_processor/mesh.py b/mesh_processor/mesh.py index c1181f3..f8b61e4 100644 --- a/mesh_processor/mesh.py +++ b/mesh_processor/mesh.py @@ -1130,404 +1130,6 @@ class Mesh: return pcd -class FastMesh: - """ - Fast implementation of mesh based on FastGLB without dependency on trimesh. - Uses direct loading of GLB/GLTF through mesh_processor. - """ - - def __init__( - self, - v: Optional[Tensor] = None, - f: Optional[Tensor] = None, - vn: Optional[Tensor] = None, - fn: Optional[Tensor] = None, - vt: Optional[Tensor] = None, - ft: Optional[Tensor] = None, - vc: Optional[Tensor] = None, - albedo: Optional[Tensor] = None, - metallicRoughness: Optional[Tensor] = None, - device: Optional[torch.device] = None, - ): - """Initialize FastMesh with the same parameters as Mesh""" - self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu") - self.v = v - self.vn = vn - self.vt = vt - self.f = f - self.fn = fn - self.ft = ft - self.vc = vc - self.albedo = albedo - self.metallicRoughness = metallicRoughness - self.ori_center = 0 - self.ori_scale = 1 - - @classmethod - def load(cls, path, resize=True, renormal=True, retex=False, clean=False, bound=0.5, front_dir='+z', **kwargs): - """Loading mesh through FastGLB approach""" - - device = kwargs.get('device') - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - if path.endswith(".obj"): - mesh = cls._load_obj_fast(path, device=device, **kwargs) - elif path.endswith((".glb", ".gltf")): - mesh = cls._load_gltf_fast(path, device=device, **kwargs) - else: - standard_mesh = Mesh.load(path, resize=False, renormal=False, retex=False, clean=False, **kwargs) - mesh = cls._from_standard_mesh(standard_mesh) - - if mesh is None: - return None - - if clean: - mesh._clean_mesh() - - print(f"[FastMesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}") - - if resize: - mesh.auto_size(bound=bound) - - if renormal or mesh.vn is None: - mesh.auto_normal() - print(f"[FastMesh loading] vn: {mesh.vn.shape}, fn: {mesh.fn.shape}") - - if retex or (mesh.albedo is not None and mesh.vt is None): - mesh.auto_uv(cache_path=path) - print(f"[FastMesh loading] vt: {mesh.vt.shape}, ft: {mesh.ft.shape}") - - if front_dir != "+z": - mesh._apply_front_dir_rotation(front_dir) - - return mesh - - @classmethod - def _load_gltf_fast(cls, path, device=None, **kwargs): - """Fast loading of GLB/GLTF through mesh_processor""" - - try: - print(f"[FastMesh] Loading GLB/GLTF: {path}") - - # Load GLTF document - from .io_gltf import load_gltf_or_glb, get_binary_data - from .mesh_ops import get_all_meshes_triangles - from .accessors import access_data - - doc = load_gltf_or_glb(path) - print(f"[FastMesh] GLTF document: {len(doc.meshes())} meshes") - - # Extract geometry - vertices, faces, mesh_groups = get_all_meshes_triangles(doc, transform_to_global=True) - - if len(vertices) == 0 or len(faces) == 0: - print(f"[FastMesh] No geometry in {path}") - return None - - mesh = cls(device=device) - mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device) - mesh.f = torch.tensor(faces, dtype=torch.int32, device=device) - - # Extract additional data - cls._extract_gltf_attributes(mesh, doc, device) - cls._extract_gltf_textures(mesh, doc, device) - - print(f"[FastMesh] Loaded: {len(vertices)} vertices, {len(faces)} faces") - return mesh - - except Exception as e: - print(f"[FastMesh] Error loading GLB/GLTF {path}: {e}") - return None - - @classmethod - def _extract_gltf_attributes(cls, mesh, doc, device): - """Extract attributes from GLTF (UV, normals, colors)""" - - if not doc.meshes(): - return - - first_mesh = doc.meshes()[0] - for primitive in first_mesh.get("primitives", []): - attributes = primitive.get("attributes", {}) - - # UV coordinates - if "TEXCOORD_0" in attributes and mesh.vt is None: - try: - from .accessors import access_data - uv_accessor_idx = int(attributes["TEXCOORD_0"]) - uv_data = access_data(doc, uv_accessor_idx).astype(np.float32) - uv_data[:, 1] = 1.0 - uv_data[:, 1] # Flip V - mesh.vt = torch.tensor(uv_data, dtype=torch.float32, device=device) - mesh.ft = mesh.f - print(f"[FastMesh] UV: {mesh.vt.shape}") - except Exception as e: - print(f"[FastMesh] Error UV: {e}") - - # Normals - if "NORMAL" in attributes and mesh.vn is None: - try: - from .accessors import access_data - normal_accessor_idx = int(attributes["NORMAL"]) - normal_data = access_data(doc, normal_accessor_idx).astype(np.float32) - mesh.vn = torch.tensor(normal_data, dtype=torch.float32, device=device) - mesh.fn = mesh.f - print(f"[FastMesh] Normals: {mesh.vn.shape}") - except Exception as e: - print(f"[FastMesh] Error normals: {e}") - - # Vertex colors - if "COLOR_0" in attributes and mesh.vc is None: - try: - from .accessors import access_data - color_accessor_idx = int(attributes["COLOR_0"]) - color_data = access_data(doc, color_accessor_idx).astype(np.float32) - if color_data.shape[1] > 3: - color_data = color_data[:, :3] - mesh.vc = torch.tensor(color_data, dtype=torch.float32, device=device) - print(f"[FastMesh] Colors: {mesh.vc.shape}") - except Exception as e: - print(f"[FastMesh] Error colors: {e}") - - @classmethod - def _extract_gltf_textures(cls, mesh, doc, device): - """Extract textures from GLTF""" - - should_create_empty = True - - try: - if not (doc.materials() and doc.textures() and doc.images()): - mesh._create_empty_albedo() - return - - material = doc.materials()[0] - pbr = material.get("pbrMetallicRoughness", {}) - - # Albedo texture - base_color_texture = pbr.get("baseColorTexture") - if base_color_texture: - albedo_texture = cls._extract_texture_by_index(doc, base_color_texture.get("index", 0)) - if albedo_texture is not None: - albedo_float = albedo_texture.astype(np.float32) / 255.0 - mesh.albedo = torch.tensor(albedo_float, dtype=torch.float32, device=device).contiguous() - should_create_empty = False - print(f"[FastMesh] Albedo: {mesh.albedo.shape}") - - # Metallic-Roughness texture - mr_texture_info = pbr.get("metallicRoughnessTexture") - if mr_texture_info and not should_create_empty: - mr_texture = cls._extract_texture_by_index(doc, mr_texture_info.get("index", 0)) - if mr_texture is not None: - mr_float = mr_texture.astype(np.float32) / 255.0 - mesh.metallicRoughness = torch.tensor(mr_float, dtype=torch.float32, device=device).contiguous() - print(f"[FastMesh] MetallicRoughness: {mesh.metallicRoughness.shape}") - - except Exception as e: - print(f"[FastMesh] Error extracting textures: {e}") - - if should_create_empty: - mesh._create_empty_albedo() - - @classmethod - def _extract_texture_by_index(cls, doc, texture_idx): - """Extract texture by index from GLTF""" - - try: - from .io_gltf import get_binary_data - - if texture_idx >= len(doc.textures()): - return None - - texture = doc.textures()[texture_idx] - image_idx = texture.get("source", 0) - - if image_idx >= len(doc.images()): - return None - - image_info = doc.images()[image_idx] - buffer_view_idx = image_info.get("bufferView") - - if buffer_view_idx is None: - return None - - buffer_view = doc.bufferViews()[buffer_view_idx] - buffer_idx = buffer_view.get("buffer", 0) - byte_offset = buffer_view.get("byteOffset", 0) - byte_length = buffer_view.get("byteLength", 0) - - binary_data = get_binary_data(doc, buffer_idx) - image_bytes = binary_data[byte_offset:byte_offset + byte_length] - - # Decode image - image_array = np.frombuffer(image_bytes, dtype=np.uint8) - texture_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) - - if texture_image is not None: - texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB) - return texture_image - - except Exception as e: - print(f"[FastMesh] Error extracting texture {texture_idx}: {e}") - - return None - - @classmethod - def _load_obj_fast(cls, path, device=None, **kwargs): - """Fast loading of OBJ (fallback to standard method)""" - standard_mesh = Mesh.load_obj(path, device=device, **kwargs) - return cls._from_standard_mesh(standard_mesh) - - @classmethod - def _from_standard_mesh(cls, standard_mesh): - """Convert standard Mesh to FastMesh""" - if standard_mesh is None: - return None - - return cls( - v=standard_mesh.v, - f=standard_mesh.f, - vn=standard_mesh.vn, - fn=standard_mesh.fn, - vt=standard_mesh.vt, - ft=standard_mesh.ft, - vc=standard_mesh.vc, - albedo=standard_mesh.albedo, - metallicRoughness=standard_mesh.metallicRoughness, - device=standard_mesh.device - ) - - def _create_empty_albedo(self): - """Create empty albedo texture""" - texture = np.ones((1024, 1024, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5]) - self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device) - print(f"[FastMesh] Empty texture: {self.albedo.shape}") - - def _clean_mesh(self): - """Clean mesh (simplified version)""" - try: - from kiui.mesh_utils import clean_mesh - vertices = self.v.detach().cpu().numpy() - triangles = self.f.detach().cpu().numpy() - vertices, triangles = clean_mesh(vertices, triangles, remesh=False) - self.v = torch.from_numpy(vertices).contiguous().float().to(self.device) - self.f = torch.from_numpy(triangles).contiguous().int().to(self.device) - except Exception as e: - print(f"[FastMesh] Error cleaning: {e}") - - def _apply_front_dir_rotation(self, front_dir): - """Apply rotation for aligning front_dir with +z""" - # Use the same logic as in standard Mesh - if front_dir == "+z": - return - - # Define transformation matrix - if "-z" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], device=self.device, dtype=torch.float32) - elif "+x" in front_dir: - T = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32) - elif "-x" in front_dir: - T = torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]], device=self.device, dtype=torch.float32) - elif "+y" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], device=self.device, dtype=torch.float32) - elif "-y" in front_dir: - T = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], device=self.device, dtype=torch.float32) - else: - T = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) - - # Additional rotations - if '1' in front_dir: - T @= torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) - elif '2' in front_dir: - T @= torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) - elif '3' in front_dir: - T @= torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], device=self.device, dtype=torch.float32) - - # Apply transformation - if self.v is not None: - self.v @= T - if self.vn is not None: - self.vn @= T - - # Use methods from standard Mesh - def aabb(self): - """AABB mesh""" - return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values - - @torch.no_grad() - def auto_size(self, bound=0.9): - """Automatic size change""" - vmin, vmax = self.aabb() - self.ori_center = (vmax + vmin) / 2 - self.ori_scale = 2 * bound / torch.max(vmax - vmin).item() - self.v = (self.v - self.ori_center) * self.ori_scale - - def auto_normal(self): - """Automatic calculation of normals""" - i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long() - v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :] - - face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) - - vn = torch.zeros_like(self.v) - vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals) - vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals) - vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals) - - vn = torch.where( - dot(vn, vn) > 1e-20, - vn, - torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device), - ) - vn = safe_normalize(vn) - - self.vn = vn - self.fn = self.f - - def auto_uv(self, cache_path=None, vmap=True): - """Automatic calculation of UV (use method from standard Mesh)""" - # Temporarily delegate to standard Mesh - temp_mesh = Mesh( - v=self.v, f=self.f, vn=self.vn, fn=self.fn, - vt=self.vt, ft=self.ft, vc=self.vc, - albedo=self.albedo, metallicRoughness=self.metallicRoughness, - device=self.device - ) - temp_mesh.auto_uv(cache_path=cache_path, vmap=vmap) - - # Copy result back - self.vt = temp_mesh.vt - self.ft = temp_mesh.ft - if vmap: - self.v = temp_mesh.v - self.f = temp_mesh.f - self.vn = temp_mesh.vn - self.fn = temp_mesh.fn - if temp_mesh.vc is not None: - self.vc = temp_mesh.vc - - def to(self, device): - """Move to another device""" - self.device = device - for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo", "vc", "metallicRoughness"]: - tensor = getattr(self, name) - if tensor is not None: - setattr(self, name, tensor.to(device)) - return self - - def write(self, path): - """Saving to file (use methods from standard Mesh)""" - # Create temporary standard Mesh for saving - temp_mesh = Mesh( - v=self.v, f=self.f, vn=self.vn, fn=self.fn, - vt=self.vt, ft=self.ft, vc=self.vc, - albedo=self.albedo, metallicRoughness=self.metallicRoughness, - device=self.device - ) - temp_mesh.ori_center = self.ori_center - temp_mesh.ori_scale = self.ori_scale - return temp_mesh.write(path) - - class PointCloud(NamedTuple): points: np.array colors: np.array diff --git a/nodes.py b/nodes.py index 5f3f43b..22803ce 100644 --- a/nodes.py +++ b/nodes.py @@ -38,7 +38,8 @@ import trimesh from PIL import Image from mesh_processor.export_utils import export_to_fastmesh, export_to_mesh -from mesh_processor.mesh import FastMesh, Mesh +from mesh_processor.fastmesh import FastMesh +from mesh_processor.mesh import Mesh from mesh_processor.mesh_utils import ( ply_to_points_cloud, get_target_axis_and_scale, @@ -288,6 +289,7 @@ class Load_3D_Mesh: return { "required": { "mesh_file_path": ("STRING", {"default": '', "multiline": False}), + "use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Use FastMesh for GLB/GLTF files"}), "resize": ("BOOLEAN", {"default": False},), "renormal": ("BOOLEAN", {"default": True},), "retex": ("BOOLEAN", {"default": False},), @@ -306,7 +308,7 @@ class Load_3D_Mesh: FUNCTION = "load_mesh" CATEGORY = "Comfy3D/Import|Export" - def load_mesh(self, mesh_file_path, resize, renormal, retex, optimizable, clean, resize_bound): + def load_mesh(self, mesh_file_path, use_fastmesh, resize, renormal, retex, optimizable, clean, resize_bound): mesh = None if not os.path.isabs(mesh_file_path): @@ -316,7 +318,10 @@ class Load_3D_Mesh: folder, filename = os.path.split(mesh_file_path) if filename.lower().endswith(SUPPORTED_3D_EXTENSIONS): with torch.inference_mode(not optimizable): - mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound) + if use_fastmesh and mesh_file_path.endswith((".glb", ".gltf")): + mesh = FastMesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound) + else: + mesh = Mesh.load(mesh_file_path, resize, renormal, retex, clean, resize_bound) else: cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}").error.print() else: @@ -366,6 +371,7 @@ class Save_3D_Mesh: "required": { "mesh": ("MESH",), "save_path": ("STRING", {"default": 'Mesh_%Y-%m-%d-%M-%S-%f.glb', "multiline": False}), + "use_fastmesh": ("BOOLEAN", {"default": True, "tooltip": "Use FastMesh for fast saving GLB"}), }, } @@ -379,11 +385,26 @@ class Save_3D_Mesh: FUNCTION = "save_mesh" CATEGORY = "Comfy3D/Import|Export" - def save_mesh(self, mesh, save_path): + def save_mesh(self, mesh, save_path, use_fastmesh): save_path = parse_save_filename(save_path, comfy_paths.output_directory, SUPPORTED_3D_EXTENSIONS, self.__class__.__name__) if save_path is not None: - mesh.write(save_path) + + if use_fastmesh and save_path.endswith((".glb", ".ply", ".obj")): + if hasattr(mesh, '_write_glb_fast'): + mesh.write(save_path) + else: + fast_mesh = FastMesh( + v=mesh.v, f=mesh.f, vn=mesh.vn, fn=mesh.fn, + vt=mesh.vt, ft=mesh.ft, vc=mesh.vc, + albedo=mesh.albedo, metallicRoughness=getattr(mesh, 'metallicRoughness', None), + device=mesh.device + ) + fast_mesh.ori_center = getattr(mesh, 'ori_center', 0) + fast_mesh.ori_scale = getattr(mesh, 'ori_scale', 1) + fast_mesh.write(save_path) + else: + mesh.write(save_path) return (save_path, )