MrForExample-ComfyUI-3D-Pack/mesh_processor/export_utils.py

"""
Utilities for exporting from pipelines to our mesh objects
Direct functions WITHOUT trimesh dependency
"""

import torch
import numpy as np
from .mesh import Mesh
from .fastmesh import FastMesh


def export_to_fastmesh(mesh_output, device=None):
    """
    Export from pipeline to FastMesh 
    
    Args:
        mesh_output: object with mesh_v and mesh_f from pipeline
        device: torch device
    
    Returns:
        FastMesh object or list of FastMesh objects
    """
    
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    if isinstance(mesh_output, list):
        outputs = []
        for mesh in mesh_output:
            if mesh is None:
                outputs.append(None)
            else:
                fastmesh = _convert_single_mesh_to_fastmesh(mesh, device)
                outputs.append(fastmesh)
        return outputs
    else:
        return _convert_single_mesh_to_fastmesh(mesh_output, device)


def export_to_mesh(mesh_output, device=None):
    """
    Export from pipeline to standard Mesh
    
    Args:
        mesh_output: object with mesh_v and mesh_f from pipeline
        device: torch device
    
    Returns:
        Mesh object or list of Mesh objects
    """
    
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    if isinstance(mesh_output, list):
        outputs = []
        for mesh in mesh_output:
            if mesh is None:
                outputs.append(None)
            else:
                standard_mesh = _convert_single_mesh_to_mesh(mesh, device)
                outputs.append(standard_mesh)
        return outputs
    else:
        return _convert_single_mesh_to_mesh(mesh_output, device)


def _convert_single_mesh_to_fastmesh(mesh_obj, device):
    """Convert one mesh object to FastMesh"""
    
    try:
        # Extract vertices and faces
        vertices = mesh_obj.mesh_v
        faces = mesh_obj.mesh_f
        
        # Flip faces (as in original function)
        faces = faces[:, ::-1]
        
        # Convert to numpy if needed
        if hasattr(vertices, 'cpu'):
            vertices = vertices.cpu().numpy()
        if hasattr(faces, 'cpu'):
            faces = faces.cpu().numpy()
        
        vertices = np.array(vertices, dtype=np.float32)
        faces = np.array(faces, dtype=np.int32)
        
        print(f"[export_to_fastmesh] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
        
        fastmesh = FastMesh(device=device)
        fastmesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
        fastmesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
        
        fastmesh.auto_normal()
        
        fastmesh._create_empty_albedo()
        
        print(f"[export_to_fastmesh] FastMesh created successfully")
        return fastmesh
        
    except Exception as e:
        print(f"[export_to_fastmesh] Error: {e}")
        return None


def _convert_single_mesh_to_mesh(mesh_obj, device):
    """Convert one mesh object to standard Mesh"""
    
    try:
        vertices = mesh_obj.mesh_v
        faces = mesh_obj.mesh_f
        
        faces = faces[:, ::-1]
        
        if hasattr(vertices, 'cpu'):
            vertices = vertices.cpu().numpy()
        if hasattr(faces, 'cpu'):
            faces = faces.cpu().numpy()
        
        vertices = np.array(vertices, dtype=np.float32)
        faces = np.array(faces, dtype=np.int32)
        
        print(f"[export_to_mesh] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
        
        mesh = Mesh(device=device)
        mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
        mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
        
        mesh.auto_normal()
        
        mesh._create_empty_albedo_fast()
        
        print(f"[export_to_mesh] Mesh created successfully")
        return mesh
        
    except Exception as e:
        print(f"[export_to_mesh] Error: {e}")
        return None


def export_to_mesh_ultra_fast(mesh_output, device=None):
    """
    Ultra-fast export to standard Mesh WITHOUT any slow operations
    """
    
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    try:
        # Process list or single object
        if isinstance(mesh_output, list):
            return [_convert_single_mesh_to_mesh_ultra_fast(mesh, device) for mesh in mesh_output]
        else:
            return _convert_single_mesh_to_mesh_ultra_fast(mesh_output, device)
    
    except Exception as e:
        print(f"[export_to_mesh_ultra_fast] Error: {e}")
        return None


def _convert_single_mesh_to_mesh_ultra_fast(mesh_obj, device):
    """Convert one mesh object to standard Mesh Ultra-fast"""
    
    try:
        # Extract vertices and faces
        vertices = mesh_obj.mesh_v
        faces = mesh_obj.mesh_f
        
        # Flip faces (as in original function)
        faces = faces[:, ::-1]
        
        # Convert to numpy if needed
        if hasattr(vertices, 'cpu'):
            vertices = vertices.cpu().numpy()
        if hasattr(faces, 'cpu'):
            faces = faces.cpu().numpy()
        
        vertices = np.array(vertices, dtype=np.float32)
        faces = np.array(faces, dtype=np.int32)
        
        print(f"[export_to_mesh_ultra_fast] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
        
        # Create standard Mesh Minimal
        mesh = Mesh(device=device)
        mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
        mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
        
        print(f"[export_to_mesh_ultra_fast] Mesh created instantly")
        return mesh
        
    except Exception as e:
        print(f"[export_to_mesh_ultra_fast] Error: {e}")
        return None