styletransferandClaude Opus 4.5
5d0cc8aa0a
Features: - 8 ComfyUI nodes for TRELLIS2 video-to-3D processing - Onion skinning with configurable frame interval - Multiple rendering modes (PBR, wireframe, cel, normal, depth) - 4 lighting presets (studio, outdoor, dramatic, flat) - Post-processing effects (bloom, DoF, motion blur, outlines) - Progress bars and auto model download - 5 example workflows 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
315 lines
9.1 KiB
Python
315 lines
9.1 KiB
Python
"""Mesh I/O utilities for TRELLIS2 meshes."""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Try to import trimesh for mesh I/O
|
|
try:
|
|
import trimesh
|
|
TRIMESH_AVAILABLE = True
|
|
except ImportError:
|
|
TRIMESH_AVAILABLE = False
|
|
logger.warning("trimesh not available. Install with: pip install trimesh")
|
|
|
|
|
|
def load_mesh(
|
|
path: Union[str, Path],
|
|
device: str = "cuda"
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Load a mesh from file and convert to tensor format.
|
|
|
|
Args:
|
|
path: Path to mesh file (supports .obj, .ply, .glb, .gltf)
|
|
device: Device for output tensors
|
|
|
|
Returns:
|
|
Dictionary containing mesh data:
|
|
- vertices: [N, 3] tensor of vertex positions
|
|
- faces: [F, 3] tensor of face indices
|
|
- normals: [N, 3] tensor of vertex normals (if available)
|
|
- uvs: [N, 2] tensor of texture coordinates (if available)
|
|
- colors: [N, 3/4] tensor of vertex colors (if available)
|
|
- textures: dict of texture images (if available)
|
|
"""
|
|
if not TRIMESH_AVAILABLE:
|
|
raise ImportError("trimesh is required for mesh loading. Install with: pip install trimesh")
|
|
|
|
path = Path(path)
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Mesh file not found: {path}")
|
|
|
|
logger.info(f"Loading mesh from {path}")
|
|
|
|
# Load mesh
|
|
mesh = trimesh.load(str(path), force='mesh')
|
|
|
|
result = {}
|
|
|
|
# Vertices
|
|
vertices = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
|
|
result["vertices"] = vertices
|
|
|
|
# Faces
|
|
faces = torch.tensor(mesh.faces, dtype=torch.int64, device=device)
|
|
result["faces"] = faces
|
|
|
|
# Normals
|
|
if hasattr(mesh, 'vertex_normals') and mesh.vertex_normals is not None:
|
|
normals = torch.tensor(mesh.vertex_normals, dtype=torch.float32, device=device)
|
|
result["normals"] = normals
|
|
|
|
# UVs
|
|
if hasattr(mesh.visual, 'uv') and mesh.visual.uv is not None:
|
|
uvs = torch.tensor(mesh.visual.uv, dtype=torch.float32, device=device)
|
|
result["uvs"] = uvs
|
|
|
|
# Vertex colors
|
|
if hasattr(mesh.visual, 'vertex_colors') and mesh.visual.vertex_colors is not None:
|
|
colors = torch.tensor(
|
|
mesh.visual.vertex_colors[:, :3] / 255.0,
|
|
dtype=torch.float32,
|
|
device=device
|
|
)
|
|
result["colors"] = colors
|
|
|
|
# Textures (if PBR material)
|
|
if hasattr(mesh.visual, 'material'):
|
|
material = mesh.visual.material
|
|
textures = {}
|
|
|
|
if hasattr(material, 'baseColorTexture') and material.baseColorTexture is not None:
|
|
textures["albedo"] = np.array(material.baseColorTexture)
|
|
|
|
if hasattr(material, 'metallicRoughnessTexture') and material.metallicRoughnessTexture is not None:
|
|
textures["metallic_roughness"] = np.array(material.metallicRoughnessTexture)
|
|
|
|
if hasattr(material, 'normalTexture') and material.normalTexture is not None:
|
|
textures["normal"] = np.array(material.normalTexture)
|
|
|
|
if textures:
|
|
result["textures"] = textures
|
|
|
|
# Store original mesh reference for advanced operations
|
|
result["_trimesh"] = mesh
|
|
|
|
logger.info(
|
|
f"Loaded mesh: {len(vertices)} vertices, {len(faces)} faces"
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def save_mesh(
|
|
mesh_data: Dict[str, Any],
|
|
path: Union[str, Path],
|
|
file_format: Optional[str] = None
|
|
) -> None:
|
|
"""
|
|
Save mesh data to file.
|
|
|
|
Args:
|
|
mesh_data: Dictionary containing mesh tensors
|
|
path: Output file path
|
|
file_format: Optional format override (obj, ply, glb, gltf)
|
|
"""
|
|
if not TRIMESH_AVAILABLE:
|
|
raise ImportError("trimesh is required for mesh saving. Install with: pip install trimesh")
|
|
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# If we have the original trimesh, use it
|
|
if "_trimesh" in mesh_data:
|
|
mesh = mesh_data["_trimesh"]
|
|
else:
|
|
# Reconstruct from tensors
|
|
vertices = mesh_data["vertices"].cpu().numpy()
|
|
faces = mesh_data["faces"].cpu().numpy()
|
|
|
|
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
|
|
|
|
if "normals" in mesh_data:
|
|
mesh.vertex_normals = mesh_data["normals"].cpu().numpy()
|
|
|
|
if "colors" in mesh_data:
|
|
colors = mesh_data["colors"].cpu().numpy()
|
|
# Convert to 0-255 range and add alpha
|
|
colors_255 = (colors * 255).astype(np.uint8)
|
|
if colors_255.shape[1] == 3:
|
|
alpha = np.full((len(colors_255), 1), 255, dtype=np.uint8)
|
|
colors_255 = np.concatenate([colors_255, alpha], axis=1)
|
|
mesh.visual.vertex_colors = colors_255
|
|
|
|
# Save
|
|
mesh.export(str(path), file_type=file_format)
|
|
logger.info(f"Saved mesh to {path}")
|
|
|
|
|
|
def mesh_to_tensors(
|
|
mesh: Any,
|
|
device: str = "cuda"
|
|
) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Convert a trimesh/TRELLIS mesh object to tensor dictionary.
|
|
|
|
Args:
|
|
mesh: Mesh object (trimesh.Trimesh or TRELLIS mesh)
|
|
device: Target device for tensors
|
|
|
|
Returns:
|
|
Dictionary of mesh tensors
|
|
"""
|
|
result = {}
|
|
|
|
# Handle different mesh types
|
|
if hasattr(mesh, 'vertices'):
|
|
# trimesh or similar
|
|
if isinstance(mesh.vertices, torch.Tensor):
|
|
result["vertices"] = mesh.vertices.to(device)
|
|
else:
|
|
result["vertices"] = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
|
|
|
|
if hasattr(mesh, 'faces'):
|
|
if isinstance(mesh.faces, torch.Tensor):
|
|
result["faces"] = mesh.faces.to(device)
|
|
else:
|
|
result["faces"] = torch.tensor(mesh.faces, dtype=torch.int64, device=device)
|
|
|
|
# Handle TRELLIS-specific mesh format
|
|
elif hasattr(mesh, 'v_pos'):
|
|
result["vertices"] = mesh.v_pos.to(device)
|
|
if hasattr(mesh, 't_pos_idx'):
|
|
result["faces"] = mesh.t_pos_idx.to(device)
|
|
|
|
# Handle dictionary input
|
|
elif isinstance(mesh, dict):
|
|
for key, value in mesh.items():
|
|
if isinstance(value, torch.Tensor):
|
|
result[key] = value.to(device)
|
|
elif isinstance(value, np.ndarray):
|
|
result[key] = torch.tensor(value, device=device)
|
|
|
|
return result
|
|
|
|
|
|
def get_mesh_info(mesh_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Get information about a mesh.
|
|
|
|
Args:
|
|
mesh_data: Dictionary containing mesh tensors
|
|
|
|
Returns:
|
|
Dictionary with mesh properties
|
|
"""
|
|
info = {}
|
|
|
|
if "vertices" in mesh_data:
|
|
vertices = mesh_data["vertices"]
|
|
info["num_vertices"] = len(vertices)
|
|
info["vertex_min"] = vertices.min(dim=0).values.cpu().tolist()
|
|
info["vertex_max"] = vertices.max(dim=0).values.cpu().tolist()
|
|
info["vertex_center"] = vertices.mean(dim=0).cpu().tolist()
|
|
|
|
# Compute bounding box size
|
|
bbox_min = vertices.min(dim=0).values
|
|
bbox_max = vertices.max(dim=0).values
|
|
bbox_size = bbox_max - bbox_min
|
|
info["bbox_size"] = bbox_size.cpu().tolist()
|
|
|
|
if "faces" in mesh_data:
|
|
info["num_faces"] = len(mesh_data["faces"])
|
|
|
|
if "normals" in mesh_data:
|
|
info["has_normals"] = True
|
|
|
|
if "uvs" in mesh_data:
|
|
info["has_uvs"] = True
|
|
|
|
if "colors" in mesh_data:
|
|
info["has_colors"] = True
|
|
|
|
if "textures" in mesh_data:
|
|
info["texture_types"] = list(mesh_data["textures"].keys())
|
|
|
|
return info
|
|
|
|
|
|
def normalize_mesh(
|
|
mesh_data: Dict[str, torch.Tensor],
|
|
target_scale: float = 1.0
|
|
) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Normalize mesh to fit within a unit cube centered at origin.
|
|
|
|
Args:
|
|
mesh_data: Dictionary containing mesh tensors
|
|
target_scale: Scale factor for the normalized mesh
|
|
|
|
Returns:
|
|
Normalized mesh data
|
|
"""
|
|
vertices = mesh_data["vertices"]
|
|
|
|
# Compute center and scale
|
|
center = vertices.mean(dim=0, keepdim=True)
|
|
vertices_centered = vertices - center
|
|
|
|
max_extent = vertices_centered.abs().max()
|
|
scale = target_scale / max_extent
|
|
|
|
# Apply normalization
|
|
result = mesh_data.copy()
|
|
result["vertices"] = vertices_centered * scale
|
|
|
|
return result
|
|
|
|
|
|
def combine_meshes(
|
|
meshes: list,
|
|
device: str = "cuda"
|
|
) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Combine multiple meshes into a single mesh.
|
|
|
|
Args:
|
|
meshes: List of mesh dictionaries
|
|
device: Target device
|
|
|
|
Returns:
|
|
Combined mesh dictionary
|
|
"""
|
|
all_vertices = []
|
|
all_faces = []
|
|
vertex_offset = 0
|
|
|
|
for mesh in meshes:
|
|
vertices = mesh["vertices"]
|
|
faces = mesh["faces"]
|
|
|
|
all_vertices.append(vertices)
|
|
all_faces.append(faces + vertex_offset)
|
|
|
|
vertex_offset += len(vertices)
|
|
|
|
result = {
|
|
"vertices": torch.cat(all_vertices, dim=0).to(device),
|
|
"faces": torch.cat(all_faces, dim=0).to(device),
|
|
}
|
|
|
|
# Combine other attributes if present
|
|
if all(["normals" in m for m in meshes]):
|
|
result["normals"] = torch.cat([m["normals"] for m in meshes], dim=0).to(device)
|
|
|
|
if all(["colors" in m for m in meshes]):
|
|
result["colors"] = torch.cat([m["colors"] for m in meshes], dim=0).to(device)
|
|
|
|
return result
|