diff --git a/.gitignore b/.gitignore index 000fa71..6803485 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,12 @@ env/ Thumbs.db tmp +__pycache__/ +.venv/ +deps/ +autotune_cache.json +output/ +*.glb +tmp/ +runpod/*.tar.gz +RUNPOD_SESSION.md diff --git a/README_MACOS.md b/README_MACOS.md new file mode 100644 index 0000000..fd93386 --- /dev/null +++ b/README_MACOS.md @@ -0,0 +1,118 @@ +# Pixal3D sur macOS Apple Silicon + +Ce dépôt contient le code officiel Pixal3D et un port d’exécution MPS pour le +Mac M3 Max 36 Go. Le port reprend les backends Metal validés dans +`../trellis2-macos` : `mtldiffrast`, `mtlmesh`, `mtlgemm`, `mtlbvh` et le fork +Apple de `o_voxel`. + +## Installation + +```bash +xcodebuild -downloadComponent MetalToolchain +bash setup_macos.sh +source .venv/bin/activate +``` + +Les poids sont téléchargés à la demande par Hugging Face. Pour les mettre en +cache avant le premier calcul : + +```bash +python scripts/download_models.py +``` + +## Génération iso-qualité CUDA + +Le profil `cuda-parity` conserve la cascade neurale 1536, le volume PBR +1536, une cible d’environ un million de faces et les textures PBR 4096 px. +Le dual-contouring utilise par défaut une grille 512³ : c’est le profil +validé avec marge sur 36 Go, et son résultat passe les contrôles structuraux +face au GLB CUDA de référence. + +```bash +python inference.py \ + --image assets/images/0_img.png \ + --output output/0_cuda_parity.glb \ + --low_vram \ + --resolution 1536 \ + --export-profile cuda-parity +``` + +Avant l’export, le programme sauvegarde automatiquement +`output/0_cuda_parity.decoded.pt`. Ce checkpoint contient le maillage décodé +et son volume PBR sparse, sans les poids des modèles. Si le remeshing ou la +texture 4096 échoue, l’export peut être repris sans relancer la génération : + +```bash +python inference.py \ + --image assets/images/0_img.png \ + --decoded-checkpoint output/0_cuda_parity.decoded.pt \ + --output output/0_cuda_parity.glb \ + --export-profile cuda-parity +``` + +`--remesh-resolution` permet d’expérimenter avec une grille plus dense, mais +la mémoire du simplificateur croît rapidement au-delà de 512. Le profil +historique léger reste accessible avec `--export-profile portable`. + +La pression sur les 36 Go de mémoire unifiée est limitée de quatre façons : + +- DINOv3 et NAF sont partagés entre les quatre conditionneurs ; +- chaque flow/decoder est supprimé après sa dernière étape ; +- le maillage et le volume décodés passent sur CPU avant l’export ; +- le BVH source de 18 millions de triangles est réduit exactement sur des + hiérarchies successives de 250 000 faces ; +- le remesh, le nettoyage/simplification et le bake sont des étapes séparées, + ce qui libère leurs allocations Metal entre elles ; +- l’échantillonnage du volume est découpé en lots ; +- seuls les sommets d’échantillonnage puis les rares texels invalides sont + reprojetés sur la surface source. + +La grille dual-contouring n’est volontairement pas tuilée : des blocs +indépendants créeraient des raccords. C’est le BVH de distance qui est +découpé, puis réduit par minimum global, donc sans fissure aux frontières. + +### Validation de référence + +Le cas `output/inputs/0_img_2048.png`, seed 42, a été comparé au GLB produit +par la branche `main` sur une RTX A5000 RunPod : + +- CUDA : 937 343 faces, 5 arêtes de bord, 99,47 % dans la composante + principale, aire 5,197 ; +- MPS final : 989 941 faces, 7 arêtes de bord, 99,46 % dans la composante + principale, aire 4,944 ; +- les deux fichiers ont une texture 4096², un matériau opaque/simple face et + un alpha p01 de 254. + +Le contrôle automatisé se relance avec : + +```bash +python -m scripts.compare_glb_quality \ + --reference output/pixal3d_main_cuda_a5000_2048_lowvram.glb \ + --candidate output/pixal3d_mps_1536_cuda_parity_final.glb +``` + +Sur ce Mac, la génération neurale 1536 mesurée prend 1 979,85 s et l’export +final intégré 128,65 s, soit environ 35 min 09 s au total. Le même cas avait +pris environ 13 min 25 s sur l’A5000. + +## Interface + +Pour lancer l’interface locale : + +```bash +python app.py --low_vram +``` + +Le CLI utilise les convolutions sparse `flex_gemm`, SDPA sur MPS pour les +longues séquences et le vrai modèle NAF appris. Le noyau d’attention Metal +fusionné reste disponible, mais il n’est pas retenu par défaut : il régresse +fortement vers 40 000 tokens malgré ses bons résultats sur les petites +séquences. Seule l’opération NATTEN CUDA de NAF est remplacée par une +implémentation MPS équivalente, découpée par lignes. Ces chemins ont des tests +numériques FP16/BF16 face à leurs références PyTorch. + +Les ajouts spécifiques sont listés dans `requirements-macos.txt`. Le setup +retient la version récente de `utils3d` requise par MoGe ; la wheel 0.0.2 +indiquée dans la fiche Pixal3D est trop ancienne pour cette API. +`requirements-hfdemo.txt` est réservé au Space Hugging Face et ne doit pas +être utilisé ici. diff --git a/app.py b/app.py index 48479df..7332236 100644 --- a/app.py +++ b/app.py @@ -1,11 +1,9 @@ import os -import subprocess import argparse import math import time import shutil import cv2 -import torch import numpy as np import base64 import io @@ -25,12 +23,24 @@ init_lock = threading.Lock() os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" -os.environ.setdefault("ATTN_BACKEND", "flash_attn") +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +os.environ.setdefault("ATTN_BACKEND", "sdpa") +os.environ.setdefault("SPARSE_ATTN_BACKEND", "sdpa") +os.environ.setdefault("SPARSE_CONV_BACKEND", "none") os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json') -os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1' +from macos_compat import configure + +DEVICE = configure() +import torch -import spaces +try: + import spaces +except ImportError: + class _LocalSpaces: + @staticmethod + def GPU(**_kwargs): + return lambda fn: fn + spaces = _LocalSpaces() from gradio import Server from gradio.data_classes import FileData from fastapi.responses import HTMLResponse @@ -41,6 +51,7 @@ from pixal3d.renderers import EnvMap from pixal3d.utils import render_utils import o_voxel +from o_voxel import postprocess_cpu # ============================================================================ # Constants & Defaults @@ -110,7 +121,7 @@ def build_image_cond_model(config: dict): model.eval() return model -def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME): +def load_moge_model(device=DEVICE, model_name=MOGE_MODEL_NAME): from moge.model.v2 import MoGeModel moge_model = MoGeModel.from_pretrained(model_name).to(device) moge_model.eval() @@ -120,7 +131,7 @@ def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME): pipeline = None moge_model = None envmap = None -LOW_VRAM = os.environ.get("LOW_VRAM", "0") == "1" +LOW_VRAM = os.environ.get("LOW_VRAM", "1") == "1" def init_models(): global pipeline, moge_model, envmap @@ -166,17 +177,17 @@ def init_models(): m = getattr(pipeline, attr, None) if m is not None and getattr(m, 'use_naf_upsample', False): m._load_naf() - pipeline._device = torch.device("cuda") + pipeline._device = torch.device(DEVICE) pipeline.low_vram = True print("[Pipeline] Low-VRAM mode enabled.") else: # Standard mode: all models loaded to GPU at once. pipeline.low_vram = False - pipeline.cuda() - pipeline.image_cond_model_ss.cuda() - pipeline.image_cond_model_shape_512.cuda() - pipeline.image_cond_model_shape_1024.cuda() - pipeline.image_cond_model_tex_1024.cuda() + pipeline.to(DEVICE) + pipeline.image_cond_model_ss.to(DEVICE) + pipeline.image_cond_model_shape_512.to(DEVICE) + pipeline.image_cond_model_shape_1024.to(DEVICE) + pipeline.image_cond_model_tex_1024.to(DEVICE) print("[NAF] Pre-loading NAF upsampler model...") for attr in ['image_cond_model_ss', 'image_cond_model_shape_512', 'image_cond_model_shape_1024', 'image_cond_model_tex_1024']: @@ -190,11 +201,11 @@ def init_models(): moge_model = load_moge_model(device="cpu") print("[MoGe-2] Low-VRAM mode: MoGe stays on CPU, loaded to GPU on-demand.") else: - moge_model = load_moge_model(device="cuda") + moge_model = load_moge_model(device=DEVICE) print("[EnvMap] Loading environment maps...") _base = os.path.dirname(os.path.abspath(__file__)) - _envmap_device = 'cpu' if LOW_VRAM else 'cuda' + _envmap_device = 'cpu' if LOW_VRAM else DEVICE envmap = { 'forest': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/forest.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)), 'sunset': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/sunset.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)), @@ -222,7 +233,7 @@ def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, imag distance_x = f_pixels * xw / x_ndc - yw return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)} -def get_camera_params_wild_moge(image_path, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512): +def get_camera_params_wild_moge(image_path, device=DEVICE, mesh_scale=1.0, extend_pixel=0, image_resolution=512): pil_image = Image.open(image_path).convert("RGB") width, height = pil_image.size image_np = np.array(pil_image).astype(np.float32) / 255.0 @@ -431,7 +442,7 @@ def generate_3d( print(f"[Camera] Using manual FOV: {fov_deg:.2f}° ({camera_angle_x:.4f} rad), distance: {distance:.4f}") else: camera_params = get_camera_params_wild_moge( - temp_processed_path, device="cuda", + temp_processed_path, device=DEVICE, mesh_scale=WILD_MESH_SCALE, extend_pixel=WILD_EXTEND_PIXEL, image_resolution=WILD_IMAGE_RESOLUTION, ) @@ -468,7 +479,7 @@ def generate_3d( far = cam_dist + 10.0 if LOW_VRAM: for v in envmap.values(): - v.image = v.image.cuda() + v.image = v.image.to(DEVICE) if hasattr(v, '_nvdiffrec_envlight'): del v._nvdiffrec_envlight renders = render_utils.render_proj_aligned_video( @@ -514,12 +525,15 @@ def extract_glb_api(state_path: str, decimation_target: int, texture_size: int, mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0] _update_progress("Decoding latent", 1, 1) - glb = o_voxel.postprocess.to_glb( + # Use the portable macOS exporter: the Metal cumesh remesher can stall on + # large decoded meshes. Keep the web demo responsive with a bounded profile. + glb = postprocess_cpu.to_glb( vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs, coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout, grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], - decimation_target=decimation_target, texture_size=texture_size, - remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True, + decimation_target=min(decimation_target, 50000), + texture_size=min(texture_size, 256), + remesh=False, remesh_band=1, remesh_project=0, use_tqdm=True, ) rot = np.array([ [-1, 0, 0, 0], @@ -547,13 +561,7 @@ def extract_glb_api(state_path: str, decimation_target: int, texture_size: int, if args.low_vram: LOW_VRAM = True - # Re-install utils3d as in original app.py - subprocess.run([ - sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", - "https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl" - ], check=True) - # Pre-initialize models before launching the server init_models() - app.launch(show_error=True, share=True) \ No newline at end of file + app.launch(show_error=True, share=False) diff --git a/backends/chunked_mtlbvh.py b/backends/chunked_mtlbvh.py new file mode 100644 index 0000000..5710a2b --- /dev/null +++ b/backends/chunked_mtlbvh.py @@ -0,0 +1,144 @@ +"""Exact-enough Metal BVH queries for very large decoded meshes. + +MtlBVH loses several voxels of accuracy when a single hierarchy is built over +the roughly 18 million tiny triangles emitted by Pixal3D's 1536 decoder. The +same implementation stays within a fraction of a voxel around 250k faces. +This adapter therefore evaluates consecutive source-face chunks and keeps the +closest result for every query point. + +Only one native hierarchy is alive at a time. This trades runtime for bounded +unified-memory pressure and lets the remesher see the complete decoded surface. +""" + +from __future__ import annotations + +from collections.abc import Callable +import gc +from typing import Any + +import torch + + +class ChunkedMtlBVH: + """Minimum-distance reduction across bounded source and query chunks.""" + + def __init__( + self, + bvh_factory: Callable[[torch.Tensor, torch.Tensor], Any], + vertices: torch.Tensor, + faces: torch.Tensor, + *, + source_face_chunk_size: int = 250_000, + query_chunk_size: int = 262_144, + ) -> None: + if source_face_chunk_size <= 8: + raise ValueError("source_face_chunk_size must be greater than 8") + if query_chunk_size <= 0: + raise ValueError("query_chunk_size must be positive") + + self._factory = bvh_factory + self._vertices = vertices.detach().cpu().float().contiguous() + self._faces = faces.detach().cpu().int().contiguous() + self.source_face_chunk_size = int(source_face_chunk_size) + self.query_chunk_size = int(query_chunk_size) + self._single = None + if len(self._faces) <= self.source_face_chunk_size: + self._single = self._factory(self._vertices, self._faces) + + @property + def source_chunks(self) -> int: + return ( + len(self._faces) + self.source_face_chunk_size - 1 + ) // self.source_face_chunk_size + + def _query_one( + self, + bvh: Any, + positions: torch.Tensor, + *, + return_uvw: bool, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + outputs: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]] = [] + for start in range(0, len(positions), self.query_chunk_size): + outputs.append( + bvh.unsigned_distance( + positions[start : start + self.query_chunk_size], + return_uvw=return_uvw, + **kwargs, + ) + ) + if not outputs: + return bvh.unsigned_distance( + positions, + return_uvw=return_uvw, + **kwargs, + ) + distances = torch.cat([output[0] for output in outputs], dim=0) + face_ids = torch.cat([output[1] for output in outputs], dim=0) + uvw = None + if return_uvw: + uvw = torch.cat([output[2] for output in outputs], dim=0) + return distances, face_ids, uvw + + def unsigned_distance( + self, + positions: torch.Tensor, + return_uvw: bool = False, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + positions = positions.detach().cpu().float().contiguous() + if self._single is not None: + return self._query_one( + self._single, + positions, + return_uvw=return_uvw, + **kwargs, + ) + + best_distances: torch.Tensor | None = None + best_face_ids: torch.Tensor | None = None + best_uvw: torch.Tensor | None = None + + for face_start in range( + 0, + len(self._faces), + self.source_face_chunk_size, + ): + face_stop = min( + face_start + self.source_face_chunk_size, + len(self._faces), + ) + source_faces = self._faces[face_start:face_stop].contiguous() + if len(source_faces) <= 8: + continue + bvh = self._factory(self._vertices, source_faces) + distances, local_face_ids, uvw = self._query_one( + bvh, + positions, + return_uvw=return_uvw, + **kwargs, + ) + global_face_ids = local_face_ids + face_start + + if best_distances is None: + best_distances = distances.clone() + best_face_ids = global_face_ids.clone() + if return_uvw: + assert uvw is not None + best_uvw = uvw.clone() + else: + better = distances < best_distances + best_distances[better] = distances[better] + assert best_face_ids is not None + best_face_ids[better] = global_face_ids[better] + if return_uvw: + assert best_uvw is not None and uvw is not None + best_uvw[better] = uvw[better] + + del bvh, source_faces, distances, local_face_ids, global_face_ids, uvw + + gc.collect() + if best_distances is None or best_face_ids is None: + raise RuntimeError("No valid source-face chunk was available") + return best_distances, best_face_ids, best_uvw diff --git a/backends/conv_none.py b/backends/conv_none.py new file mode 100644 index 0000000..766b372 --- /dev/null +++ b/backends/conv_none.py @@ -0,0 +1,133 @@ +""" +Pure-PyTorch sparse 3D convolution backend. + +Implements submanifold sparse convolution by gathering neighbor features, +applying convolution weights via matrix multiply, and scatter-adding results. +No CUDA extensions needed — works on MPS and CPU. + +Slower than flex_gemm/spconv but fully portable. +""" + +import math +import torch +import torch.nn as nn +from .. import SparseTensor + + +def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None): + assert stride == 1 and (padding is None), \ + "Naive implementation only supports submanifold sparse convolution (stride=1, padding=None)" + + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = tuple(kernel_size) if isinstance(kernel_size, (list, tuple)) else (kernel_size,) * 3 + self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride,) * 3 + self.dilation = tuple(dilation) if isinstance(dilation, (list, tuple)) else (dilation,) * 3 + + self.weight = nn.Parameter(torch.empty((out_channels, in_channels, *self.kernel_size))) + if bias: + self.bias = nn.Parameter(torch.empty(out_channels)) + else: + self.register_parameter("bias", None) + + torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + torch.nn.init.uniform_(self.bias, -bound, bound) + + # Match flex_gemm weight layout: (Co, Ci, Kd, Kh, Kw) -> (Co, Kd, Kh, Kw, Ci) + self.weight = nn.Parameter(self.weight.permute(0, 2, 3, 4, 1).contiguous()) + + +def sparse_conv3d_forward(self, x: SparseTensor) -> SparseTensor: + """ + Submanifold sparse 3D convolution via gather-scatter. + + For each active voxel, gather features from its kernel-sized neighborhood + (only where other active voxels exist), multiply by the corresponding + kernel weight, and scatter-add results back. + """ + Co, Kd, Kh, Kw, Ci = self.weight.shape + device = x.feats.device + dtype = x.feats.dtype + + coords = x.coords # [N, 4] (batch_idx, z, y, x) + feats = x.feats # [N, Ci] + N = coords.shape[0] + + # Build neighbor index cache (reused across forward passes for same coords) + cache_key = f'SubMConv3d_naive_neighbor_{Kw}x{Kh}x{Kd}_dilation{self.dilation}' + neighbor_cache = x.get_spatial_cache(cache_key) + + if neighbor_cache is None: + # Build spatial hash: coord tuple -> voxel index + coord_to_idx = {} + coords_cpu = coords.cpu() + for i in range(N): + key = tuple(coords_cpu[i].tolist()) + coord_to_idx[key] = i + + # For each kernel position, find (source, target) voxel pairs + dz, dy, dx = self.dilation + src_indices = [] + tgt_indices = [] + kernel_indices = [] + + for kz in range(Kd): + for ky in range(Kh): + for kx in range(Kw): + oz = (kz - Kd // 2) * dz + oy = (ky - Kh // 2) * dy + ox = (kx - Kw // 2) * dx + k_idx = kz * Kh * Kw + ky * Kw + kx + + for i in range(N): + b, z, y, xc = coords_cpu[i].tolist() + neighbor_key = (b, z + oz, y + oy, xc + ox) + if neighbor_key in coord_to_idx: + j = coord_to_idx[neighbor_key] + src_indices.append(j) + tgt_indices.append(i) + kernel_indices.append(k_idx) + + neighbor_cache = ( + torch.tensor(src_indices, dtype=torch.long, device=device), + torch.tensor(tgt_indices, dtype=torch.long, device=device), + torch.tensor(kernel_indices, dtype=torch.long, device=device), + ) + x.register_spatial_cache(cache_key, neighbor_cache) + + src_idx, tgt_idx, k_idx = neighbor_cache + + # Reshape weight: (Co, Kd, Kh, Kw, Ci) -> (K, Ci, Co) + K_total = Kd * Kh * Kw + w = self.weight.reshape(Co, K_total, Ci).permute(1, 2, 0) # (K, Ci, Co) + + out = torch.zeros(N, Co, device=device, dtype=dtype) + + if len(src_idx) > 0: + # Process each kernel position to keep memory bounded + for k in range(K_total): + mask = (k_idx == k) + if not mask.any(): + continue + s_idx = src_idx[mask] + t_idx = tgt_idx[mask] + src_f = feats[s_idx] # [E_k, Ci] + edge_out = src_f @ w[k] # [E_k, Co] + out.scatter_add_(0, t_idx.unsqueeze(1).expand(-1, Co), edge_out) + + if self.bias is not None: + out = out + self.bias + + return x.replace(out) + + +def sparse_inverse_conv3d_init(self, *args, **kwargs): + raise NotImplementedError("SparseInverseConv3d with naive backend is not implemented") + + +def sparse_inverse_conv3d_forward(self, x: SparseTensor) -> SparseTensor: + raise NotImplementedError("SparseInverseConv3d with naive backend is not implemented") diff --git a/backends/cuda_parity_export.py b/backends/cuda_parity_export.py new file mode 100644 index 0000000..12bcb19 --- /dev/null +++ b/backends/cuda_parity_export.py @@ -0,0 +1,429 @@ +"""CUDA-quality GLB export using the native Apple-Silicon ``o_voxel`` stack. + +Metal's monolithic BVH loses several voxels of accuracy on Pixal3D's roughly +18-million-triangle decoded mesh. The visible symptom is a fragmented, +point-cloud-like GLB even though the neural output is healthy. This exporter +uses a bounded two-stage path instead: + +* query consecutive 250k-face source BVHs while running one continuous + dual-contouring grid (the grid itself is never tiled); +* clean, close and simplify that result to about one million faces; +* project only texture-sampling vertices, plus any remaining invalid texels, + back to the decoded surface through the same accurate chunked BVH; +* bake the 4096px sparse PBR volume without mutating the prepared geometry. + +The 512-cell remesh default is the highest profile validated with comfortable +headroom on a 36GB M3 Max. It retains the 1536 neural cascade and PBR volume; +only the intermediate dual-contouring grid is reduced. +""" + +from __future__ import annotations + +from contextlib import contextmanager +import gc +import importlib +import time +from types import ModuleType +from typing import Any, Iterator + +import torch + +from backends.chunked_mtlbvh import ChunkedMtlBVH +from backends.metal_preserve import ( + texture_projection_kwargs, + use_geometry_preserving_backend, +) + + +def _cat_optional_tensors( + chunks: list[tuple[Any, ...]], +) -> tuple[Any, ...]: + outputs: list[Any] = [] + for values in zip(*chunks): + first = values[0] + if first is None: + outputs.append(None) + elif torch.is_tensor(first): + outputs.append(torch.cat(list(values), dim=0)) + else: + raise TypeError( + "Chunked BVH returned an unsupported value of type " + f"{type(first).__name__}" + ) + return tuple(outputs) + + +@contextmanager +def memory_bounded_o_voxel( + postprocess_module: ModuleType | Any, + *, + bvh_chunk_size: int = 262_144, + grid_chunk_size: int = 262_144, + source_resolution: int, + remesh_resolution: int | None = None, +) -> Iterator[None]: + """Temporarily bound native Metal query allocations. + + ``o_voxel.postprocess`` exposes its selected backends as module globals. + The patch is therefore process-global and intended for the serial CLI. + Every original object is restored even when export fails. + """ + + if bvh_chunk_size <= 0 or grid_chunk_size <= 0: + raise ValueError("Chunk sizes must be positive") + if source_resolution <= 0: + raise ValueError("source_resolution must be positive") + if remesh_resolution is not None and remesh_resolution <= 0: + raise ValueError("remesh_resolution must be positive when provided") + + original_bvh = postprocess_module._BVH + original_grid_sample = postprocess_module._grid_sample_3d + original_remesh = postprocess_module._remesh_narrow_band_dc + + class ChunkedBVH: + def __init__(self, vertices, faces): + self.inner = original_bvh(vertices, faces) + + def unsigned_distance(self, positions, return_uvw=False, **kwargs): + count = int(positions.shape[0]) + if count <= bvh_chunk_size: + return self.inner.unsigned_distance( + positions, + return_uvw=return_uvw, + **kwargs, + ) + chunks = [ + self.inner.unsigned_distance( + positions[start : start + bvh_chunk_size], + return_uvw=return_uvw, + **kwargs, + ) + for start in range(0, count, bvh_chunk_size) + ] + return _cat_optional_tensors(chunks) + + def __getattr__(self, name): + return getattr(self.inner, name) + + def chunked_grid_sample(feats, coords, shape, grid, mode="trilinear"): + length = int(grid.shape[1]) + outputs = [ + original_grid_sample( + feats, + coords, + shape, + grid[:, start : start + grid_chunk_size], + mode=mode, + ) + for start in range(0, length, grid_chunk_size) + ] + if not outputs: + return original_grid_sample(feats, coords, shape, grid, mode=mode) + output = torch.cat(outputs, dim=1) + # The Metal flex_gemm API returns [B, L, C], while the pinned + # postprocessor assigns a single batch into a [L, C] texture view. + if output.ndim == 3 and output.shape[0] == 1: + output = output[0] + return output + + def bounded_remesh(*args, **kwargs): + bvh = kwargs.get("bvh") + if isinstance(bvh, ChunkedBVH): + kwargs["bvh"] = bvh.inner + + requested_resolution = remesh_resolution + native_resolution = int(kwargs.get("resolution", source_resolution)) + if ( + requested_resolution is not None + and requested_resolution != native_resolution + ): + band = float(kwargs.get("band", 1.0)) + expanded_scale = float(kwargs["scale"]) + base_scale = ( + expanded_scale + * native_resolution + / (native_resolution + 3 * band) + ) + kwargs["resolution"] = int(requested_resolution) + kwargs["scale"] = ( + (requested_resolution + 3 * band) + / requested_resolution + * base_scale + ) + return original_remesh(*args, **kwargs) + + postprocess_module._BVH = ChunkedBVH + postprocess_module._grid_sample_3d = chunked_grid_sample + postprocess_module._remesh_narrow_band_dc = bounded_remesh + try: + yield + finally: + postprocess_module._BVH = original_bvh + postprocess_module._grid_sample_3d = original_grid_sample + postprocess_module._remesh_narrow_band_dc = original_remesh + + +def force_cuda_material_semantics(mesh: Any) -> None: + """Match the official opaque, single-sided glTF material flags.""" + + geometries = ( + mesh.geometry.values() + if hasattr(mesh, "geometry") + else [mesh] + ) + for geometry in geometries: + visual = getattr(geometry, "visual", None) + material = getattr(visual, "material", None) + if material is None: + continue + material.alphaMode = "OPAQUE" + material.doubleSided = False + + +def _release_metal_temporaries() -> None: + gc.collect() + if torch.backends.mps.is_available(): + torch.mps.synchronize() + torch.mps.empty_cache() + + +def _project_to_source( + source_bvh: ChunkedMtlBVH, + source_vertices: torch.Tensor, + source_faces: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + """Return closest source positions while bounding native BVH memory.""" + + distances, face_ids, uvw = source_bvh.unsigned_distance( + positions, + return_uvw=True, + ) + assert uvw is not None + source_triangles = source_vertices[ + source_faces[face_ids.long()].long() + ] + projected = ( + source_triangles * uvw.unsqueeze(-1) + ).sum(dim=1).contiguous() + del distances, face_ids, uvw, source_triangles + return projected + + +def _prepare_metal_geometry( + *, + source_vertices: torch.Tensor, + source_faces: torch.Tensor, + source_resolution: int, + remesh_resolution: int, + decimation_target: int, + source_face_chunk_size: int, + query_chunk_size: int, + verbose: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Accurately remesh and close a decoded mesh within unified memory.""" + + from cumesh import CuMesh, remeshing + from mtlbvh import MtlBVH + + remesh_bvh = ChunkedMtlBVH( + MtlBVH, + source_vertices, + source_faces, + source_face_chunk_size=source_face_chunk_size, + query_chunk_size=query_chunk_size, + ) + if verbose: + print( + "[Export/Geometry] Accurate Metal remesh: " + f"{remesh_resolution}³, " + f"{remesh_bvh.source_chunks} source BVH chunks." + ) + remesh_started = time.perf_counter() + candidate_vertices, candidate_faces = ( + remeshing.remesh_narrow_band_dc( + source_vertices, + source_faces, + center=torch.zeros(3, dtype=torch.float32), + scale=(remesh_resolution + 3) / remesh_resolution, + resolution=remesh_resolution, + band=1, + project_back=0, + verbose=verbose, + bvh=remesh_bvh, + ) + ) + if verbose: + print( + "[Export/Geometry] Raw remesh: " + f"{len(candidate_vertices):,} vertices, " + f"{len(candidate_faces):,} faces in " + f"{time.perf_counter() - remesh_started:.2f} s." + ) + del remesh_bvh + _release_metal_temporaries() + + cleanup_started = time.perf_counter() + mesh = CuMesh() + mesh.init(candidate_vertices, candidate_faces) + del candidate_vertices, candidate_faces + mesh.remove_duplicate_faces() + mesh.remove_degenerate_faces() + mesh.repair_non_manifold_edges() + mesh.remove_small_connected_components(1e-5) + mesh.fill_holes(max_hole_perimeter=3e-2) + mesh.simplify(decimation_target, verbose=verbose) + mesh.remove_duplicate_faces() + mesh.remove_degenerate_faces() + mesh.repair_non_manifold_edges() + mesh.remove_small_connected_components(1e-5) + # The remaining loops are tiny reconstruction defects. Closing all of + # them removes the interior-view effect without changing the silhouette. + mesh.fill_holes(max_hole_perimeter=10.0) + prepared_vertices, prepared_faces = mesh.read() + prepared_vertices = prepared_vertices.detach().cpu().float().contiguous() + prepared_faces = prepared_faces.detach().cpu().int().contiguous() + del mesh + _release_metal_temporaries() + if verbose: + print( + "[Export/Geometry] Prepared mesh: " + f"{len(prepared_vertices):,} vertices, " + f"{len(prepared_faces):,} faces in " + f"{time.perf_counter() - cleanup_started:.2f} s." + ) + return prepared_vertices, prepared_faces + + +def to_glb_cuda_parity( + *, + vertices: torch.Tensor, + faces: torch.Tensor, + attr_volume: torch.Tensor, + coords: torch.Tensor, + attr_layout: dict[str, slice], + resolution: int, + decimation_target: int = 1_000_000, + texture_size: int = 4096, + bvh_chunk_size: int = 262_144, + grid_chunk_size: int = 262_144, + source_face_chunk_size: int = 250_000, + remesh_resolution: int | None = None, + verbose: bool = True, + use_tqdm: bool = True, +): + """Run the validated high-quality profile through native Metal.""" + + if resolution <= 0: + raise ValueError("resolution must be positive") + if decimation_target <= 0 or texture_size <= 0: + raise ValueError("decimation_target and texture_size must be positive") + if source_face_chunk_size <= 8: + raise ValueError("source_face_chunk_size must be greater than 8") + + postprocess = importlib.import_module("o_voxel.postprocess") + if not getattr(postprocess, "_HAS_GPU_DEPS", False): + raise RuntimeError( + "CUDA-parity export requires mtldiffrast, cumesh/mtlmesh and mtlbvh" + ) + if getattr(postprocess, "_BACKEND", None) != "metal": + raise RuntimeError( + "CUDA-parity macOS export expected o_voxel's Metal backend" + ) + + effective_remesh_resolution = ( + int(remesh_resolution) + if remesh_resolution is not None + else min(512, resolution) + ) + source_vertices = vertices.detach().cpu().float().contiguous() + source_faces = faces.detach().cpu().int().contiguous() + prepared_vertices, prepared_faces = _prepare_metal_geometry( + source_vertices=source_vertices, + source_faces=source_faces, + source_resolution=resolution, + remesh_resolution=effective_remesh_resolution, + decimation_target=decimation_target, + source_face_chunk_size=source_face_chunk_size, + query_chunk_size=bvh_chunk_size, + verbose=verbose, + ) + + from mtlbvh import MtlBVH + + texture_bvh = ChunkedMtlBVH( + MtlBVH, + source_vertices, + source_faces, + source_face_chunk_size=source_face_chunk_size, + query_chunk_size=bvh_chunk_size, + ) + projection_started = time.perf_counter() + texture_sample_vertices = _project_to_source( + texture_bvh, + source_vertices, + source_faces, + prepared_vertices, + ) + if verbose: + print( + "[Export/Texture] Projected " + f"{len(texture_sample_vertices):,} sampling vertices in " + f"{time.perf_counter() - projection_started:.2f} s." + ) + + fallback_queries = 0 + + def project_invalid_texels(positions: torch.Tensor) -> torch.Tensor: + nonlocal fallback_queries + fallback_queries += len(positions) + return _project_to_source( + texture_bvh, + source_vertices, + source_faces, + positions, + ) + + projection_kwargs = texture_projection_kwargs(postprocess, "preserve") + with ( + memory_bounded_o_voxel( + postprocess, + bvh_chunk_size=bvh_chunk_size, + grid_chunk_size=grid_chunk_size, + source_resolution=resolution, + ), + use_geometry_preserving_backend(postprocess), + ): + result = postprocess.to_glb( + vertices=prepared_vertices, + faces=prepared_faces, + attr_volume=attr_volume.detach().cpu().contiguous(), + coords=coords.detach().cpu().contiguous(), + attr_layout=attr_layout, + grid_size=resolution, + aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + decimation_target=max(decimation_target, len(prepared_faces)), + texture_size=texture_size, + remesh=False, + remesh_band=1, + remesh_project=0, + verbose=verbose, + use_tqdm=use_tqdm, + texture_sample_vertices=texture_sample_vertices, + texture_fallback_projector=project_invalid_texels, + **projection_kwargs, + ) + if verbose: + print( + "[Export/Texture] Exact fallback projection: " + f"{fallback_queries:,} texels." + ) + del ( + texture_bvh, + texture_sample_vertices, + prepared_vertices, + prepared_faces, + ) + _release_metal_temporaries() + force_cuda_material_semantics(result) + return result diff --git a/backends/decoded_checkpoint.py b/backends/decoded_checkpoint.py new file mode 100644 index 0000000..d574432 --- /dev/null +++ b/backends/decoded_checkpoint.py @@ -0,0 +1,103 @@ +"""Portable checkpoints for Pixal3D's decoded mesh and sparse PBR volume. + +The expensive neural stages finish before ``o_voxel`` starts remeshing and +texture baking. Persisting that boundary makes high-quality export retries +cheap and lets a failed 4096px bake resume without re-running the diffusion +pipeline. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Mapping + +import torch + +from pixal3d.representations import MeshWithVoxel + + +CHECKPOINT_VERSION = 1 + + +def _serialize_layout(layout: Mapping[str, slice]) -> dict[str, list[int | None]]: + return { + name: [value.start, value.stop, value.step] + for name, value in layout.items() + } + + +def _deserialize_layout( + layout: Mapping[str, list[int | None]], +) -> dict[str, slice]: + return { + name: slice(*value) + for name, value in layout.items() + } + + +def save_decoded_checkpoint( + path: str | Path, + mesh: MeshWithVoxel, + *, + resolution: int, + metadata: Mapping[str, Any] | None = None, +) -> Path: + """Atomically save CPU tensors needed for all subsequent export steps.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_name(f".{output.name}.tmp") + payload = { + "version": CHECKPOINT_VERSION, + "resolution": int(resolution), + "vertices": mesh.vertices.detach().cpu().contiguous(), + "faces": mesh.faces.detach().cpu().contiguous(), + "coords": mesh.coords.detach().cpu().contiguous(), + "attrs": mesh.attrs.detach().cpu().contiguous(), + "origin": mesh.origin.detach().cpu().tolist(), + "voxel_size": float(mesh.voxel_size), + "voxel_shape": list(mesh.voxel_shape), + "layout": _serialize_layout(mesh.layout), + "metadata": dict(metadata or {}), + } + try: + torch.save(payload, temporary) + os.replace(temporary, output) + finally: + temporary.unlink(missing_ok=True) + return output + + +def load_decoded_checkpoint( + path: str | Path, +) -> tuple[MeshWithVoxel, int, dict[str, Any]]: + """Load a decoded checkpoint without importing arbitrary Python objects.""" + + source = Path(path) + payload = torch.load( + source, + map_location="cpu", + weights_only=True, + ) + version = int(payload.get("version", -1)) + if version != CHECKPOINT_VERSION: + raise ValueError( + f"Unsupported decoded checkpoint version {version}; " + f"expected {CHECKPOINT_VERSION}" + ) + resolution = int(payload["resolution"]) + if resolution <= 0: + raise ValueError("Decoded checkpoint resolution must be positive") + + mesh = MeshWithVoxel( + vertices=payload["vertices"], + faces=payload["faces"], + origin=payload["origin"], + voxel_size=float(payload["voxel_size"]), + coords=payload["coords"], + attrs=payload["attrs"], + voxel_shape=torch.Size(payload["voxel_shape"]), + layout=_deserialize_layout(payload["layout"]), + ) + return mesh, resolution, dict(payload.get("metadata", {})) diff --git a/backends/export_normals.py b/backends/export_normals.py new file mode 100644 index 0000000..5e79dd3 --- /dev/null +++ b/backends/export_normals.py @@ -0,0 +1,868 @@ +"""Deterministic normals for indexed meshes and UV-split exports. + +The indexed mesh must remain the source of truth for smooth shading. UV +parameterization duplicates vertices along chart boundaries; recomputing +normals afterwards turns those boundaries into unintended hard edges. This +module can therefore compute normals on the indexed, pre-UV mesh and transfer +them to an exact triangle-preserving export without welding positions. +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from dataclasses import asdict, dataclass +from typing import Any + +import numpy as np + + +@dataclass(frozen=True, slots=True) +class NormalRecomputeReport: + geometries: int + vertices: int + faces: int + degenerate_faces: int + cancellation_vertices_repaired: int + radial_fallback_vertices: int + + def to_dict(self) -> dict[str, int]: + return asdict(self) + + +@dataclass(frozen=True, slots=True) +class NormalTransferReport: + """Summary of a successful indexed-to-UV normal transfer.""" + + geometries: int + instances: int + reference_vertices: int + reference_faces: int + asset_vertices: int + asset_instance_vertices: int + asset_faces: int + uv_split_vertices: int + degenerate_reference_faces: int + cancellation_vertices_repaired: int + radial_fallback_vertices: int + locally_repaired_vertices: int + residual_locally_opposed_vertices: int + smoothing_groups: int + smoothed_asset_vertices: int + hard_split_reference_vertices: int + + def to_dict(self) -> dict[str, int]: + return asdict(self) + + +class NormalTransferError(ValueError): + """The exported asset cannot be matched to the reference unambiguously.""" + + +@dataclass(frozen=True, slots=True) +class _AssetInstance: + geometry_name: str + geometry: Any + local_to_world: np.ndarray + + +def recompute_vertex_normals( + vertices: np.ndarray, + faces: np.ndarray, +) -> tuple[np.ndarray, dict[str, int]]: + """Return finite unit normals consistent with the supplied face winding. + + Incident unit face normals are weighted by their corner angle, accumulated + per vertex and normalized. At non-manifold singularities they can cancel; + those rare vertices take the first deterministic non-degenerate incident + face normal. A radial fallback is reserved for vertices incident only to + degenerate triangles. + """ + + vertices = np.asarray(vertices, dtype=np.float64) + faces = np.asarray(faces, dtype=np.int64) + if vertices.ndim != 2 or vertices.shape[1] != 3 or not len(vertices): + raise ValueError("vertices must be a non-empty [N, 3] array") + if faces.ndim != 2 or faces.shape[1] != 3 or not len(faces): + raise ValueError("faces must be a non-empty [F, 3] array") + if not np.isfinite(vertices).all(): + raise ValueError("vertices contain non-finite values") + if np.min(faces) < 0 or np.max(faces) >= len(vertices): + raise ValueError("faces contain out-of-range vertex indices") + + triangles = vertices[faces] + face_normals = np.cross( + triangles[:, 1] - triangles[:, 0], + triangles[:, 2] - triangles[:, 0], + ) + face_lengths = np.linalg.norm(face_normals, axis=1) + nondegenerate = face_lengths > 1e-20 + unit_faces = np.zeros_like(face_normals) + unit_faces[nondegenerate] = ( + face_normals[nondegenerate] / face_lengths[nondegenerate, None] + ) + + normals = np.zeros_like(vertices) + for corner in range(3): + first_edge = triangles[:, (corner + 1) % 3] - triangles[:, corner] + second_edge = triangles[:, (corner + 2) % 3] - triangles[:, corner] + first_length = np.linalg.norm(first_edge, axis=1) + second_length = np.linalg.norm(second_edge, axis=1) + valid_corner = ( + nondegenerate & (first_length > 1e-20) & (second_length > 1e-20) + ) + cosine = np.ones(len(faces), dtype=np.float64) + cosine[valid_corner] = np.einsum( + "ij,ij->i", + first_edge[valid_corner] / first_length[valid_corner, None], + second_edge[valid_corner] / second_length[valid_corner, None], + optimize=False, + ) + corner_angle = np.zeros(len(faces), dtype=np.float64) + corner_angle[valid_corner] = np.arccos( + np.clip(cosine[valid_corner], -1.0, 1.0) + ) + np.add.at( + normals, + faces[:, corner], + unit_faces * corner_angle[:, None], + ) + lengths = np.linalg.norm(normals, axis=1) + valid = lengths > 1e-12 + normals[valid] /= lengths[valid, None] + cancellation = ~valid + cancellation_count = int(np.count_nonzero(cancellation)) + + if cancellation_count: + # Stable face order makes this fallback deterministic. UV seam vertices + # commonly have only one incident chart face, so they resolve directly. + fallback_face = np.full(len(vertices), -1, dtype=np.int64) + for corner in range(3): + vertex_ids = faces[nondegenerate, corner] + face_ids = np.flatnonzero(nondegenerate) + missing = fallback_face[vertex_ids] < 0 + fallback_face[vertex_ids[missing]] = face_ids[missing] + use_face = cancellation & (fallback_face >= 0) + normals[use_face] = unit_faces[fallback_face[use_face]] + cancellation &= ~use_face + + radial_fallback_count = int(np.count_nonzero(cancellation)) + if radial_fallback_count: + centre = (vertices.min(axis=0) + vertices.max(axis=0)) * 0.5 + radial = vertices[cancellation] - centre + radial_lengths = np.linalg.norm(radial, axis=1) + usable = radial_lengths > 1e-12 + radial[usable] /= radial_lengths[usable, None] + radial[~usable] = np.array([0.0, 0.0, 1.0]) + normals[cancellation] = radial + + final_lengths = np.linalg.norm(normals, axis=1) + if not np.isfinite(normals).all() or not np.allclose( + final_lengths, 1.0, atol=5e-6 + ): + raise RuntimeError("failed to produce finite unit vertex normals") + return np.asarray(normals, dtype=np.float32), { + "vertices": int(len(vertices)), + "faces": int(len(faces)), + "degenerate_faces": int(np.count_nonzero(~nondegenerate)), + "cancellation_vertices_repaired": cancellation_count, + "radial_fallback_vertices": radial_fallback_count, + } + + +def trellis_z_up_to_gltf_y_up_transform() -> np.ndarray: + """Return the homogeneous rotation from TRELLIS to glTF coordinates. + + This maps ``(x, y, z)`` to ``(x, z, -y)``. A fresh array is returned so a + caller cannot mutate module-global state accidentally. + """ + + return np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float64, + ) + + +def _mesh_arrays( + vertices: np.ndarray, + faces: np.ndarray, + *, + label: str, +) -> tuple[np.ndarray, np.ndarray]: + vertices = np.asarray(vertices, dtype=np.float32) + faces = np.asarray(faces, dtype=np.int64) + if vertices.ndim != 2 or vertices.shape[1] != 3 or not len(vertices): + raise NormalTransferError(f"{label} vertices must be non-empty [N, 3]") + if faces.ndim != 2 or faces.shape[1] != 3 or not len(faces): + raise NormalTransferError(f"{label} faces must be non-empty [F, 3]") + if not np.isfinite(vertices).all(): + raise NormalTransferError(f"{label} vertices contain non-finite values") + if np.min(faces) < 0 or np.max(faces) >= len(vertices): + raise NormalTransferError( + f"{label} faces contain out-of-range vertex indices" + ) + return ( + np.ascontiguousarray(vertices, dtype=np.float32), + np.ascontiguousarray(faces, dtype=np.int64), + ) + + +def _affine_transform( + transform: np.ndarray | None, + *, + label: str, +) -> np.ndarray: + if transform is None: + return np.eye(4, dtype=np.float64) + matrix = np.asarray(transform, dtype=np.float64) + if matrix.shape != (4, 4) or not np.isfinite(matrix).all(): + raise NormalTransferError(f"{label} must be a finite [4, 4] matrix") + if not np.allclose(matrix[3], [0.0, 0.0, 0.0, 1.0], atol=1e-12): + raise NormalTransferError(f"{label} must be affine, not projective") + determinant = float(np.linalg.det(matrix[:3, :3])) + if not np.isfinite(determinant) or abs(determinant) <= 1e-12: + raise NormalTransferError(f"{label} has a singular linear transform") + return matrix + + +def _transform_points(vertices: np.ndarray, transform: np.ndarray) -> np.ndarray: + transformed = ( + np.asarray(vertices, dtype=np.float64) @ transform[:3, :3].T + + transform[:3, 3] + ) + if not np.isfinite(transformed).all(): + raise NormalTransferError("coordinate transform produced non-finite vertices") + return np.ascontiguousarray(transformed, dtype=np.float32) + + +def _normalise_rows(normals: np.ndarray, *, label: str) -> np.ndarray: + normals = np.asarray(normals, dtype=np.float64) + lengths = np.linalg.norm(normals, axis=1) + if ( + not np.isfinite(normals).all() + or not np.isfinite(lengths).all() + or np.any(lengths <= 1e-12) + ): + raise NormalTransferError(f"{label} produced invalid normals") + return normals / lengths[:, None] + + +def _world_to_local_normals( + normals: np.ndarray, + local_to_world: np.ndarray, +) -> np.ndarray: + """Transform winding-consistent world normals into geometry-local space.""" + + linear = local_to_world[:3, :3] + # For row vectors, A.T maps a world normal back to local coordinates. A + # reflection also reverses triangle winding, hence the determinant sign. + determinant_sign = 1.0 if np.linalg.det(linear) > 0.0 else -1.0 + local = np.asarray(normals, dtype=np.float64) @ linear + return _normalise_rows( + local * determinant_sign, + label="scene inverse normal transform", + ) + + +def _asset_instances(asset: Any) -> tuple[list[_AssetInstance], dict[str, Any]]: + """Return mesh instances while preserving scene-node transforms.""" + + if not hasattr(asset, "geometry"): + return ( + [ + _AssetInstance( + geometry_name="mesh", + geometry=asset, + local_to_world=np.eye(4, dtype=np.float64), + ) + ], + {"mesh": asset}, + ) + + geometries = dict(asset.geometry) + graph = getattr(asset, "graph", None) + nodes = list(getattr(graph, "nodes_geometry", ())) + if not geometries or not nodes: + raise NormalTransferError( + "asset scene must contain at least one instanced mesh geometry" + ) + + instances: list[_AssetInstance] = [] + referenced: set[str] = set() + for node in nodes: + try: + transform, geometry_name = graph.get(node) + geometry = geometries[geometry_name] + except (KeyError, TypeError, ValueError) as exc: + raise NormalTransferError( + f"cannot resolve scene geometry for node {node!r}" + ) from exc + local_to_world = _affine_transform( + transform, + label=f"scene transform for node {node!r}", + ) + geometry_name = str(geometry_name) + referenced.add(geometry_name) + instances.append( + _AssetInstance( + geometry_name=geometry_name, + geometry=geometry, + local_to_world=local_to_world, + ) + ) + + unused = set(geometries).difference(referenced) + if unused: + raise NormalTransferError( + "asset scene contains uninstanced geometries: " + + ", ".join(sorted(map(str, unused))) + ) + return instances, geometries + + +def _position_bits(vertices: np.ndarray) -> np.ndarray: + """Return canonical exact-float32 keys, treating signed zero as equal.""" + + canonical = np.array(vertices, dtype=np.float32, order="C", copy=True) + canonical[canonical == 0.0] = 0.0 + return canonical.view(np.uint32).reshape((-1, 3)) + + +def _oriented_face_key( + vertex_bits: np.ndarray, + face: np.ndarray, +) -> tuple[tuple[int, int, int], ...]: + corners = tuple( + tuple(int(value) for value in vertex_bits[int(vertex)]) for vertex in face + ) + rotations = ( + corners, + corners[1:] + corners[:1], + corners[2:] + corners[:2], + ) + return min(rotations) + + +def _aligned_reference_indices( + output_corner_bits: np.ndarray, + candidate_faces: list[int], + reference_bits: np.ndarray, + reference_faces: np.ndarray, +) -> list[np.ndarray]: + """Enumerate every exact orientation-preserving corner correspondence.""" + + aligned: list[np.ndarray] = [] + for face_id in candidate_faces: + reference_face = reference_faces[face_id] + for offset in range(3): + indices = reference_face[ + np.array([(offset + corner) % 3 for corner in range(3)]) + ] + if np.array_equal(reference_bits[indices], output_corner_bits): + aligned.append(indices) + return aligned + + +def _incident_face_cones( + vertices: np.ndarray, + faces: np.ndarray, +) -> tuple[list[np.ndarray], np.ndarray, np.ndarray]: + """Return unit face constraints and angle-weighted per-vertex accumulators.""" + + triangles = np.asarray(vertices, dtype=np.float64)[faces] + area_normals = np.cross( + triangles[:, 1] - triangles[:, 0], + triangles[:, 2] - triangles[:, 0], + ) + lengths = np.linalg.norm(area_normals, axis=1) + valid = lengths > 1e-20 + unit_faces = np.zeros_like(area_normals) + unit_faces[valid] = area_normals[valid] / lengths[valid, None] + + accumulator = np.zeros((len(vertices), 3), dtype=np.float64) + weights = np.zeros(len(vertices), dtype=np.float64) + incident_lists: list[list[np.ndarray]] = [[] for _ in range(len(vertices))] + valid_face_ids = np.flatnonzero(valid) + for corner in range(3): + first_edge = triangles[:, (corner + 1) % 3] - triangles[:, corner] + second_edge = triangles[:, (corner + 2) % 3] - triangles[:, corner] + first_length = np.linalg.norm(first_edge, axis=1) + second_length = np.linalg.norm(second_edge, axis=1) + valid_corner = valid & (first_length > 1e-20) & (second_length > 1e-20) + angles = np.zeros(len(faces), dtype=np.float64) + cosine = np.einsum( + "ij,ij->i", + first_edge[valid_corner] / first_length[valid_corner, None], + second_edge[valid_corner] / second_length[valid_corner, None], + optimize=False, + ) + angles[valid_corner] = np.arccos(np.clip(cosine, -1.0, 1.0)) + np.add.at(accumulator, faces[:, corner], unit_faces * angles[:, None]) + np.add.at(weights, faces[:, corner], angles) + for face_id in valid_face_ids: + incident_lists[int(faces[face_id, corner])].append(unit_faces[face_id]) + + incident = [ + np.asarray(normals, dtype=np.float64).reshape((-1, 3)) + for normals in incident_lists + ] + return incident, accumulator, weights + + +def _nonopposed( + normal: np.ndarray, + constraints: np.ndarray, + *, + tolerance: float = 0.0, +) -> bool: + return not len(constraints) or bool( + np.all(np.asarray(constraints) @ np.asarray(normal) >= -tolerance) + ) + + +def _repair_to_face_cone( + normal: np.ndarray, + constraints: np.ndarray, +) -> np.ndarray | None: + """Project a rare opposed local normal into its incident-face cone.""" + + candidate = _normalise_rows( + np.asarray(normal, dtype=np.float64).reshape((1, 3)), + label="local post-UV normal", + )[0] + constraints = np.asarray(constraints, dtype=np.float64).reshape((-1, 3)) + if _nonopposed(candidate, constraints): + return candidate + if not len(constraints): + return candidate + + # Alternating projections are only needed for the handful of singular + # UV vertices whose angle-weighted sum lies just outside one constraint. + # A small positive margin avoids a numerically negative dot in the GLB + # validator after the final float32 conversion. + margin = 2e-7 + starts = [candidate, constraints.sum(axis=0)] + starts.extend(constraint for constraint in constraints) + for start in starts: + length = float(np.linalg.norm(start)) + if length <= 1e-12: + continue + projected = np.asarray(start, dtype=np.float64) / length + for _ in range(256): + dots = constraints @ projected + worst = int(np.argmin(dots)) + if dots[worst] >= margin: + break + projected = projected + (margin - dots[worst]) * constraints[worst] + length = float(np.linalg.norm(projected)) + if length <= 1e-12: + break + projected /= length + projected32 = np.asarray(projected, dtype=np.float32).astype(np.float64) + projected32 /= np.linalg.norm(projected32) + if _nonopposed(projected32, constraints): + return projected32 + return None + + +@dataclass(slots=True) +class _SmoothingGroup: + members: list[int] + raw_normal: np.ndarray + normal: np.ndarray + constraints: np.ndarray + + +def _smooth_source_copies( + copy_ids: list[int], + reference_normal: np.ndarray, + incident: list[np.ndarray], + raw_normals: np.ndarray, + local_normals: np.ndarray, +) -> tuple[dict[int, np.ndarray], int, int, bool, int, int]: + """Partition UV copies into safe smoothing groups for one source vertex.""" + + constraint_sets = [ + incident[copy_id] for copy_id in copy_ids if len(incident[copy_id]) + ] + all_constraints = ( + np.concatenate(constraint_sets, axis=0) + if constraint_sets + else np.empty((0, 3), dtype=np.float64) + ) + reference_normal = _normalise_rows( + np.asarray(reference_normal).reshape((1, 3)), + label="pre-UV reference normal", + )[0] + if _nonopposed(reference_normal, all_constraints): + return ( + {copy_id: reference_normal for copy_id in copy_ids}, + 1, + len(copy_ids) if len(copy_ids) > 1 else 0, + False, + 0, + 0, + ) + + groups: list[_SmoothingGroup] = [] + locally_repaired = 0 + for copy_id in copy_ids: + constraints = incident[copy_id] + repaired = _repair_to_face_cone(local_normals[copy_id], constraints) + if repaired is None: + # Some non-manifold fans have no non-zero vector in the intersection + # of all incident hemispheres. Keep the already validated post-UV + # normal for that hard group; importantly, never spread it to a + # second chart unless the merged result satisfies every constraint. + repaired = local_normals[copy_id] + elif np.linalg.norm(repaired - local_normals[copy_id]) > 5e-6: + locally_repaired += 1 + raw = np.asarray(raw_normals[copy_id], dtype=np.float64) + if np.linalg.norm(raw) <= 1e-12: + raw = repaired.copy() + groups.append( + _SmoothingGroup( + members=[copy_id], + raw_normal=raw, + normal=repaired, + constraints=constraints, + ) + ) + + # Deterministic agglomeration favours the most similarly shaded charts. + # Pairwise normal compatibility prevents a coincident back-facing sheet + # from being averaged merely because the position is identical. + while len(groups) > 1: + candidates: list[tuple[float, int, int, np.ndarray, np.ndarray]] = [] + for left in range(len(groups)): + for right in range(left + 1, len(groups)): + similarity = float(groups[left].normal @ groups[right].normal) + if similarity < 0.0: + continue + raw = groups[left].raw_normal + groups[right].raw_normal + length = float(np.linalg.norm(raw)) + if length <= 1e-12: + continue + candidate = raw / length + constraints = np.concatenate( + (groups[left].constraints, groups[right].constraints), axis=0 + ) + candidate32 = np.asarray(candidate, dtype=np.float32).astype( + np.float64 + ) + candidate32 /= np.linalg.norm(candidate32) + if _nonopposed(candidate32, constraints): + candidates.append( + (similarity, left, right, candidate32, constraints) + ) + if not candidates: + break + _, left, right, candidate, constraints = max( + candidates, + key=lambda item: (item[0], -item[1], -item[2]), + ) + merged = _SmoothingGroup( + members=groups[left].members + groups[right].members, + raw_normal=groups[left].raw_normal + groups[right].raw_normal, + normal=candidate, + constraints=constraints, + ) + groups = [ + group + for index, group in enumerate(groups) + if index not in (left, right) + ] + groups.append(merged) + + result: dict[int, np.ndarray] = {} + smoothed = 0 + residual_opposed = 0 + for group in groups: + if len(group.members) > 1: + smoothed += len(group.members) + for copy_id in group.members: + result[copy_id] = group.normal + residual_opposed += int( + not _nonopposed(group.normal, incident[copy_id]) + ) + return ( + result, + len(groups), + smoothed, + len(groups) > 1, + locally_repaired, + residual_opposed, + ) + + +def transfer_reference_vertex_normals( + asset: Any, + reference_vertices: np.ndarray, + reference_faces: np.ndarray, + *, + reference_to_asset: np.ndarray | None = None, + normal_tolerance: float = 5e-6, +) -> NormalTransferReport: + """Transfer pre-UV indexed normals to an exact UV-split trimesh asset. + + ``reference_to_asset`` maps reference positions into the asset scene's + world coordinate system. Pass + :func:`trellis_z_up_to_gltf_y_up_transform` for the Metal baker's glTF + output. Scene-node transforms are respected and the resulting normals are + stored in each geometry's local coordinates. + + Matching uses exact float32, orientation-preserving triangle positions. + Positions are never welded, and every exported corner must resolve to one + unique indexed reference vertex. Compatible UV copies receive their + shared pre-UV normal. On non-manifold vertices where that normal would be + opposed to an incident face, copies are partitioned into deterministic + smoothing groups and hard splits are retained. Any topology change, + reversed triangle, isolated output vertex, or identity ambiguity raises + :class:`NormalTransferError` before the asset is mutated. + """ + + if ( + not np.isfinite(normal_tolerance) + or normal_tolerance <= 0.0 + or normal_tolerance >= 1.0 + ): + raise ValueError("normal_tolerance must be finite and between zero and one") + + reference_vertices, reference_faces = _mesh_arrays( + reference_vertices, + reference_faces, + label="reference", + ) + reference_transform = _affine_transform( + reference_to_asset, + label="reference_to_asset", + ) + reference_world = _transform_points(reference_vertices, reference_transform) + reference_normals, normal_report = recompute_vertex_normals( + reference_world, + reference_faces, + ) + reference_normals = np.asarray(reference_normals, dtype=np.float64) + reference_bits = _position_bits(reference_world) + + reference_by_triangle: dict[ + tuple[tuple[int, int, int], ...], list[int] + ] = defaultdict(list) + for face_id, face in enumerate(reference_faces): + reference_by_triangle[_oriented_face_key(reference_bits, face)].append( + face_id + ) + reference_counts = Counter( + {key: len(face_ids) for key, face_ids in reference_by_triangle.items()} + ) + + instances, geometries = _asset_instances(asset) + geometry_arrays: dict[str, tuple[np.ndarray, np.ndarray]] = {} + for geometry_name, geometry in geometries.items(): + geometry_arrays[str(geometry_name)] = _mesh_arrays( + np.asarray(geometry.vertices), + np.asarray(geometry.faces), + label=f"asset geometry {geometry_name!r}", + ) + + instance_world: list[tuple[_AssetInstance, np.ndarray, np.ndarray]] = [] + asset_counts: Counter[tuple[tuple[int, int, int], ...]] = Counter() + asset_instance_vertices = 0 + asset_faces = 0 + for instance in instances: + local_vertices, faces = geometry_arrays[instance.geometry_name] + world_vertices = _transform_points(local_vertices, instance.local_to_world) + world_bits = _position_bits(world_vertices) + instance_world.append((instance, world_vertices, world_bits)) + asset_instance_vertices += len(local_vertices) + asset_faces += len(faces) + for face in faces: + asset_counts[_oriented_face_key(world_bits, face)] += 1 + + if asset_counts != reference_counts: + missing = sum((reference_counts - asset_counts).values()) + extra = sum((asset_counts - reference_counts).values()) + reversed_hint = "" + if missing == extra and missing: + reversed_hint = " (triangles may have reversed winding)" + raise NormalTransferError( + "asset triangles do not exactly match the oriented reference " + f"multiset: {missing} missing, {extra} extra{reversed_hint}" + ) + + pending_normals: dict[str, np.ndarray] = { + name: np.zeros((len(arrays[0]), 3), dtype=np.float64) + for name, arrays in geometry_arrays.items() + } + assigned: dict[str, np.ndarray] = { + name: np.zeros(len(arrays[0]), dtype=bool) + for name, arrays in geometry_arrays.items() + } + locally_repaired_vertices = 0 + smoothing_groups = 0 + smoothed_asset_vertices = 0 + hard_split_reference_vertices = 0 + residual_locally_opposed_vertices = 0 + + for instance, world_vertices, world_bits in instance_world: + _, faces = geometry_arrays[instance.geometry_name] + geometry_normals = pending_normals[instance.geometry_name] + geometry_assigned = assigned[instance.geometry_name] + source_ids = np.full(len(world_vertices), -1, dtype=np.int64) + for face in faces: + key = _oriented_face_key(world_bits, face) + variants = _aligned_reference_indices( + world_bits[face], + reference_by_triangle[key], + reference_bits, + reference_faces, + ) + if not variants: + raise NormalTransferError( + "internal error: matched triangle has no corner correspondence" + ) + chosen_source = variants[0] + for variant in variants[1:]: + if not np.array_equal(variant, chosen_source): + raise NormalTransferError( + "ambiguous coincident reference triangles map to " + "different indexed proxy identities" + ) + for corner, vertex_id in enumerate(face): + vertex_id = int(vertex_id) + source_id = int(chosen_source[corner]) + if source_ids[vertex_id] >= 0: + if source_ids[vertex_id] != source_id: + raise NormalTransferError( + "one exported vertex maps to multiple indexed proxy " + "identities; the asset may have merged distinct sheets" + ) + else: + source_ids[vertex_id] = source_id + + if np.any(source_ids < 0): + raise NormalTransferError( + "asset contains vertices not referenced by any triangle" + ) + + local_world_normals, _ = recompute_vertex_normals( + world_vertices, + faces, + ) + local_world_normals = np.asarray(local_world_normals, dtype=np.float64) + incident, raw_normals, _ = _incident_face_cones(world_vertices, faces) + copies_by_source: dict[int, list[int]] = defaultdict(list) + for copy_id, source_id in enumerate(source_ids): + copies_by_source[int(source_id)].append(copy_id) + + instance_normals = np.zeros_like(local_world_normals) + for source_id, copy_ids in copies_by_source.items(): + ( + smoothed, + group_count, + smoothed_count, + hard_split, + repaired_count, + residual_opposed_count, + ) = _smooth_source_copies( + copy_ids, + reference_normals[source_id], + incident, + raw_normals, + local_world_normals, + ) + smoothing_groups += group_count + smoothed_asset_vertices += smoothed_count + hard_split_reference_vertices += int(hard_split) + locally_repaired_vertices += repaired_count + residual_locally_opposed_vertices += residual_opposed_count + for copy_id, normal in smoothed.items(): + instance_normals[copy_id] = normal + + chosen_local = _world_to_local_normals( + instance_normals, + instance.local_to_world, + ) + for vertex_id, normal in enumerate(chosen_local): + if geometry_assigned[vertex_id]: + disagreement = np.linalg.norm( + geometry_normals[vertex_id] - normal + ) + if disagreement > normal_tolerance: + raise NormalTransferError( + "an instanced geometry requires incompatible local normals" + ) + else: + geometry_normals[vertex_id] = normal + geometry_assigned[vertex_id] = True + + unassigned = { + name: int(np.count_nonzero(~mask)) for name, mask in assigned.items() + } + unassigned = {name: count for name, count in unassigned.items() if count} + if unassigned: + detail = ", ".join( + f"{name}: {count}" for name, count in sorted(unassigned.items()) + ) + raise NormalTransferError( + f"asset contains vertices not referenced by any triangle ({detail})" + ) + + # Mutate only after every geometry and instance has passed the strict gate. + for geometry_name, geometry in geometries.items(): + geometry.vertex_normals = np.asarray( + pending_normals[str(geometry_name)], + dtype=np.float32, + ) + + asset_vertices = sum(len(arrays[0]) for arrays in geometry_arrays.values()) + return NormalTransferReport( + geometries=len(geometries), + instances=len(instances), + reference_vertices=int(len(reference_vertices)), + reference_faces=int(len(reference_faces)), + asset_vertices=int(asset_vertices), + asset_instance_vertices=int(asset_instance_vertices), + asset_faces=int(asset_faces), + uv_split_vertices=max(0, int(asset_instance_vertices - len(reference_vertices))), + degenerate_reference_faces=normal_report["degenerate_faces"], + cancellation_vertices_repaired=normal_report[ + "cancellation_vertices_repaired" + ], + radial_fallback_vertices=normal_report["radial_fallback_vertices"], + locally_repaired_vertices=locally_repaired_vertices, + residual_locally_opposed_vertices=residual_locally_opposed_vertices, + smoothing_groups=smoothing_groups, + smoothed_asset_vertices=smoothed_asset_vertices, + hard_split_reference_vertices=hard_split_reference_vertices, + ) + + +def recompute_asset_vertex_normals(asset: Any) -> NormalRecomputeReport: + """Replace normals on every mesh in a trimesh asset without changing it.""" + + geometries = ( + list(asset.geometry.values()) + if hasattr(asset, "geometry") + else [asset] + ) + totals = { + "vertices": 0, + "faces": 0, + "degenerate_faces": 0, + "cancellation_vertices_repaired": 0, + "radial_fallback_vertices": 0, + } + for geometry in geometries: + normals, report = recompute_vertex_normals( + np.asarray(geometry.vertices), + np.asarray(geometry.faces), + ) + geometry.vertex_normals = normals + for key in totals: + totals[key] += report[key] + return NormalRecomputeReport(geometries=len(geometries), **totals) diff --git a/backends/export_validation.py b/backends/export_validation.py new file mode 100644 index 0000000..1ef5108 --- /dev/null +++ b/backends/export_validation.py @@ -0,0 +1,790 @@ +"""Strict geometry validation for textured GLB exports. + +TRELLIS meshes are Z-up while glTF stores Y-up geometry. This module reloads +an exported GLB, converts it back to TRELLIS coordinates, welds only vertices +whose float32 positions are exactly equal, and compares the resulting triangle +multiset with the geometry handed to the texture baker. GLB validation also +requires explicit finite unit vertex normals and reports their agreement with +the final triangle winding. + +The exact weld is intentional: UV charts duplicate vertices at seams, but a +tolerance-based weld could hide a crack by joining two genuinely distinct, +nearby surface layers. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Literal + +import numpy as np + + +CoordinateSystem = Literal["gltf_y_up", "trellis_z_up"] + + +@dataclass(frozen=True, slots=True) +class ExportMeshMetrics: + """Topology measurements after welding exact float32 positions.""" + + source_vertices: int + welded_vertices: int + faces: int + boundary_edges: int + boundary_length: float + boundary_components: int + closed_boundary_loops: int + small_closed_boundary_loops: int + small_closed_boundary_edges: int + nonmanifold_edges: int + + +@dataclass(frozen=True, slots=True) +class ExportNormalMetrics: + """Validation and diagnostics for explicitly exported vertex normals.""" + + required: bool + present: bool + shape_valid: bool + expected_count: int + count: int + nonfinite: int + zero_length: int + nonunit: int + opposed_corner_normals: int + faces_with_any_opposed_corner_normals: int + faces_with_all_opposed_corner_normals: int + max_opposed_corner_normals: int + max_faces_with_all_opposed_corner_normals: int + unit_tolerance: float + + @property + def passed(self) -> bool: + if not self.required: + return True + return ( + self.present + and self.shape_valid + and self.count == self.expected_count + and self.nonfinite == 0 + and self.zero_length == 0 + and self.nonunit == 0 + and self.faces_with_all_opposed_corner_normals + <= self.max_faces_with_all_opposed_corner_normals + and self.opposed_corner_normals <= self.max_opposed_corner_normals + ) + + +@dataclass(frozen=True, slots=True) +class ExportValidationResult: + """Result of comparing a baked export with its reference geometry.""" + + reference: ExportMeshMetrics + exported: ExportMeshMetrics + triangle_multiset_matches: bool + oriented_triangle_multiset_matches: bool + matched_triangles: int + same_winding_triangles: int + reversed_winding_triangles: int + missing_triangles: int + extra_triangles: int + normals: ExportNormalMetrics + reasons: tuple[str, ...] + + @property + def passed(self) -> bool: + """Whether the export preserved every reference triangle exactly.""" + + return ( + self.triangle_multiset_matches + and self.oriented_triangle_multiset_matches + and self.normals.passed + ) + + def raise_for_error(self) -> None: + """Raise :class:`ExportValidationError` when validation failed.""" + + if not self.passed: + raise ExportValidationError(self) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable report.""" + + report = asdict(self) + report["normals"]["passed"] = self.normals.passed + report["passed"] = self.passed + return report + + +class ExportValidationError(RuntimeError): + """Raised when a GLB does not preserve its bake-input geometry.""" + + def __init__(self, result: ExportValidationResult): + self.result = result + detail = "; ".join(result.reasons) or "export validation failed" + super().__init__(detail) + + +def _mesh_arrays( + vertices: np.ndarray, faces: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + vertices = np.asarray(vertices, dtype=np.float32) + faces = np.asarray(faces, dtype=np.int64) + if vertices.ndim != 2 or vertices.shape[1] != 3 or not len(vertices): + raise ValueError("vertices must be a non-empty [N, 3] array") + if faces.ndim != 2 or faces.shape[1] != 3 or not len(faces): + raise ValueError("faces must be a non-empty [F, 3] array") + if not np.isfinite(vertices).all(): + raise ValueError("vertices contain non-finite values") + if np.min(faces) < 0 or np.max(faces) >= len(vertices): + raise ValueError("faces contain out-of-range vertex indices") + return np.ascontiguousarray(vertices), np.ascontiguousarray(faces) + + +def trellis_z_up_to_gltf_y_up(vertices: np.ndarray) -> np.ndarray: + """Rotate TRELLIS ``(x, y, z)`` positions to glTF ``(x, z, -y)``.""" + + vertices = np.asarray(vertices, dtype=np.float32) + if vertices.ndim != 2 or vertices.shape[1] != 3: + raise ValueError("vertices must be an [N, 3] array") + converted = np.empty_like(vertices) + converted[:, 0] = vertices[:, 0] + converted[:, 1] = vertices[:, 2] + converted[:, 2] = -vertices[:, 1] + return converted + + +def gltf_y_up_to_trellis_z_up(vertices: np.ndarray) -> np.ndarray: + """Rotate glTF ``(x, y, z)`` positions back to TRELLIS ``(x, -z, y)``.""" + + vertices = np.asarray(vertices, dtype=np.float32) + if vertices.ndim != 2 or vertices.shape[1] != 3: + raise ValueError("vertices must be an [N, 3] array") + converted = np.empty_like(vertices) + converted[:, 0] = vertices[:, 0] + converted[:, 1] = -vertices[:, 2] + converted[:, 2] = vertices[:, 1] + return converted + + +def weld_exact_float32( + vertices: np.ndarray, faces: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Merge vertices only when all three float32 coordinates compare equal.""" + + vertices, faces = _mesh_arrays(vertices, faces) + welded_vertices, inverse = np.unique(vertices, axis=0, return_inverse=True) + welded_faces = inverse[faces] + return ( + np.ascontiguousarray(welded_vertices, dtype=np.float32), + np.ascontiguousarray(welded_faces, dtype=np.int64), + ) + + +def _boundary_components( + boundary: np.ndarray, vertex_count: int, small_loop_max_edges: int +) -> tuple[int, int, int, int]: + if not len(boundary): + return 0, 0, 0, 0 + + parent = np.arange(vertex_count, dtype=np.int64) + rank = np.zeros(vertex_count, dtype=np.uint8) + + def find(node: int) -> int: + root = node + while parent[root] != root: + root = int(parent[root]) + while parent[node] != node: + next_node = int(parent[node]) + parent[node] = root + node = next_node + return root + + def union(left: int, right: int) -> None: + left_root = find(left) + right_root = find(right) + if left_root == right_root: + return + if rank[left_root] < rank[right_root]: + left_root, right_root = right_root, left_root + parent[right_root] = left_root + if rank[left_root] == rank[right_root]: + rank[left_root] += 1 + + for left, right in boundary: + union(int(left), int(right)) + + edge_roots = np.fromiter( + (find(int(left)) for left in boundary[:, 0]), + dtype=np.int64, + count=len(boundary), + ) + component_ids, edge_component = np.unique(edge_roots, return_inverse=True) + component_edge_counts = np.bincount( + edge_component, minlength=len(component_ids) + ) + + degree = np.bincount(boundary.reshape(-1), minlength=vertex_count) + boundary_vertices = np.flatnonzero(degree) + root_to_component = { + int(root): index for index, root in enumerate(component_ids.tolist()) + } + non_loop_vertices = np.zeros(len(component_ids), dtype=np.int64) + for vertex in boundary_vertices: + if degree[vertex] != 2: + non_loop_vertices[root_to_component[find(int(vertex))]] += 1 + + closed = non_loop_vertices == 0 + small = closed & (component_edge_counts <= small_loop_max_edges) + return ( + int(len(component_ids)), + int(np.count_nonzero(closed)), + int(np.count_nonzero(small)), + int(component_edge_counts[small].sum()), + ) + + +def export_mesh_metrics( + vertices: np.ndarray, + faces: np.ndarray, + *, + small_loop_max_edges: int = 12, +) -> ExportMeshMetrics: + """Measure topology after an exact float32 position weld.""" + + if small_loop_max_edges < 1: + raise ValueError("small_loop_max_edges must be positive") + source_vertices = len(np.asarray(vertices)) + vertices, faces = weld_exact_float32(vertices, faces) + + edges = np.concatenate( + (faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]), axis=0 + ) + edges.sort(axis=1) + unique_edges, edge_counts = np.unique(edges, axis=0, return_counts=True) + boundary = unique_edges[edge_counts == 1] + boundary_length = float( + np.linalg.norm( + vertices[boundary[:, 0]].astype(np.float64) + - vertices[boundary[:, 1]].astype(np.float64), + axis=1, + ).sum() + ) + ( + boundary_components, + closed_boundary_loops, + small_closed_boundary_loops, + small_closed_boundary_edges, + ) = _boundary_components(boundary, len(vertices), small_loop_max_edges) + + return ExportMeshMetrics( + source_vertices=int(source_vertices), + welded_vertices=int(len(vertices)), + faces=int(len(faces)), + boundary_edges=int(len(boundary)), + boundary_length=boundary_length, + boundary_components=boundary_components, + closed_boundary_loops=closed_boundary_loops, + small_closed_boundary_loops=small_closed_boundary_loops, + small_closed_boundary_edges=small_closed_boundary_edges, + nonmanifold_edges=int(np.count_nonzero(edge_counts > 2)), + ) + + +def export_normal_metrics( + vertices: np.ndarray, + faces: np.ndarray, + normals: np.ndarray | None, + *, + required: bool, + unit_tolerance: float = 5e-3, + max_opposed_corner_fraction: float = 1e-4, + max_all_opposed_face_fraction: float = 1e-5, +) -> ExportNormalMetrics: + """Validate explicit vertex normals and measure face/corner disagreement.""" + + if not np.isfinite(unit_tolerance) or unit_tolerance <= 0: + raise ValueError("unit_tolerance must be finite and positive") + if ( + not np.isfinite(max_opposed_corner_fraction) + or not 0 <= max_opposed_corner_fraction <= 1 + ): + raise ValueError( + "max_opposed_corner_fraction must be between zero and one" + ) + if ( + not np.isfinite(max_all_opposed_face_fraction) + or not 0 <= max_all_opposed_face_fraction <= 1 + ): + raise ValueError( + "max_all_opposed_face_fraction must be between zero and one" + ) + vertices, faces = _mesh_arrays(vertices, faces) + expected_count = len(vertices) + max_opposed_corners = max( + 3, + int(np.ceil(len(faces) * 3 * max_opposed_corner_fraction)), + ) + max_all_opposed_faces = int( + np.floor(len(faces) * max_all_opposed_face_fraction) + ) + if normals is None: + return ExportNormalMetrics( + required=required, + present=False, + shape_valid=False, + expected_count=expected_count, + count=0, + nonfinite=0, + zero_length=0, + nonunit=0, + opposed_corner_normals=0, + faces_with_any_opposed_corner_normals=0, + faces_with_all_opposed_corner_normals=0, + max_opposed_corner_normals=max_opposed_corners, + max_faces_with_all_opposed_corner_normals=max_all_opposed_faces, + unit_tolerance=float(unit_tolerance), + ) + + normals_array = np.asarray(normals) + shape_valid = normals_array.ndim == 2 and normals_array.shape[1:] == (3,) + count = int(len(normals_array)) if normals_array.ndim else 0 + if not shape_valid: + return ExportNormalMetrics( + required=required, + present=True, + shape_valid=False, + expected_count=expected_count, + count=count, + nonfinite=0, + zero_length=0, + nonunit=0, + opposed_corner_normals=0, + faces_with_any_opposed_corner_normals=0, + faces_with_all_opposed_corner_normals=0, + max_opposed_corner_normals=max_opposed_corners, + max_faces_with_all_opposed_corner_normals=max_all_opposed_faces, + unit_tolerance=float(unit_tolerance), + ) + + normals_array = np.asarray(normals_array, dtype=np.float64) + finite_rows = np.isfinite(normals_array).all(axis=1) + nonfinite = int(np.count_nonzero(~finite_rows)) + lengths = np.zeros(len(normals_array), dtype=np.float64) + lengths[finite_rows] = np.linalg.norm(normals_array[finite_rows], axis=1) + zero_length = int(np.count_nonzero(finite_rows & (lengths <= 1e-12))) + nonunit = int( + np.count_nonzero( + finite_rows + & (lengths > 1e-12) + & (np.abs(lengths - 1.0) > unit_tolerance) + ) + ) + + opposed_corners = 0 + faces_with_any_opposed = 0 + faces_with_all_opposed = 0 + if count == expected_count and nonfinite == 0: + triangles = vertices.astype(np.float64)[faces] + face_normals = np.cross( + triangles[:, 1] - triangles[:, 0], + triangles[:, 2] - triangles[:, 0], + ) + valid_faces = np.linalg.norm(face_normals, axis=1) > 1e-12 + corner_alignment = np.einsum( + "fci,fi->fc", + normals_array[faces], + face_normals, + optimize=False, + ) + opposed = (corner_alignment < 0.0) & valid_faces[:, None] + opposed_corners = int(np.count_nonzero(opposed)) + faces_with_any_opposed = int(np.count_nonzero(np.any(opposed, axis=1))) + faces_with_all_opposed = int(np.count_nonzero(np.all(opposed, axis=1))) + + return ExportNormalMetrics( + required=required, + present=True, + shape_valid=True, + expected_count=expected_count, + count=count, + nonfinite=nonfinite, + zero_length=zero_length, + nonunit=nonunit, + opposed_corner_normals=opposed_corners, + faces_with_any_opposed_corner_normals=faces_with_any_opposed, + faces_with_all_opposed_corner_normals=faces_with_all_opposed, + max_opposed_corner_normals=max_opposed_corners, + max_faces_with_all_opposed_corner_normals=max_all_opposed_faces, + unit_tolerance=float(unit_tolerance), + ) + + +def _compare_triangle_multisets( + reference_vertices: np.ndarray, + reference_faces: np.ndarray, + exported_vertices: np.ndarray, + exported_faces: np.ndarray, +) -> tuple[int, int, int, int, int, int]: + """Compare unoriented and oriented triangle occurrence multisets.""" + + combined_vertices = np.concatenate( + (reference_vertices, exported_vertices), axis=0 + ) + _, inverse = np.unique(combined_vertices, axis=0, return_inverse=True) + reference_ids = inverse[: len(reference_vertices)] + exported_ids = inverse[len(reference_vertices) :] + + reference_oriented = reference_ids[reference_faces] + exported_oriented = exported_ids[exported_faces] + reference_triangles = np.sort(reference_oriented, axis=1) + exported_triangles = np.sort(exported_oriented, axis=1) + + def canonical_orientation(triangles: np.ndarray) -> np.ndarray: + # Cyclic rotations preserve winding. Pick the lexicographically smallest + # rotation; reversed winding keeps the final two IDs swapped. Comparing + # all rotations also gives deterministic behavior for degenerate faces. + rotations = ( + triangles, + triangles[:, [1, 2, 0]], + triangles[:, [2, 0, 1]], + ) + canonical = rotations[0].copy() + + def row_is_less(left: np.ndarray, right: np.ndarray) -> np.ndarray: + return ( + (left[:, 0] < right[:, 0]) + | ( + (left[:, 0] == right[:, 0]) + & (left[:, 1] < right[:, 1]) + ) + | ( + (left[:, 0] == right[:, 0]) + & (left[:, 1] == right[:, 1]) + & (left[:, 2] < right[:, 2]) + ) + ) + + for rotation in rotations[1:]: + replace = row_is_less(rotation, canonical) + canonical[replace] = rotation[replace] + return canonical + + def overlap_counts( + reference: np.ndarray, exported: np.ndarray + ) -> tuple[int, int, int]: + all_triangles = np.concatenate((reference, exported), axis=0) + _, triangle_inverse = np.unique(all_triangles, axis=0, return_inverse=True) + reference_counts = np.bincount( + triangle_inverse[: len(reference)] + ) + exported_counts = np.bincount( + triangle_inverse[len(reference) :], + minlength=len(reference_counts), + ) + if len(exported_counts) > len(reference_counts): + reference_counts = np.pad( + reference_counts, (0, len(exported_counts) - len(reference_counts)) + ) + elif len(reference_counts) > len(exported_counts): + exported_counts = np.pad( + exported_counts, (0, len(reference_counts) - len(exported_counts)) + ) + + matched = int(np.minimum(reference_counts, exported_counts).sum()) + missing = int(np.maximum(reference_counts - exported_counts, 0).sum()) + extra = int(np.maximum(exported_counts - reference_counts, 0).sum()) + return matched, missing, extra + + matched, missing, extra = overlap_counts( + reference_triangles, exported_triangles + ) + oriented_matched, oriented_missing, oriented_extra = overlap_counts( + canonical_orientation(reference_oriented), + canonical_orientation(exported_oriented), + ) + return ( + matched, + missing, + extra, + oriented_matched, + oriented_missing, + oriented_extra, + ) + + +def validate_export_mesh( + reference_vertices: np.ndarray, + reference_faces: np.ndarray, + exported_vertices: np.ndarray, + exported_faces: np.ndarray, + *, + exported_normals: np.ndarray | None = None, + require_normals: bool = False, + normal_unit_tolerance: float = 5e-3, + max_opposed_corner_fraction: float = 1e-4, + max_all_opposed_face_fraction: float = 1e-5, + exported_coordinates: CoordinateSystem = "gltf_y_up", + small_loop_max_edges: int = 12, +) -> ExportValidationResult: + """Compare exported geometry with a TRELLIS Z-up reference mesh. + + Face order, vertex order and exact UV-seam duplicates are ignored. Every + triangle occurrence, float32 position and face winding must be preserved. + """ + + reference_vertices, reference_faces = _mesh_arrays( + reference_vertices, reference_faces + ) + exported_vertices, exported_faces = _mesh_arrays( + exported_vertices, exported_faces + ) + transformed_normals = exported_normals + if exported_coordinates == "gltf_y_up": + exported_vertices = gltf_y_up_to_trellis_z_up(exported_vertices) + candidate_normals = np.asarray(exported_normals) if exported_normals is not None else None + if ( + candidate_normals is not None + and candidate_normals.ndim == 2 + and candidate_normals.shape[1:] == (3,) + ): + transformed_normals = gltf_y_up_to_trellis_z_up(candidate_normals) + elif exported_coordinates != "trellis_z_up": + raise ValueError(f"unsupported coordinate system: {exported_coordinates!r}") + + reference_welded_vertices, reference_welded_faces = weld_exact_float32( + reference_vertices, reference_faces + ) + exported_welded_vertices, exported_welded_faces = weld_exact_float32( + exported_vertices, exported_faces + ) + ( + matched, + missing, + extra, + oriented_matched, + oriented_missing, + oriented_extra, + ) = _compare_triangle_multisets( + reference_welded_vertices, + reference_welded_faces, + exported_welded_vertices, + exported_welded_faces, + ) + + reference_metrics = export_mesh_metrics( + reference_vertices, + reference_faces, + small_loop_max_edges=small_loop_max_edges, + ) + exported_metrics = export_mesh_metrics( + exported_vertices, + exported_faces, + small_loop_max_edges=small_loop_max_edges, + ) + normal_metrics = export_normal_metrics( + exported_vertices, + exported_faces, + transformed_normals, + required=require_normals, + unit_tolerance=normal_unit_tolerance, + max_opposed_corner_fraction=max_opposed_corner_fraction, + max_all_opposed_face_fraction=max_all_opposed_face_fraction, + ) + matches = missing == 0 and extra == 0 + oriented_matches = oriented_missing == 0 and oriented_extra == 0 + reversed_winding = max(0, matched - oriented_matched) + reasons: list[str] = [] + if reference_metrics.faces != exported_metrics.faces: + reasons.append( + "face count changed " + f"({reference_metrics.faces} -> {exported_metrics.faces})" + ) + if not matches: + reasons.append( + "exact float32 triangle multiset differs " + f"({missing} missing, {extra} extra)" + ) + if reversed_winding: + reasons.append( + f"face winding changed ({reversed_winding} triangles reversed)" + ) + if exported_metrics.boundary_edges > reference_metrics.boundary_edges: + reasons.append( + "boundary edge count increased " + f"({reference_metrics.boundary_edges} -> " + f"{exported_metrics.boundary_edges})" + ) + if ( + exported_metrics.small_closed_boundary_loops + > reference_metrics.small_closed_boundary_loops + ): + reasons.append( + "small closed boundary loop count increased " + f"({reference_metrics.small_closed_boundary_loops} -> " + f"{exported_metrics.small_closed_boundary_loops})" + ) + if exported_metrics.nonmanifold_edges > reference_metrics.nonmanifold_edges: + reasons.append( + "non-manifold edge count increased " + f"({reference_metrics.nonmanifold_edges} -> " + f"{exported_metrics.nonmanifold_edges})" + ) + if require_normals and not normal_metrics.present: + reasons.append("GLB has no explicit vertex normals") + elif require_normals and not normal_metrics.shape_valid: + reasons.append("GLB vertex normals are not an [N, 3] array") + elif require_normals: + if normal_metrics.count != normal_metrics.expected_count: + reasons.append( + "vertex normal count differs from position count " + f"({normal_metrics.count} != {normal_metrics.expected_count})" + ) + if normal_metrics.nonfinite: + reasons.append( + f"vertex normals contain {normal_metrics.nonfinite} non-finite rows" + ) + if normal_metrics.zero_length: + reasons.append( + f"vertex normals contain {normal_metrics.zero_length} zero vectors" + ) + if normal_metrics.nonunit: + reasons.append( + f"vertex normals contain {normal_metrics.nonunit} non-unit vectors " + f"(tolerance {normal_metrics.unit_tolerance:g})" + ) + if ( + normal_metrics.faces_with_all_opposed_corner_normals + > normal_metrics.max_faces_with_all_opposed_corner_normals + ): + reasons.append( + "too many faces have inward normals at all three corners " + f"({normal_metrics.faces_with_all_opposed_corner_normals} > " + f"{normal_metrics.max_faces_with_all_opposed_corner_normals})" + ) + if ( + normal_metrics.opposed_corner_normals + > normal_metrics.max_opposed_corner_normals + ): + reasons.append( + "too many vertex normals oppose their incident face winding " + f"({normal_metrics.opposed_corner_normals} > " + f"{normal_metrics.max_opposed_corner_normals} corners)" + ) + + return ExportValidationResult( + reference=reference_metrics, + exported=exported_metrics, + triangle_multiset_matches=matches, + oriented_triangle_multiset_matches=oriented_matches, + matched_triangles=matched, + same_winding_triangles=oriented_matched, + reversed_winding_triangles=reversed_winding, + missing_triangles=missing, + extra_triangles=extra, + normals=normal_metrics, + reasons=tuple(reasons), + ) + + +def _load_glb_mesh( + path: str | Path, +) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: + import trimesh + + loaded = trimesh.load(str(path), force=None, process=False) + + def cached_normals(mesh: Any) -> np.ndarray | None: + cache = getattr(getattr(mesh, "_cache", None), "cache", {}) + value = cache.get("vertex_normals") + return None if value is None else np.asarray(value) + + if not isinstance(loaded, trimesh.Scene): + if not hasattr(loaded, "vertices") or not hasattr(loaded, "faces"): + raise ValueError(f"GLB does not contain a triangular mesh: {path}") + return ( + np.asarray(loaded.vertices), + np.asarray(loaded.faces), + cached_normals(loaded), + ) + + if not loaded.geometry or not loaded.graph.nodes_geometry: + raise ValueError(f"GLB scene has no geometry: {path}") + + vertices_parts: list[np.ndarray] = [] + faces_parts: list[np.ndarray] = [] + normal_parts: list[np.ndarray] = [] + all_normals_present = True + vertex_offset = 0 + for node_name in loaded.graph.nodes_geometry: + transform, geometry_name = loaded.graph[node_name] + geometry = loaded.geometry[geometry_name] + local_vertices = np.asarray(geometry.vertices) + transformed_vertices = trimesh.transformations.transform_points( + local_vertices, transform + ) + vertices_parts.append(transformed_vertices) + faces_parts.append(np.asarray(geometry.faces) + vertex_offset) + vertex_offset += len(local_vertices) + + local_normals = cached_normals(geometry) + if local_normals is None: + all_normals_present = False + continue + linear = np.asarray(transform, dtype=np.float64)[:3, :3] + try: + normal_transform = np.linalg.inv(linear) + except np.linalg.LinAlgError as exc: + raise ValueError( + f"GLB geometry node has a singular transform: {node_name}" + ) from exc + local_normals = np.asarray(local_normals, dtype=np.float64) + transformed_normals = local_normals @ normal_transform + # glTF renderers normalize after applying the inverse-transpose normal + # matrix. Preserve the accessor's original magnitude so a non-unit + # NORMAL remains detectable, while removing node-scale distortion. + local_lengths = np.linalg.norm(local_normals, axis=1) + transformed_lengths = np.linalg.norm(transformed_normals, axis=1) + usable = ( + np.isfinite(local_lengths) + & np.isfinite(transformed_lengths) + & (local_lengths > 0.0) + & (transformed_lengths > 0.0) + ) + transformed_normals[usable] *= ( + local_lengths[usable] / transformed_lengths[usable] + )[:, None] + normal_parts.append(transformed_normals) + + return ( + np.concatenate(vertices_parts, axis=0), + np.concatenate(faces_parts, axis=0), + np.concatenate(normal_parts, axis=0) if all_normals_present else None, + ) + + +def validate_glb_export( + reference_vertices: np.ndarray, + reference_faces: np.ndarray, + glb_path: str | Path, + *, + small_loop_max_edges: int = 12, + max_opposed_corner_fraction: float = 1e-4, + max_all_opposed_face_fraction: float = 1e-5, +) -> ExportValidationResult: + """Reload and strictly validate a Y-up GLB against TRELLIS geometry.""" + + exported_vertices, exported_faces, exported_normals = _load_glb_mesh(glb_path) + return validate_export_mesh( + reference_vertices, + reference_faces, + exported_vertices, + exported_faces, + exported_normals=exported_normals, + require_normals=True, + max_opposed_corner_fraction=max_opposed_corner_fraction, + max_all_opposed_face_fraction=max_all_opposed_face_fraction, + exported_coordinates="gltf_y_up", + small_loop_max_edges=small_loop_max_edges, + ) diff --git a/backends/face_budget.py b/backends/face_budget.py new file mode 100644 index 0000000..81aafb4 --- /dev/null +++ b/backends/face_budget.py @@ -0,0 +1,99 @@ +"""Resolve a textured-export face budget without penalising denser inputs. + +The historical 200k default was calibrated on the mono 512 control, whose +decoded mesh contains 549,814 faces. Reusing that absolute cap for a denser +multi-view mesh discards substantially more geometry. The adaptive default +therefore retains the control's exact ratio while preserving the two existing +explicit CLI meanings:: + + --pbr-face-target omitted -> adaptive control ratio + --pbr-face-target N -> explicit cap of N faces + --pbr-face-target 0 -> full-detail export + +For CLI integration, use ``None`` as the argparse default. Existing commands +that pass an integer remain backward compatible, including the old behaviour +via an explicit ``--pbr-face-target 200000``. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Literal + + +MONO_CONTROL_SOURCE_FACES = 549_814 +MONO_CONTROL_TARGET_FACES = 200_000 +DEFAULT_RETENTION_RATIO = ( + MONO_CONTROL_TARGET_FACES / MONO_CONTROL_SOURCE_FACES +) + +FaceBudgetMode = Literal["adaptive", "explicit", "full"] + + +@dataclass(frozen=True, slots=True) +class FaceBudget: + """Resolved simplification budget plus enough provenance to report it.""" + + source_faces: int + target_faces: int + mode: FaceBudgetMode + requested_target: int | None + + @property + def retention_ratio(self) -> float: + return self.target_faces / self.source_faces + + def to_dict(self) -> dict[str, int | float | str | None]: + report = asdict(self) + report["retention_ratio"] = self.retention_ratio + return report + + +def _require_non_negative_int(name: str, value: int) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if value < 0: + raise ValueError(f"{name} must be zero or positive") + return value + + +def resolve_pbr_face_budget( + source_faces: int, + requested_target: int | None = None, +) -> FaceBudget: + """Resolve the face target for one PBR bake proxy. + + ``requested_target=None`` applies the mono-control retention ratio. The + integer calculation rounds to the nearest face without floating-point + drift, so the known 769,286-face multi-view input resolves to 279,835. + A positive explicit request remains a cap and ``0`` keeps the full mesh. + """ + + source_faces = _require_non_negative_int("source_faces", source_faces) + if source_faces == 0: + raise ValueError("source_faces must be positive") + + if requested_target is None: + numerator = source_faces * MONO_CONTROL_TARGET_FACES + target_faces = ( + numerator + MONO_CONTROL_SOURCE_FACES // 2 + ) // MONO_CONTROL_SOURCE_FACES + target_faces = min(source_faces, max(1, target_faces)) + mode: FaceBudgetMode = "adaptive" + else: + requested_target = _require_non_negative_int( + "requested_target", requested_target + ) + if requested_target == 0: + target_faces = source_faces + mode = "full" + else: + target_faces = min(source_faces, requested_target) + mode = "explicit" + + return FaceBudget( + source_faces=source_faces, + target_faces=target_faces, + mode=mode, + requested_target=requested_target, + ) diff --git a/backends/face_orientation.py b/backends/face_orientation.py new file mode 100644 index 0000000..7359515 --- /dev/null +++ b/backends/face_orientation.py @@ -0,0 +1,646 @@ +"""Optional face-winding repair for roughly closed solid assets. + +The decoded TRELLIS mesh can contain many edge-disconnected or locally +inconsistent surface patches. Preserving that winding exactly is important as +an export invariant, but it does not imply that the normals point outwards. + +``orient_faces_radially`` first makes every manifold edge-connected component +locally consistent, then chooses the component's global orientation from an +area-weighted radial score around the mesh bounding-box centre. It changes +only the order of indices inside existing triangles: vertices, topology and +triangle occurrences remain immutable. + +This heuristic is deliberately opt-in. It works well for roughly closed, +star-shaped objects such as a human body, but a radial direction is not a +reliable definition of "outside" for open sheets, deep cavities or strongly +concave assets. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import asdict, dataclass +from typing import Any + +import numpy as np + + +@dataclass(frozen=True, slots=True) +class RadialOrientationReport: + """Diagnostics for a topology-preserving radial winding pass.""" + + faces: int + vertices: int + manifold_adjacencies: int + edge_components: int + singleton_components: int + degenerate_faces: int + duplicate_triangle_occurrences: int + manifold_winding_conflicts_before: int + manifold_winding_conflicts_after: int + local_faces_flipped: int + radial_components_flipped: int + radial_faces_flipped: int + ambiguous_components: int + low_confidence_components: int + low_confidence_faces: int + faces_flipped_from_input: int + radial_score_epsilon: float + confidence_threshold: float + median_component_confidence: float + radially_outward_face_fraction_before: float + radially_outward_face_fraction_after: float + radially_outward_area_fraction_before: float + radially_outward_area_fraction_after: float + topology_changed: bool = False + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable report.""" + + return asdict(self) + + +def _mesh_arrays( + vertices: np.ndarray, faces: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + vertices = np.asarray(vertices, dtype=np.float32) + input_faces = np.asarray(faces) + if vertices.ndim != 2 or vertices.shape[1] != 3 or not len(vertices): + raise ValueError("vertices must be a non-empty [N, 3] array") + if input_faces.ndim != 2 or input_faces.shape[1] != 3 or not len(input_faces): + raise ValueError("faces must be a non-empty [F, 3] array") + if not np.issubdtype(input_faces.dtype, np.integer): + raise ValueError("faces must contain integer vertex indices") + if not np.isfinite(vertices).all(): + raise ValueError("vertices contain non-finite values") + faces64 = np.asarray(input_faces, dtype=np.int64) + if np.min(faces64) < 0 or np.max(faces64) >= len(vertices): + raise ValueError("faces contain out-of-range vertex indices") + return np.ascontiguousarray(vertices), np.ascontiguousarray(input_faces) + + +def _manifold_adjacencies( + faces: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return face pairs and their required relative flip parity. + + Only edges used by exactly two distinct faces are constraints. Edges used + once are boundaries; edges used more than twice are non-manifold and do not + have one unambiguous opposite-face relationship. + """ + + face_count = len(faces) + directed = np.concatenate( + ( + faces[:, [0, 1]], + faces[:, [1, 2]], + faces[:, [2, 0]], + ), + axis=0, + ) + face_ids = np.tile(np.arange(face_count, dtype=np.int64), 3) + nondegenerate = directed[:, 0] != directed[:, 1] + directed = directed[nondegenerate] + face_ids = face_ids[nondegenerate] + undirected = np.sort(directed, axis=1) + + order = np.lexsort((undirected[:, 1], undirected[:, 0])) + ordered_edges = undirected[order] + if not len(ordered_edges): + empty = np.empty(0, dtype=np.int64) + return empty, empty, np.empty(0, dtype=bool) + + group_start = np.r_[ + 0, + np.flatnonzero(np.any(ordered_edges[1:] != ordered_edges[:-1], axis=1)) + + 1, + ] + group_end = np.r_[group_start[1:], len(ordered_edges)] + paired_groups = np.flatnonzero(group_end - group_start == 2) + first_occurrence = order[group_start[paired_groups]] + second_occurrence = order[group_start[paired_groups] + 1] + + left = face_ids[first_occurrence] + right = face_ids[second_occurrence] + distinct_faces = left != right + left = left[distinct_faces] + right = right[distinct_faces] + first_occurrence = first_occurrence[distinct_faces] + second_occurrence = second_occurrence[distinct_faces] + + # When both faces traverse their shared edge in the same direction, exactly + # one must be flipped. Opposite directions require equal flip parity. + requires_different_parity = ( + directed[first_occurrence, 0] == directed[second_occurrence, 0] + ) + return left, right, requires_different_parity + + +def _local_orientation( + face_count: int, + left: np.ndarray, + right: np.ndarray, + requires_different_parity: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, int, int]: + """Solve manifold edge constraints with deterministic local refinement. + + A breadth-first spanning forest gives an exact solution for orientable + components. Non-orientable or otherwise contradictory components depend + on which constraints landed outside that forest, however, and a single + local defect can consequently be reported as many residual conflicts. + + After the forest pass, ``_refine_parity_with_bridge_cuts`` moves those + defects across deterministic cuts of the already-satisfied graph. Every + accepted cut strictly lowers the number of violated constraints, so the + refinement cannot regress a component that the forest already solved. + """ + + degree = np.bincount( + np.concatenate((left, right)), minlength=face_count + ).astype(np.int64, copy=False) + offsets = np.empty(face_count + 1, dtype=np.int64) + offsets[0] = 0 + np.cumsum(degree, out=offsets[1:]) + neighbours = np.empty(offsets[-1], dtype=np.int64) + constraints = np.empty(offsets[-1], dtype=bool) + cursor = offsets[:-1].copy() + for edge_index in range(len(left)): + left_face = int(left[edge_index]) + right_face = int(right[edge_index]) + constraint = bool(requires_different_parity[edge_index]) + + position = int(cursor[left_face]) + neighbours[position] = right_face + constraints[position] = constraint + cursor[left_face] += 1 + + position = int(cursor[right_face]) + neighbours[position] = left_face + constraints[position] = constraint + cursor[right_face] += 1 + + parity = np.full(face_count, -1, dtype=np.int8) + component_ids = np.full(face_count, -1, dtype=np.int64) + component_sizes: list[int] = [] + queue: deque[int] = deque() + + for seed in range(face_count): + if parity[seed] >= 0: + continue + component_index = len(component_sizes) + parity[seed] = 0 + component_ids[seed] = component_index + queue.append(seed) + component_size = 0 + + while queue: + face = queue.popleft() + component_size += 1 + for position in range(int(offsets[face]), int(offsets[face + 1])): + neighbour = int(neighbours[position]) + expected = int(parity[face]) ^ int(constraints[position]) + if parity[neighbour] < 0: + parity[neighbour] = expected + component_ids[neighbour] = component_index + queue.append(neighbour) + + component_sizes.append(component_size) + + parity, constraint_violations = _refine_parity_with_bridge_cuts( + parity.astype(bool), + left, + right, + requires_different_parity, + ) + sizes = np.asarray(component_sizes, dtype=np.int64) + return ( + parity, + component_ids, + int(np.count_nonzero(sizes == 1)), + constraint_violations, + ) + + +def _satisfied_forest( + face_count: int, + left: np.ndarray, + right: np.ndarray, + satisfied: np.ndarray, +) -> tuple[ + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, +]: + """Build a DFS forest and identify bridges in the satisfied-edge graph.""" + + satisfied_edges = np.flatnonzero(satisfied) + selected_left = left[satisfied_edges] + selected_right = right[satisfied_edges] + degree = np.bincount( + np.concatenate((selected_left, selected_right)), + minlength=face_count, + ).astype(np.int64, copy=False) + offsets = np.empty(face_count + 1, dtype=np.int64) + offsets[0] = 0 + np.cumsum(degree, out=offsets[1:]) + directed_faces = np.concatenate((selected_left, selected_right)) + neighbours = np.concatenate((selected_right, selected_left)) + edge_indices = np.concatenate((satisfied_edges, satisfied_edges)) + adjacency_order = np.lexsort((edge_indices, directed_faces)) + neighbours = neighbours[adjacency_order] + edge_indices = edge_indices[adjacency_order] + + discovery = np.full(face_count, -1, dtype=np.int64) + low = np.full(face_count, -1, dtype=np.int64) + parent = np.full(face_count, -1, dtype=np.int64) + parent_edge = np.full(face_count, -1, dtype=np.int64) + depth = np.zeros(face_count, dtype=np.int64) + subtree_size = np.ones(face_count, dtype=np.int64) + root = np.full(face_count, -1, dtype=np.int64) + bridge_child = np.zeros(face_count, dtype=bool) + order = np.empty(face_count, dtype=np.int64) + next_position = offsets[:-1].copy() + order_size = 0 + + for seed in range(face_count): + if discovery[seed] >= 0: + continue + discovery[seed] = order_size + low[seed] = order_size + root[seed] = seed + order[order_size] = seed + order_size += 1 + stack = [seed] + + while stack: + face = stack[-1] + if next_position[face] < offsets[face + 1]: + position = int(next_position[face]) + next_position[face] += 1 + edge_index = int(edge_indices[position]) + if edge_index == parent_edge[face]: + continue + neighbour = int(neighbours[position]) + if discovery[neighbour] < 0: + parent[neighbour] = face + parent_edge[neighbour] = edge_index + depth[neighbour] = depth[face] + 1 + root[neighbour] = root[face] + discovery[neighbour] = order_size + low[neighbour] = order_size + order[order_size] = neighbour + order_size += 1 + stack.append(neighbour) + else: + low[face] = min(low[face], discovery[neighbour]) + continue + + stack.pop() + parent_face = int(parent[face]) + if parent_face >= 0: + subtree_size[parent_face] += subtree_size[face] + low[parent_face] = min(low[parent_face], low[face]) + if low[face] > discovery[parent_face]: + bridge_child[face] = True + + if order_size != face_count: # pragma: no cover - defensive invariant + raise RuntimeError("satisfied-edge traversal missed one or more faces") + return ( + parent, + depth, + root, + discovery, + subtree_size, + bridge_child, + order, + ) + + +def _lowest_common_ancestors( + parent: np.ndarray, + depth: np.ndarray, + root: np.ndarray, + left: np.ndarray, + right: np.ndarray, +) -> np.ndarray: + """Return deterministic LCAs for pairs in the same DFS tree.""" + + if len(left) != len(right): # pragma: no cover - internal invariant + raise RuntimeError("LCA endpoint arrays have different lengths") + if not len(left): + return np.empty(0, dtype=np.int64) + if np.any(root[left] != root[right]): + raise RuntimeError("LCA endpoints belong to different trees") + + face_count = len(parent) + parent_safe = np.where( + parent >= 0, + parent, + np.arange(face_count, dtype=np.int64), + ) + ancestors = [parent_safe] + maximum_depth = int(depth.max(initial=0)) + while 1 << len(ancestors) <= maximum_depth: + previous = ancestors[-1] + ancestors.append(previous[previous]) + + first = left.astype(np.int64, copy=True) + second = right.astype(np.int64, copy=True) + swap = depth[first] < depth[second] + first[swap], second[swap] = second[swap], first[swap].copy() + depth_difference = depth[first] - depth[second] + for level, ancestor in enumerate(ancestors): + move = (depth_difference & (1 << level)) != 0 + first[move] = ancestor[first[move]] + + different = first != second + for ancestor in reversed(ancestors): + move = different & (ancestor[first] != ancestor[second]) + first[move] = ancestor[first[move]] + second[move] = ancestor[second[move]] + return np.where(different, parent_safe[first], first) + + +def _refine_parity_with_bridge_cuts( + parity: np.ndarray, + left: np.ndarray, + right: np.ndarray, + requires_different_parity: np.ndarray, +) -> tuple[np.ndarray, int]: + """Strictly reduce residual XOR conflicts using satisfied-graph bridges. + + Flipping one side of a bridge replaces its one satisfied crossing edge by + every currently violated edge crossing the same cut. The move is accepted + only when at least two violated edges cross, reducing the objective by at + least one. The bounded pass count prevents adversarial meshes from turning + a quality safeguard into unbounded export work. + """ + + parity = np.asarray(parity, dtype=bool).copy() + face_count = len(parity) + residual = ( + (parity[left] ^ parity[right]) != requires_different_parity + ) + residual_count = int(np.count_nonzero(residual)) + maximum_passes = min(residual_count, 64) + + for _ in range(maximum_passes): + if residual_count < 2: + break + ( + parent, + depth, + root, + discovery, + subtree_size, + bridge_child, + order, + ) = _satisfied_forest(face_count, left, right, ~residual) + residual_edges = np.flatnonzero(residual) + residual_left = left[residual_edges] + residual_right = right[residual_edges] + + # A disconnected satisfied graph offers a zero-cost cut. This is not + # expected after the initial BFS forest, but handling it keeps the + # refinement correct for future callers and after unusual multiedges. + crosses_root = root[residual_left] != root[residual_right] + if np.any(crosses_root): + crossing_roots = np.concatenate( + ( + root[residual_left[crosses_root]], + root[residual_right[crosses_root]], + ) + ) + roots, crossing_counts = np.unique( + crossing_roots, return_counts=True + ) + root_sizes = subtree_size[roots] + candidate_order = np.lexsort((roots, root_sizes, -crossing_counts)) + chosen_root = int(roots[candidate_order[0]]) + flip_faces = np.flatnonzero(root == chosen_root) + else: + lcas = _lowest_common_ancestors( + parent, + depth, + root, + residual_left, + residual_right, + ) + path_counts = np.zeros(face_count, dtype=np.int64) + np.add.at(path_counts, residual_left, 1) + np.add.at(path_counts, residual_right, 1) + np.add.at(path_counts, lcas, -2) + for face in order[::-1]: + parent_face = int(parent[face]) + if parent_face >= 0: + path_counts[parent_face] += path_counts[face] + + candidates = np.flatnonzero( + bridge_child & (path_counts >= 2) + ) + if not len(candidates): + break + # Cuts in different satisfied components are independent. Apply + # the best cut from every such component in the same pass, while + # retaining one-at-a-time behavior for nested cuts in a component. + candidate_order = np.lexsort( + ( + candidates, + subtree_size[candidates], + -path_counts[candidates], + root[candidates], + ) + ) + ordered_candidates = candidates[candidate_order] + _, first_per_root = np.unique( + root[ordered_candidates], return_index=True + ) + chosen_faces = ordered_candidates[np.sort(first_per_root)] + flip_ranges = [] + for chosen in chosen_faces: + start = int(discovery[chosen]) + stop = start + int(subtree_size[chosen]) + flip_ranges.append(order[start:stop]) + flip_faces = np.concatenate(flip_ranges) + + previous_count = residual_count + parity[flip_faces] ^= True + residual = ( + (parity[left] ^ parity[right]) != requires_different_parity + ) + residual_count = int(np.count_nonzero(residual)) + if residual_count >= previous_count: # pragma: no cover - invariant + raise RuntimeError("bridge-cut refinement did not reduce conflicts") + + return parity, residual_count + + +def _radial_statistics( + vertices: np.ndarray, + faces: np.ndarray, + centre: np.ndarray, +) -> tuple[np.ndarray, float, float]: + triangles = vertices[faces] + area_normals = np.cross( + triangles[:, 1] - triangles[:, 0], + triangles[:, 2] - triangles[:, 0], + ) + centroids = triangles.mean(axis=1) + scores = np.einsum( + "ij,ij->i", area_normals, centroids - centre, optimize=False + ) + positive = scores > 0.0 + face_fraction = float(np.count_nonzero(positive) / len(faces)) + double_area = np.linalg.norm(area_normals, axis=1) + total_area = float(double_area.sum()) + area_fraction = ( + float(double_area[positive].sum() / total_area) + if total_area > 0.0 + else 0.0 + ) + return scores, face_fraction, area_fraction + + +def orient_faces_radially( + vertices: np.ndarray, + faces: np.ndarray, + *, + score_epsilon_scale: float = 1e-12, + confidence_threshold: float = 0.1, +) -> tuple[np.ndarray, RadialOrientationReport]: + """Locally unify and radially orient existing triangles. + + Args: + vertices: Mesh positions shaped ``[N, 3]``. + faces: Integer triangle indices shaped ``[F, 3]``. + score_epsilon_scale: Dimensionless ambiguity threshold, scaled by the + cube of the mesh bounding-box diagonal because radial scores have + units of volume. + confidence_threshold: Report components below this normalized radial + agreement score as low-confidence without changing their result. + + Returns: + A contiguous face array with the input dtype and a diagnostic report. + + The returned triangles have the same indices as the input triangles; only + winding can differ. Connectivity comes strictly from shared indices. This + function is intended to run before UV seams duplicate vertices, avoiding + accidental joins between coincident but semantically separate shells. + """ + + if not np.isfinite(score_epsilon_scale) or score_epsilon_scale < 0: + raise ValueError("score_epsilon_scale must be finite and non-negative") + if ( + not np.isfinite(confidence_threshold) + or not 0 <= confidence_threshold <= 1 + ): + raise ValueError("confidence_threshold must be between zero and one") + vertices32, input_faces = _mesh_arrays(vertices, faces) + face_dtype = input_faces.dtype + faces64 = np.asarray(input_faces, dtype=np.int64) + duplicate_triangle_occurrences = int( + len(faces64) - len(np.unique(np.sort(faces64, axis=1), axis=0)) + ) + + # This pass runs before UV unwrapping, so connectivity must come from the + # authored indices. Welding coincident positions here could incorrectly + # merge separate shells that merely touch or overlap. + left, right, constraints = _manifold_adjacencies(faces64) + ( + local_flip, + component_ids, + singleton_components, + constraint_violations, + ) = _local_orientation(len(faces64), left, right, constraints) + + local_faces = faces64.copy() + local_faces[local_flip] = local_faces[local_flip][:, [0, 2, 1]] + + vertices64 = vertices32.astype(np.float64) + input_triangles = vertices64[faces64] + input_area_normals = np.cross( + input_triangles[:, 1] - input_triangles[:, 0], + input_triangles[:, 2] - input_triangles[:, 0], + ) + degenerate_faces = int( + np.count_nonzero(np.linalg.norm(input_area_normals, axis=1) == 0.0) + ) + bounds_min = vertices64.min(axis=0) + bounds_max = vertices64.max(axis=0) + centre = (bounds_min + bounds_max) * 0.5 + diagonal = float(np.linalg.norm(bounds_max - bounds_min)) + if not np.isfinite(diagonal) or diagonal <= 0.0: + raise ValueError("vertices must span a non-zero bounding box") + score_epsilon = float(score_epsilon_scale * diagonal**3) + + local_scores, _, _ = _radial_statistics(vertices64, local_faces, centre) + component_count = int(component_ids.max()) + 1 + component_scores = np.bincount( + component_ids, + weights=local_scores, + minlength=component_count, + ) + component_score_magnitudes = np.bincount( + component_ids, + weights=np.abs(local_scores), + minlength=component_count, + ) + component_confidence = np.divide( + np.abs(component_scores), + component_score_magnitudes, + out=np.zeros(component_count, dtype=np.float64), + where=component_score_magnitudes > 0.0, + ) + flip_components = component_scores < -score_epsilon + ambiguous_components = np.abs(component_scores) <= score_epsilon + low_confidence_components = component_confidence < confidence_threshold + radial_flip = flip_components[component_ids] + final_flip = local_flip ^ radial_flip + + oriented_faces = faces64.copy() + oriented_faces[final_flip] = oriented_faces[final_flip][:, [0, 2, 1]] + if not np.array_equal( + np.sort(oriented_faces, axis=1), np.sort(faces64, axis=1) + ): + raise RuntimeError("radial orientation changed the triangle multiset") + + _, face_fraction_before, area_fraction_before = _radial_statistics( + vertices64, faces64, centre + ) + _, face_fraction_after, area_fraction_after = _radial_statistics( + vertices64, oriented_faces, centre + ) + report = RadialOrientationReport( + faces=int(len(faces64)), + vertices=int(len(vertices32)), + manifold_adjacencies=int(len(left)), + edge_components=component_count, + singleton_components=singleton_components, + degenerate_faces=degenerate_faces, + duplicate_triangle_occurrences=duplicate_triangle_occurrences, + manifold_winding_conflicts_before=int(np.count_nonzero(constraints)), + manifold_winding_conflicts_after=constraint_violations, + local_faces_flipped=int(np.count_nonzero(local_flip)), + radial_components_flipped=int(np.count_nonzero(flip_components)), + radial_faces_flipped=int(np.count_nonzero(radial_flip)), + ambiguous_components=int(np.count_nonzero(ambiguous_components)), + low_confidence_components=int( + np.count_nonzero(low_confidence_components) + ), + low_confidence_faces=int( + np.count_nonzero(low_confidence_components[component_ids]) + ), + faces_flipped_from_input=int(np.count_nonzero(final_flip)), + radial_score_epsilon=score_epsilon, + confidence_threshold=float(confidence_threshold), + median_component_confidence=float(np.median(component_confidence)), + radially_outward_face_fraction_before=face_fraction_before, + radially_outward_face_fraction_after=face_fraction_after, + radially_outward_area_fraction_before=area_fraction_before, + radially_outward_area_fraction_after=area_fraction_after, + ) + return np.ascontiguousarray(oriented_faces, dtype=face_dtype), report diff --git a/backends/memory.py b/backends/memory.py new file mode 100644 index 0000000..7aa280b --- /dev/null +++ b/backends/memory.py @@ -0,0 +1,62 @@ +"""Unified-memory lifecycle helpers for Apple-Silicon inference.""" + +from __future__ import annotations + +import gc +import os +from typing import Any + + +def memory_snapshot() -> dict[str, float]: + """Return current process and MPS memory counters in GiB.""" + + import psutil + import torch + + snapshot = { + "rss_gib": psutil.Process(os.getpid()).memory_info().rss / 1024**3, + } + if torch.backends.mps.is_available(): + snapshot.update( + { + "mps_allocated_gib": torch.mps.current_allocated_memory() + / 1024**3, + "mps_driver_gib": torch.mps.driver_allocated_memory() + / 1024**3, + } + ) + return snapshot + + +def release_accelerator_memory( + label: str | None = None, + *, + verbose: bool = False, +) -> dict[str, float]: + """Synchronize, collect dead objects and release backend allocator caches.""" + + import torch + + if torch.backends.mps.is_available(): + torch.mps.synchronize() + gc.collect() + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() + snapshot = memory_snapshot() + if verbose and label: + values = ", ".join( + f"{name.removesuffix('_gib')}={value:.2f} GiB" + for name, value in snapshot.items() + ) + print(f"[Memory] {label}: {values}", flush=True) + return snapshot + + +def drop_model(container: Any, name: str) -> None: + """Remove a one-shot model from a pipeline dictionary if present.""" + + models = getattr(container, "models", None) + if isinstance(models, dict): + models.pop(name, None) diff --git a/backends/mesh_extract.py b/backends/mesh_extract.py new file mode 100644 index 0000000..b548dd5 --- /dev/null +++ b/backends/mesh_extract.py @@ -0,0 +1,146 @@ +""" +Pure-Python/PyTorch mesh extraction from sparse voxel dual-grid. + +Replaces the CUDA-only o_voxel._C hashmap operations with Python dicts. +Produces identical output to the CUDA version for inference. +""" + +import torch +import numpy as np +from typing import Union + + +def mesh_to_flexible_dual_grid(*args, **kwargs): + raise RuntimeError("mesh_to_flexible_dual_grid requires CUDA (o_voxel)") + + +# Static lookup tables (lazily initialized, cached per-device) +_edge_neighbor_voxel_offset = None +_quad_split_1 = None +_quad_split_2 = None + + +def flexible_dual_grid_to_mesh( + coords: torch.Tensor, + dual_vertices: torch.Tensor, + intersected_flag: torch.Tensor, + split_weight: Union[torch.Tensor, None], + aabb: Union[list, tuple, np.ndarray, torch.Tensor], + voxel_size: Union[float, list, tuple, np.ndarray, torch.Tensor] = None, + grid_size: Union[int, list, tuple, np.ndarray, torch.Tensor] = None, + train: bool = False, +): + """ + Extract a triangle mesh from sparse voxel dual-grid representation. + + Given a set of voxel coordinates with dual vertex positions and edge + intersection flags, builds quads connecting adjacent voxels at intersected + edges, then splits each quad into two triangles. + + Args: + coords: [N, 3] integer voxel coordinates. + dual_vertices: [N, 3] float vertex offsets within each voxel. + intersected_flag: [N, 3] bool flags indicating which edges are intersected. + split_weight: [N, 1] optional quad split weights (None = use normal alignment). + aabb: [[min_x, min_y, min_z], [max_x, max_y, max_z]] bounding box. + voxel_size: Size of each voxel (alternative to grid_size). + grid_size: Number of voxels per axis (alternative to voxel_size). + train: Must be False (training not supported in pure-Python version). + + Returns: + (vertices, triangles): mesh vertices [V, 3] and face indices [F, 3]. + """ + global _edge_neighbor_voxel_offset, _quad_split_1, _quad_split_2 + + device = coords.device + + if _edge_neighbor_voxel_offset is None or _edge_neighbor_voxel_offset.device != device: + _edge_neighbor_voxel_offset = torch.tensor([ + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], + [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], + [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]], + ], dtype=torch.int, device=device).unsqueeze(0) + _quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device) + _quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device) + + if isinstance(aabb, (list, tuple)): + aabb = np.array(aabb) + if isinstance(aabb, np.ndarray): + aabb = torch.tensor(aabb, dtype=torch.float32, device=device) + + if voxel_size is not None: + if isinstance(voxel_size, (int, float)): + voxel_size = [voxel_size] * 3 + if isinstance(voxel_size, (list, tuple, np.ndarray)): + voxel_size = torch.tensor(np.array(voxel_size), dtype=torch.float32, device=device) + grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int() + else: + if isinstance(grid_size, int): + grid_size = [grid_size] * 3 + if isinstance(grid_size, (list, tuple, np.ndarray)): + grid_size = torch.tensor(np.array(grid_size), dtype=torch.int32, device=device) + voxel_size = (aabb[1] - aabb[0]) / grid_size.float() + + N = dual_vertices.shape[0] + + # Build coordinate lookup on CPU + coords_cpu = coords.cpu() + coord_to_idx = {} + for i in range(N): + key = (coords_cpu[i, 0].item(), coords_cpu[i, 1].item(), coords_cpu[i, 2].item()) + coord_to_idx[key] = i + + # Find connected voxels for each intersected edge + edge_neighbor_voxel = coords.reshape(N, 1, 1, 3) + _edge_neighbor_voxel_offset + connected_voxel = edge_neighbor_voxel[intersected_flag] + M = connected_voxel.shape[0] + + if M == 0: + return torch.zeros(0, 3, device=device), torch.zeros(0, 3, dtype=torch.long, device=device) + + # Look up neighbor indices via dict + connected_cpu = connected_voxel.cpu().reshape(-1, 3) + indices = [] + for j in range(connected_cpu.shape[0]): + key = (connected_cpu[j, 0].item(), connected_cpu[j, 1].item(), connected_cpu[j, 2].item()) + indices.append(coord_to_idx.get(key, 0xFFFFFFFF)) + + connected_voxel_indices = torch.tensor(indices, dtype=torch.int64, device=device).reshape(M, 4) + connected_voxel_valid = (connected_voxel_indices != 0xFFFFFFFF).all(dim=1) + quad_indices = connected_voxel_indices[connected_voxel_valid].long() + L = quad_indices.shape[0] + + if L == 0: + return torch.zeros(0, 3, device=device), torch.zeros(0, 3, dtype=torch.long, device=device) + + # Compute world-space vertex positions + mesh_vertices = (coords.float() + dual_vertices) * voxel_size + aabb[0].reshape(1, 3) + + if train: + raise RuntimeError("Training mode not supported in pure-Python mesh extraction") + + # Triangulate quads: choose the diagonal split that produces better-aligned normals + if split_weight is None: + a1 = quad_indices[:, _quad_split_1] + n0 = torch.cross(mesh_vertices[a1[:, 1]] - mesh_vertices[a1[:, 0]], mesh_vertices[a1[:, 2]] - mesh_vertices[a1[:, 0]]) + n1 = torch.cross(mesh_vertices[a1[:, 2]] - mesh_vertices[a1[:, 1]], mesh_vertices[a1[:, 3]] - mesh_vertices[a1[:, 1]]) + align0 = (n0 * n1).sum(dim=1, keepdim=True).abs() + + a2 = quad_indices[:, _quad_split_2] + n0 = torch.cross(mesh_vertices[a2[:, 1]] - mesh_vertices[a2[:, 0]], mesh_vertices[a2[:, 2]] - mesh_vertices[a2[:, 0]]) + n1 = torch.cross(mesh_vertices[a2[:, 2]] - mesh_vertices[a2[:, 1]], mesh_vertices[a2[:, 3]] - mesh_vertices[a2[:, 1]]) + align1 = (n0 * n1).sum(dim=1, keepdim=True).abs() + + mesh_triangles = torch.where(align0 > align1, a1, a2).reshape(-1, 3) + else: + sw = split_weight[quad_indices] + sw_02 = (sw[:, 0] * sw[:, 2]).squeeze() + sw_13 = (sw[:, 1] * sw[:, 3]).squeeze() + cond = (sw_02 > sw_13).unsqueeze(1).expand(-1, 6) + mesh_triangles = torch.where( + cond, + quad_indices[:, _quad_split_1], + quad_indices[:, _quad_split_2], + ).reshape(-1, 3) + + return mesh_vertices, mesh_triangles diff --git a/backends/mesh_postprocess.py b/backends/mesh_postprocess.py new file mode 100644 index 0000000..5884001 --- /dev/null +++ b/backends/mesh_postprocess.py @@ -0,0 +1,383 @@ +"""Opt-in, non-destructive topology reconstruction for decoded TRELLIS meshes. + +The decoded mesh is always the source of truth. This module can build an UDF +dual-contouring candidate, measure it against the decoded surface, and select +it only when conservative quality gates pass. Callers are expected to keep +the raw mesh even when the candidate is accepted. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +import math +from pathlib import Path +import shutil +import time +from typing import Any + +import numpy as np + + +@dataclass(frozen=True) +class UDFConfig: + """Parameters and conservative acceptance thresholds for UDF remeshing.""" + + resolution: int = 512 + band: float = 1.0 + project_back: float = 0.9 + padding: float = 1.02 + min_largest_component_share: float = 0.99 + max_distance_p95_voxels: float = 1.0 + max_distance_p99_voxels: float = 2.0 + max_bounds_drift_voxels: float = 2.0 + distance_sample_limit: int = 100_000 + + def validate(self) -> None: + if self.resolution < 32 or self.resolution & (self.resolution - 1): + raise ValueError("UDF resolution must be a power of two >= 32") + if not math.isfinite(self.band) or self.band <= 0: + raise ValueError("UDF band must be a finite positive number") + if not math.isfinite(self.project_back) or not 0 <= self.project_back <= 1: + raise ValueError("UDF project_back must be between 0 and 1") + if not math.isfinite(self.padding) or self.padding <= 1: + raise ValueError("UDF padding must be a finite number greater than 1") + if not 0 < self.min_largest_component_share <= 1: + raise ValueError("Largest-component threshold must be in (0, 1]") + if self.distance_sample_limit <= 0: + raise ValueError("Distance sample limit must be positive") + + +@dataclass +class MeshPostprocessResult: + """Candidate plus the geometry selected by the fail-open quality gate.""" + + selected_vertices: np.ndarray + selected_faces: np.ndarray + candidate_vertices: np.ndarray | None + candidate_faces: np.ndarray | None + report: dict[str, Any] + + @property + def accepted(self) -> bool: + return self.report["status"] == "accepted" + + +def _validate_mesh_arrays(vertices: np.ndarray, faces: np.ndarray) -> None: + if vertices.ndim != 2 or vertices.shape[1] != 3 or not len(vertices): + raise ValueError("vertices must be a non-empty [N, 3] array") + if faces.ndim != 2 or faces.shape[1] != 3 or not len(faces): + raise ValueError("faces must be a non-empty [F, 3] array") + if not np.isfinite(vertices).all(): + raise ValueError("vertices contain non-finite values") + if np.min(faces) < 0 or np.max(faces) >= len(vertices): + raise ValueError("faces contain out-of-range vertex indices") + + +def mesh_metrics(vertices: np.ndarray, faces: np.ndarray) -> dict[str, Any]: + """Return topology metrics, explicitly using edge-connected face islands.""" + + import trimesh + + vertices = np.asarray(vertices, dtype=np.float64) + faces = np.asarray(faces, dtype=np.int64) + _validate_mesh_arrays(vertices, faces) + + edges = np.concatenate( + (faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]), axis=0 + ) + edges.sort(axis=1) + unique_edges, edge_counts = np.unique(edges, axis=0, return_counts=True) + boundary_edges = int(np.count_nonzero(edge_counts == 1)) + nonmanifold_edges = int(np.count_nonzero(edge_counts > 2)) + boundary = unique_edges[edge_counts == 1] + + mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) + face_labels = trimesh.graph.connected_component_labels( + np.asarray(mesh.face_adjacency, dtype=np.int64), node_count=len(faces) + ) + face_component_sizes = np.bincount(face_labels) + largest_faces = int(face_component_sizes.max(initial=0)) + + vertex_labels = trimesh.graph.connected_component_labels( + unique_edges, node_count=len(vertices) + ) + referenced_vertices = np.unique(faces.reshape(-1)) + vertex_components = int(len(np.unique(vertex_labels[referenced_vertices]))) + + if len(boundary): + boundary_labels = trimesh.graph.connected_component_labels( + boundary, node_count=len(vertices) + ) + component_ids, edge_component = np.unique( + boundary_labels[boundary[:, 0]], return_inverse=True + ) + component_edge_counts = np.bincount(edge_component) + boundary_degree = np.bincount( + boundary.reshape(-1), minlength=len(vertices) + ) + boundary_vertices = np.flatnonzero(boundary_degree) + vertex_component = np.searchsorted( + component_ids, boundary_labels[boundary_vertices] + ) + non_loop_vertices = np.bincount( + vertex_component, + weights=(boundary_degree[boundary_vertices] != 2), + minlength=len(component_ids), + ) + closed_loops = non_loop_vertices == 0 + small_loops = closed_loops & (component_edge_counts <= 12) + boundary_components = int(len(component_ids)) + closed_boundary_loops = int(np.count_nonzero(closed_loops)) + small_closed_boundary_loops = int(np.count_nonzero(small_loops)) + small_closed_boundary_edges = int(component_edge_counts[small_loops].sum()) + else: + boundary_components = 0 + closed_boundary_loops = 0 + small_closed_boundary_loops = 0 + small_closed_boundary_edges = 0 + + bounds_min = vertices.min(axis=0) + bounds_max = vertices.max(axis=0) + return { + "vertices": int(len(vertices)), + "faces": int(len(faces)), + "boundary_edges": boundary_edges, + "boundary_components": boundary_components, + "closed_boundary_loops": closed_boundary_loops, + "small_closed_boundary_loops": small_closed_boundary_loops, + "small_closed_boundary_edges": small_closed_boundary_edges, + "nonmanifold_edges": nonmanifold_edges, + "edge_components": int(len(face_component_sizes)), + "vertex_components": vertex_components, + "largest_edge_component_faces": largest_faces, + "largest_edge_component_share": float(largest_faces / len(faces)), + "bounds_min": bounds_min.tolist(), + "bounds_max": bounds_max.tolist(), + "extents": (bounds_max - bounds_min).tolist(), + } + + +def _sample_tensor_rows(tensor, limit: int): + if len(tensor) <= limit: + return tensor + import torch + + indices = torch.linspace( + 0, len(tensor) - 1, steps=limit, device=tensor.device + ).long() + return tensor[indices] + + +def _distance_percentiles(bvh, points, voxel_size: float, limit: int) -> dict[str, float]: + sampled = _sample_tensor_rows(points, limit) + distances = bvh.unsigned_distance(sampled)[0].detach().cpu().float().numpy() + in_voxels = distances / voxel_size + return { + "samples": int(len(distances)), + "p95_voxels": float(np.percentile(in_voxels, 95)), + "p99_voxels": float(np.percentile(in_voxels, 99)), + "max_voxels": float(np.max(in_voxels)), + } + + +def evaluate_candidate( + before: dict[str, Any], + after: dict[str, Any], + forward_distance: dict[str, float], + reverse_distance: dict[str, float], + bounds_drift_voxels: list[float], + config: UDFConfig, +) -> list[str]: + """Return rejection reasons; an empty list means the candidate is safe.""" + + reasons: list[str] = [] + before_defects = before["boundary_edges"] + before["nonmanifold_edges"] + after_defects = after["boundary_edges"] + after["nonmanifold_edges"] + + if after["largest_edge_component_share"] < config.min_largest_component_share: + reasons.append( + "largest edge-connected component is below " + f"{config.min_largest_component_share:.1%}" + ) + if before_defects == 0: + if after_defects: + reasons.append("candidate introduces boundary or non-manifold edges") + elif after_defects >= before_defects: + reasons.append("candidate does not reduce total topology defects") + if after["boundary_edges"] > before["boundary_edges"]: + reasons.append("candidate increases boundary edges") + if after.get("small_closed_boundary_loops", 0) > before.get( + "small_closed_boundary_loops", 0 + ): + reasons.append("candidate introduces additional small closed boundary loops") + if after["nonmanifold_edges"] > before["nonmanifold_edges"]: + reasons.append("candidate increases non-manifold edges") + if after["vertex_components"] > before["vertex_components"]: + reasons.append("candidate increases vertex-disconnected components") + + for direction, stats in ( + ("forward", forward_distance), + ("reverse", reverse_distance), + ): + if stats["p95_voxels"] > config.max_distance_p95_voxels: + reasons.append( + f"{direction} surface distance P95 exceeds " + f"{config.max_distance_p95_voxels:g} voxel" + ) + if stats["p99_voxels"] > config.max_distance_p99_voxels: + reasons.append( + f"{direction} surface distance P99 exceeds " + f"{config.max_distance_p99_voxels:g} voxels" + ) + + if max(bounds_drift_voxels, default=0.0) > config.max_bounds_drift_voxels: + reasons.append( + "candidate moves a silhouette bound by more than " + f"{config.max_bounds_drift_voxels:g} voxels" + ) + return reasons + + +def run_udf_postprocess( + vertices: np.ndarray, + faces: np.ndarray, + config: UDFConfig, + *, + verbose: bool = False, +) -> MeshPostprocessResult: + """Build and validate an UDF candidate, falling back to raw on failure.""" + + raw_vertices = np.ascontiguousarray(vertices, dtype=np.float32) + raw_faces = np.ascontiguousarray(faces, dtype=np.int32) + config.validate() + _validate_mesh_arrays(raw_vertices, raw_faces) + started = time.time() + + try: + import torch + # Import through cumesh's platform selector. Importing the physical + # ``cumesh.remeshing`` CUDA module directly bypasses the Darwin alias. + from cumesh import remeshing + from mtlbvh import MtlBVH + + if not torch.backends.mps.is_available(): + raise RuntimeError("UDF mesh reconstruction requires the Metal/MPS backend") + + before = mesh_metrics(raw_vertices, raw_faces) + bounds_min = raw_vertices.min(axis=0) + bounds_max = raw_vertices.max(axis=0) + center_np = (bounds_min + bounds_max) * 0.5 + scale = float(np.max(bounds_max - bounds_min) * config.padding) + voxel_size = scale / config.resolution + + # mtlmesh/mtlbvh dispatch to Metal internally from CPU tensors. Passing + # MPS tensors into the installed metal_hash extension can SIGBUS. + device = torch.device("cpu") + vertices_t = torch.from_numpy(raw_vertices).to(device=device, dtype=torch.float32) + faces_t = torch.from_numpy(raw_faces).to(device=device, dtype=torch.int32) + center_t = torch.from_numpy(center_np).to(device=device, dtype=torch.float32) + source_bvh = MtlBVH(vertices_t, faces_t) + candidate_vertices_t, candidate_faces_t = remeshing.remesh_narrow_band_dc( + vertices_t, + faces_t, + center_t, + scale, + config.resolution, + band=config.band, + project_back=config.project_back, + verbose=verbose, + bvh=source_bvh, + ) + candidate_vertices = candidate_vertices_t.detach().cpu().float().numpy() + candidate_faces = candidate_faces_t.detach().cpu().int().numpy() + _validate_mesh_arrays(candidate_vertices, candidate_faces) + + after = mesh_metrics(candidate_vertices, candidate_faces) + forward = _distance_percentiles( + source_bvh, + candidate_vertices_t, + voxel_size, + config.distance_sample_limit, + ) + candidate_bvh = MtlBVH(candidate_vertices_t, candidate_faces_t) + reverse = _distance_percentiles( + candidate_bvh, + vertices_t, + voxel_size, + config.distance_sample_limit, + ) + + candidate_bounds = np.concatenate( + (candidate_vertices.min(axis=0), candidate_vertices.max(axis=0)) + ) + raw_bounds = np.concatenate((bounds_min, bounds_max)) + bounds_drift = (np.abs(candidate_bounds - raw_bounds) / voxel_size).tolist() + reasons = evaluate_candidate( + before, after, forward, reverse, bounds_drift, config + ) + status = "accepted" if not reasons else "rejected" + report = { + "mode": "udf", + "status": status, + "reasons": reasons, + "config": asdict(config), + "domain": { + "center": center_np.tolist(), + "scale": scale, + "voxel_size": voxel_size, + }, + "before": before, + "after": after, + "surface_distance": {"forward": forward, "reverse": reverse}, + "bounds_drift_voxels": bounds_drift, + "seconds": time.time() - started, + "selected_geometry": "repaired" if status == "accepted" else "raw", + } + return MeshPostprocessResult( + selected_vertices=candidate_vertices if status == "accepted" else raw_vertices, + selected_faces=candidate_faces if status == "accepted" else raw_faces, + candidate_vertices=candidate_vertices, + candidate_faces=candidate_faces, + report=report, + ) + except Exception as exc: + return MeshPostprocessResult( + selected_vertices=raw_vertices, + selected_faces=raw_faces, + candidate_vertices=None, + candidate_faces=None, + report={ + "mode": "udf", + "status": "failed", + "reasons": [f"{type(exc).__name__}: {exc}"], + "config": asdict(config), + "seconds": time.time() - started, + "selected_geometry": "raw", + }, + ) + + +def write_obj(path: str | Path, vertices: np.ndarray, faces: np.ndarray) -> Path: + """Write float32 geometry with enough precision for an exact round trip.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("w", encoding="utf-8") as handle: + for x, y, z in np.asarray(vertices): + handle.write(f"v {x:.9g} {y:.9g} {z:.9g}\n") + for a, b, c in np.asarray(faces, dtype=np.int64): + handle.write(f"f {a + 1} {b + 1} {c + 1}\n") + return output + + +def write_raw_obj_artifacts( + output_prefix: str | Path, vertices: np.ndarray, faces: np.ndarray +) -> tuple[Path, Path]: + """Write canonical ``_raw.obj`` first, then the legacy ``.obj`` copy.""" + + prefix = Path(output_prefix) + raw_path = Path(f"{prefix}_raw.obj") + legacy_path = Path(f"{prefix}.obj") + write_obj(raw_path, vertices, faces) + shutil.copyfile(raw_path, legacy_path) + return raw_path, legacy_path diff --git a/backends/metal_preserve.py b/backends/metal_preserve.py new file mode 100644 index 0000000..628b306 --- /dev/null +++ b/backends/metal_preserve.py @@ -0,0 +1,188 @@ +"""Geometry-preserving adapter for the pinned Metal ``o_voxel`` baker. + +``o_voxel.postprocess.to_glb`` currently performs mesh repair and a second +simplification before UV unwrapping. Those operations can turn overlapping +or non-manifold input geometry into visible holes. This module provides a +small, opt-in compatibility adapter that keeps the backend's UV unwrapping and +normal generation while making the geometry-mutating cleanup calls no-ops. + +The adapter intentionally relies on the private +``o_voxel.postprocess._MeshBackend`` hook from the Apple ``o-voxel`` fork at +commit ``6055b868734af6e12769d229d90580e775fae9f0``. It must be reviewed when +that pin changes. The patch is process-global while the context manager is +active, so it is intended for the repository's serial, single-process CLI. +Concurrent calls to ``to_glb`` in the same process are not supported. + +Preservation applies to the standard ``to_glb(..., remesh=False)`` path. The +remesh path explicitly replaces the mesh through a second ``init`` call and is +therefore incompatible with geometry preservation. +""" + +from __future__ import annotations + +from contextlib import contextmanager +import importlib +import inspect +import threading +from types import ModuleType +from typing import Any, Iterator, TypeVar + + +PINNED_O_VOXEL_COMMIT = "6055b868734af6e12769d229d90580e775fae9f0" + +# Keep this list aligned with the calls in the pinned +# o_voxel.postprocess.to_glb standard path. uv_unwrap, normal generation, +# read, and init deliberately remain inherited from the real backend. +GEOMETRY_MUTATION_METHODS = ( + "fill_holes", + "simplify", + "remove_duplicate_faces", + "repair_non_manifold_edges", + "remove_small_connected_components", + # The pinned orientation pass flips disconnected/non-manifold regions + # independently. On the character regression mesh it reversed 89,980 of + # 199,997 triangles, which Blender displays as holes with backface culling. + # An explicit radial mode, when requested, runs deterministically on the + # bake proxy before this backend and must not be overridden here. + "unify_face_orientations", +) + +_BackendT = TypeVar("_BackendT", bound=type) +_PATCH_LOCK = threading.RLock() + + +class MetalPreserveCompatibilityError(RuntimeError): + """The installed Metal baker does not match this repo's pinned adapter.""" + + +def texture_projection_kwargs( + postprocess_module: ModuleType | Any, + geometry_mode: str, +) -> dict[str, bool]: + """Return the texture-projection override for an ``o_voxel`` bake. + + The historical baker projects rasterized UV positions back through a BVH + because its cleanup and simplification stages can move the surface. The + preserve adapter disables those stages, so the rasterized positions are + already exact points on the immutable bake proxy. Reprojecting them can + select a nearby overlapping/non-manifold sheet and corrupt the texture. + + ``legacy`` intentionally returns no override so upstream behaviour remains + the default. Preserve mode fails closed when the tracked dependency patch + has not been installed into the active Python environment. + """ + + if geometry_mode == "legacy": + return {} + if geometry_mode != "preserve": + raise ValueError(f"Unknown PBR geometry mode: {geometry_mode}") + + to_glb = getattr(postprocess_module, "to_glb", None) + try: + parameters = inspect.signature(to_glb).parameters + except (TypeError, ValueError) as exc: + raise MetalPreserveCompatibilityError( + "Cannot inspect o_voxel.postprocess.to_glb. The preserve texture " + "path requires the repository's pinned o-voxel patch. Re-run " + "./setup.sh before generating." + ) from exc + + if "reproject_texture_to_source" not in parameters: + raise MetalPreserveCompatibilityError( + "The active o_voxel installation does not expose " + "reproject_texture_to_source. Re-run ./setup.sh (or reinstall " + "deps/trellis2-apple/o-voxel into .venv) before using " + "--pbr-geometry-mode preserve." + ) + + return {"reproject_texture_to_source": False} + + +def _ignore_geometry_mutation(self: Any, *args: Any, **kwargs: Any) -> None: + """Match the native mutation methods' call shape without changing state.""" + + del self, args, kwargs + return None + + +def make_geometry_preserving_backend(base_backend: _BackendT) -> _BackendT: + """Return a direct subclass that only neutralizes geometry mutations. + + A dynamic direct subclass is used instead of a proxy so native backend + state and methods such as ``uv_unwrap`` remain untouched. The pinned + ``cumesh.CuMesh`` backend is expected to be subclassable; a clear error is + raised if a future dependency revision changes that contract. + """ + + if not isinstance(base_backend, type): + raise TypeError( + "o_voxel.postprocess._MeshBackend must be a class; " + f"got {type(base_backend).__name__}" + ) + + backend_name = getattr(base_backend, "__name__", "MeshBackend") + namespace = { + "__doc__": ( + f"Geometry-preserving {backend_name} used during Metal texture bake." + ), + "__module__": __name__, + "_trellis_geometry_preserving": True, + "_trellis_original_backend": base_backend, + } + namespace.update( + {method_name: _ignore_geometry_mutation for method_name in GEOMETRY_MUTATION_METHODS} + ) + + try: + preserving_backend = type( + f"GeometryPreserving{backend_name}", + (base_backend,), + namespace, + ) + except TypeError as exc: + raise TypeError( + "The installed o-voxel mesh backend cannot be subclassed. " + f"This adapter supports the pinned commit {PINNED_O_VOXEL_COMMIT}; " + "review the private backend API before updating the dependency." + ) from exc + + return preserving_backend + + +@contextmanager +def use_geometry_preserving_backend( + postprocess_module: ModuleType | Any | None = None, +) -> Iterator[type]: + """Temporarily install a geometry-preserving ``_MeshBackend`` subclass. + + Args: + postprocess_module: ``o_voxel.postprocess``. Supplying it explicitly + avoids an import during tests; when omitted it is imported lazily. + + Yields: + The temporary backend class, primarily for diagnostics and tests. + + The original backend is restored even if baking raises. Nested uses are + supported in one thread. Because the upstream hook is a module global, + all ``to_glb`` calls in a process must remain serial while this context is + active. + """ + + if postprocess_module is None: + postprocess_module = importlib.import_module("o_voxel.postprocess") + + if not hasattr(postprocess_module, "_MeshBackend"): + raise RuntimeError( + "The installed o_voxel.postprocess module has no private " + "_MeshBackend hook. The geometry-preserving adapter requires " + f"the pinned o-voxel commit {PINNED_O_VOXEL_COMMIT}." + ) + + with _PATCH_LOCK: + original_backend = postprocess_module._MeshBackend + preserving_backend = make_geometry_preserving_backend(original_backend) + postprocess_module._MeshBackend = preserving_backend + try: + yield preserving_backend + finally: + postprocess_module._MeshBackend = original_backend diff --git a/backends/multiview.py b/backends/multiview.py new file mode 100644 index 0000000..4c6e01a --- /dev/null +++ b/backends/multiview.py @@ -0,0 +1,550 @@ +"""Experimental spatial multi-view sampling for the released TRELLIS.2 weights. + +The public checkpoint is trained and packaged as a single-image model. This +module does not claim native multi-view support: it evaluates the same model +once per view at every flow step and blends the predicted velocities across +the horizontal voxel grid. Keeping this implementation outside the pinned +TRELLIS.2 checkout makes the experiment opt-in and setup-safe. + +The spatial per-view blending design is inspired by the MIT-licensed +``visualbruno/ComfyUI-Trellis2`` community implementation. This version is a +narrow rewrite for the macOS runtime; it corrects the horizontal-axis mapping +and applies classifier-free guidance once after blending the positive views. + +Azimuth convention (in raw TRELLIS latent/OBJ coordinates, Z-up): + +* 0 degrees: camera on +X +* 90 degrees: camera on +Y +* 180 degrees: camera on -X +* 270 degrees: camera on -Y +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any, Iterator, Sequence + +import numpy as np +import torch +import torch.nn as nn +from easydict import EasyDict as edict +from tqdm import tqdm + +from trellis2.pipelines.samplers.flow_euler import FlowEulerSampler + + +def _validate_view_arguments( + azimuths: Sequence[float], blend_temperature: float +) -> None: + if not azimuths: + raise ValueError("At least one view azimuth is required.") + if not np.isfinite(np.asarray(azimuths, dtype=np.float64)).all(): + raise ValueError("View azimuths must be finite numbers.") + if not np.isfinite(blend_temperature) or blend_temperature <= 0: + raise ValueError("blend_temperature must be a finite positive number.") + + +def _azimuth_directions( + azimuths: Sequence[float], *, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + radians = torch.as_tensor(azimuths, device=device, dtype=torch.float32) + radians = torch.deg2rad(radians) + # Camera direction in the horizontal XY plane; Z remains vertical. + return torch.cos(radians), torch.sin(radians) + + +def dense_view_weights( + shape: Sequence[int], + *, + device: torch.device, + azimuths: Sequence[float], + blend_temperature: float = 2.0, +) -> torch.Tensor: + """Return spatial view weights shaped ``(V, D, H, W)`` for a dense grid.""" + _validate_view_arguments(azimuths, blend_temperature) + if len(shape) != 5: + raise ValueError(f"Expected dense shape (B,C,D,H,W), got {tuple(shape)}") + + depth, height, width = int(shape[2]), int(shape[3]), int(shape[4]) + if min(depth, height, width) <= 0: + raise ValueError(f"Dense spatial dimensions must be positive, got {tuple(shape)}") + + # Voxel centres give a symmetric range without assigning the outermost + # cells exactly to the AABB boundary. + grid_x = (torch.arange(depth, device=device, dtype=torch.float32) + 0.5) + grid_x = grid_x / depth * 2.0 - 1.0 + grid_y = (torch.arange(height, device=device, dtype=torch.float32) + 0.5) + grid_y = grid_y / height * 2.0 - 1.0 + grid_x = grid_x[:, None, None].expand(depth, height, width) + grid_y = grid_y[None, :, None].expand(depth, height, width) + + direction_x, direction_y = _azimuth_directions(azimuths, device=device) + scores = ( + direction_x[:, None, None, None] * grid_x[None] + + direction_y[:, None, None, None] * grid_y[None] + ) + return torch.softmax(scores * blend_temperature, dim=0) + + +def sparse_view_weights( + coords: torch.Tensor, + *, + resolution: int, + azimuths: Sequence[float], + blend_temperature: float = 2.0, +) -> torch.Tensor: + """Return spatial view weights shaped ``(N,V)`` for TRELLIS sparse coords.""" + _validate_view_arguments(azimuths, blend_temperature) + if coords.ndim != 2 or coords.shape[1] != 4: + raise ValueError(f"Expected sparse coordinates (N,4), got {tuple(coords.shape)}") + if resolution <= 0: + raise ValueError("resolution must be positive.") + + # Sparse coordinates map directly to raw mesh (batch, X, Y, Z). + grid_x = (coords[:, 1].float() + 0.5) / resolution * 2.0 - 1.0 + grid_y = (coords[:, 2].float() + 0.5) / resolution * 2.0 - 1.0 + direction_x, direction_y = _azimuth_directions( + azimuths, device=coords.device + ) + scores = ( + grid_x[:, None] * direction_x[None] + + grid_y[:, None] * direction_y[None] + ) + return torch.softmax(scores * blend_temperature, dim=1) + + +class SpatialMultiViewFlowEulerSampler(FlowEulerSampler): + """Euler flow sampler that blends per-view velocities in one 3D latent.""" + + def __init__(self, sigma_min: float, resolution: int): + super().__init__(sigma_min) + if resolution <= 0: + raise ValueError("resolution must be positive.") + self.resolution = resolution + + @torch.no_grad() + def sample_once( + self, + model, + x_t, + t: float, + t_prev: float, + *, + conditions: Sequence[dict[str, Any]], + azimuths: Sequence[float], + blend_temperature: float = 2.0, + **kwargs, + ) -> edict: + if len(conditions) != len(azimuths): + raise ValueError( + f"Got {len(conditions)} conditions for {len(azimuths)} azimuths." + ) + + is_sparse = hasattr(x_t, "coords") and hasattr(x_t, "feats") + if is_sparse: + weights = sparse_view_weights( + x_t.coords, + resolution=self.resolution, + azimuths=azimuths, + blend_temperature=blend_temperature, + ) + else: + weights = dense_view_weights( + x_t.shape, + device=x_t.device, + azimuths=azimuths, + blend_temperature=blend_temperature, + ) + + guidance_strength = kwargs.pop("guidance_strength") + guidance_interval = kwargs.pop("guidance_interval") + guidance_rescale = kwargs.pop("guidance_rescale", 0.0) + + # Blend raw positive-conditioned velocities first, then apply CFG once. + # This preserves the released sampler semantics, particularly its + # guidance_rescale calculation. Applying CFG independently per view + # before blending is only equivalent when guidance_rescale == 0. + accumulated = None + for index, condition in enumerate(conditions): + pred_view = FlowEulerSampler._inference_model( + self, + model, + x_t, + t, + cond=condition["cond"], + **kwargs, + ) + pred_feats = pred_view.feats if is_sparse else pred_view + if is_sparse: + weight = weights[:, index].unsqueeze(1) + else: + weight = weights[index].unsqueeze(0).unsqueeze(0) + weighted = pred_feats * weight.to(dtype=pred_feats.dtype) + accumulated = weighted if accumulated is None else accumulated + weighted + + pred_pos = x_t.replace(feats=accumulated) if is_sparse else accumulated + + effective_strength = ( + guidance_strength + if guidance_interval[0] <= t <= guidance_interval[1] + else 1.0 + ) + if effective_strength == 1: + pred_v = pred_pos + else: + pred_neg = FlowEulerSampler._inference_model( + self, + model, + x_t, + t, + cond=conditions[0]["neg_cond"], + **kwargs, + ) + if effective_strength == 0: + pred_v = pred_neg + else: + pred_v = ( + effective_strength * pred_pos + + (1 - effective_strength) * pred_neg + ) + if guidance_rescale > 0: + x_0_pos = self._pred_to_xstart(x_t, t, pred_pos) + x_0_cfg = self._pred_to_xstart(x_t, t, pred_v) + dims = list(range(1, x_0_pos.ndim)) + std_pos = x_0_pos.std(dim=dims, keepdim=True) + std_cfg = x_0_cfg.std(dim=dims, keepdim=True) + x_0_rescaled = x_0_cfg * (std_pos / std_cfg) + x_0 = ( + guidance_rescale * x_0_rescaled + + (1 - guidance_rescale) * x_0_cfg + ) + pred_v = self._xstart_to_pred(x_t, t, x_0) + + pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v) + pred_x_prev = x_t - (t - t_prev) * pred_v + return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}) + + @torch.no_grad() + def sample( + self, + model, + noise, + *, + conditions: Sequence[dict[str, Any]], + azimuths: Sequence[float], + steps: int = 50, + rescale_t: float = 1.0, + guidance_strength: float = 3.0, + guidance_interval: tuple[float, float] = (0.0, 1.0), + guidance_rescale: float = 0.0, + blend_temperature: float = 2.0, + verbose: bool = True, + tqdm_desc: str = "Sampling multi-view", + **kwargs, + ) -> edict: + if len(conditions) != len(azimuths): + raise ValueError( + f"Got {len(conditions)} conditions for {len(azimuths)} azimuths." + ) + _validate_view_arguments(azimuths, blend_temperature) + + sample = noise + t_seq = np.linspace(1, 0, steps + 1) + t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq) + t_pairs = list(zip(t_seq[:-1].tolist(), t_seq[1:].tolist())) + ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []}) + for t, t_prev in tqdm(t_pairs, desc=tqdm_desc, disable=not verbose): + out = self.sample_once( + model, + sample, + t, + t_prev, + conditions=conditions, + azimuths=azimuths, + blend_temperature=blend_temperature, + guidance_strength=guidance_strength, + guidance_interval=guidance_interval, + guidance_rescale=guidance_rescale, + **kwargs, + ) + sample = out.pred_x_prev + ret.pred_x_t.append(out.pred_x_prev) + ret.pred_x_0.append(out.pred_x_0) + ret.samples = sample + return ret + + +class SpatialMultiViewGuidanceIntervalSampler(SpatialMultiViewFlowEulerSampler): + """Spatial multi-view Euler sampling with TRELLIS CFG interval semantics.""" + + +def _make_sampler(base_sampler, resolution: int): + return SpatialMultiViewGuidanceIntervalSampler( + sigma_min=base_sampler.sigma_min, + resolution=resolution, + ) + + +@contextmanager +def _model_on_pipeline_device(pipeline, model) -> Iterator[Any]: + if pipeline.low_vram: + model.to(pipeline.device) + try: + yield model + finally: + if pipeline.low_vram: + model.cpu() + + +def _denormalize_slat(slat, normalization: dict[str, Sequence[float]]): + std = torch.tensor(normalization["std"], device=slat.device)[None] + mean = torch.tensor(normalization["mean"], device=slat.device)[None] + return slat * std + mean + + +def _normalize_slat(slat, normalization: dict[str, Sequence[float]]): + std = torch.tensor(normalization["std"], device=slat.device)[None] + mean = torch.tensor(normalization["mean"], device=slat.device)[None] + return (slat - mean) / std + + +def _sample_sparse_structure( + pipeline, + conditions, + azimuths, + *, + resolution: int, + sampler_params: dict[str, Any], + blend_temperature: float, +): + flow_model = pipeline.models["sparse_structure_flow_model"] + model_resolution = flow_model.resolution + noise = torch.randn( + 1, + flow_model.in_channels, + model_resolution, + model_resolution, + model_resolution, + device=pipeline.device, + ) + params = {**pipeline.sparse_structure_sampler_params, **sampler_params} + sampler = _make_sampler(pipeline.sparse_structure_sampler, resolution) + with _model_on_pipeline_device(pipeline, flow_model): + z_s = sampler.sample( + flow_model, + noise, + conditions=conditions, + azimuths=azimuths, + blend_temperature=blend_temperature, + **params, + verbose=True, + tqdm_desc="Sampling sparse structure (multi-view)", + ).samples + + decoder = pipeline.models["sparse_structure_decoder"] + with _model_on_pipeline_device(pipeline, decoder): + decoded = decoder(z_s) > 0 + if resolution != decoded.shape[2]: + ratio = decoded.shape[2] // resolution + decoded = torch.nn.functional.max_pool3d( + decoded.float(), ratio, ratio, 0 + ) > 0.5 + return torch.argwhere(decoded)[:, [0, 2, 3, 4]].int() + + +def _sample_shape_slat( + pipeline, + conditions, + azimuths, + *, + flow_model, + coords: torch.Tensor, + coordinate_resolution: int, + sampler_params: dict[str, Any], + blend_temperature: float, +): + from trellis2.modules.sparse import SparseTensor + + noise = SparseTensor( + feats=torch.randn( + coords.shape[0], flow_model.in_channels, device=pipeline.device + ), + coords=coords, + ) + params = {**pipeline.shape_slat_sampler_params, **sampler_params} + sampler = _make_sampler(pipeline.shape_slat_sampler, coordinate_resolution) + with _model_on_pipeline_device(pipeline, flow_model): + slat = sampler.sample( + flow_model, + noise, + conditions=conditions, + azimuths=azimuths, + blend_temperature=blend_temperature, + **params, + verbose=True, + tqdm_desc="Sampling shape SLat (multi-view)", + ).samples + return _denormalize_slat(slat, pipeline.shape_slat_normalization) + + +def _sample_tex_slat( + pipeline, + conditions, + azimuths, + *, + flow_model, + shape_slat, + coordinate_resolution: int, + sampler_params: dict[str, Any], + blend_temperature: float, +): + normalized_shape = _normalize_slat( + shape_slat, pipeline.shape_slat_normalization + ) + in_channels = ( + flow_model.in_channels + if isinstance(flow_model, nn.Module) + else flow_model[0].in_channels + ) + noise = normalized_shape.replace( + feats=torch.randn( + normalized_shape.coords.shape[0], + in_channels - normalized_shape.feats.shape[1], + device=pipeline.device, + ) + ) + params = {**pipeline.tex_slat_sampler_params, **sampler_params} + sampler = _make_sampler(pipeline.tex_slat_sampler, coordinate_resolution) + with _model_on_pipeline_device(pipeline, flow_model): + slat = sampler.sample( + flow_model, + noise, + conditions=conditions, + azimuths=azimuths, + blend_temperature=blend_temperature, + concat_cond=normalized_shape, + **params, + verbose=True, + tqdm_desc="Sampling texture SLat (multi-view)", + ).samples + return _denormalize_slat(slat, pipeline.tex_slat_normalization) + + +def select_texture_views( + conditions: Sequence[dict[str, Any]], + azimuths: Sequence[float], + mode: str, +) -> tuple[Sequence[dict[str, Any]], Sequence[float]]: + """Select appearance conditioning independently from geometry views. + + ``primary`` keeps the shared multi-view geometry trajectory but conditions + the texture latent globally from the positional image only. This avoids + baking disagreements between turnaround images into the appearance latent; + it is still inference by the released mono-image model, not camera-aware + texture projection. + """ + + if len(conditions) != len(azimuths): + raise ValueError( + f"Got {len(conditions)} texture conditions for {len(azimuths)} azimuths." + ) + if mode == "blend": + return conditions, azimuths + if mode == "primary": + return conditions[:1], azimuths[:1] + raise ValueError(f"Unknown multi-view texture mode: {mode}") + + +@torch.no_grad() +def run_multiview( + pipeline, + images: Sequence[Any], + azimuths: Sequence[float], + *, + seed: int = 42, + pipeline_type: str = "512", + sparse_structure_sampler_params: dict[str, Any] | None = None, + shape_slat_sampler_params: dict[str, Any] | None = None, + tex_slat_sampler_params: dict[str, Any] | None = None, + blend_temperature: float = 2.0, + texture_mode: str = "blend", +): + """Run one shared TRELLIS diffusion trajectory conditioned by many views.""" + if len(images) != len(azimuths): + raise ValueError(f"Got {len(images)} images for {len(azimuths)} azimuths.") + if len(images) < 2: + raise ValueError("Multi-view generation requires at least two images.") + _validate_view_arguments(azimuths, blend_temperature) + if pipeline_type not in {"512", "1024"}: + raise ValueError( + "Experimental multi-view generation currently supports only " + f"the direct 512 and 1024 pipelines, got {pipeline_type}." + ) + + sparse_overrides = sparse_structure_sampler_params or {} + shape_overrides = shape_slat_sampler_params or {} + texture_overrides = tex_slat_sampler_params or {} + processed = [pipeline.preprocess_image(image) for image in images] + + torch.manual_seed(seed) + conditions_512 = [pipeline.get_cond([image], 512) for image in processed] + conditions_1024 = ( + [pipeline.get_cond([image], 1024) for image in processed] + if pipeline_type != "512" + else None + ) + + sparse_resolution = {"512": 32, "1024": 64}[pipeline_type] + coords = _sample_sparse_structure( + pipeline, + conditions_512, + azimuths, + resolution=sparse_resolution, + sampler_params=sparse_overrides, + blend_temperature=blend_temperature, + ) + + if pipeline_type == "512": + resolution = 512 + shape_slat = _sample_shape_slat( + pipeline, + conditions_512, + azimuths, + flow_model=pipeline.models["shape_slat_flow_model_512"], + coords=coords, + coordinate_resolution=32, + sampler_params=shape_overrides, + blend_temperature=blend_temperature, + ) + texture_conditions = conditions_512 + texture_model = pipeline.models["tex_slat_flow_model_512"] + else: + resolution = 1024 + shape_slat = _sample_shape_slat( + pipeline, + conditions_1024, + azimuths, + flow_model=pipeline.models["shape_slat_flow_model_1024"], + coords=coords, + coordinate_resolution=64, + sampler_params=shape_overrides, + blend_temperature=blend_temperature, + ) + texture_conditions = conditions_1024 + texture_model = pipeline.models["tex_slat_flow_model_1024"] + texture_conditions, texture_azimuths = select_texture_views( + texture_conditions, azimuths, texture_mode + ) + tex_slat = _sample_tex_slat( + pipeline, + texture_conditions, + texture_azimuths, + flow_model=texture_model, + shape_slat=shape_slat, + coordinate_resolution=resolution // 16, + sampler_params=texture_overrides, + blend_temperature=blend_temperature, + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return pipeline.decode_latent(shape_slat, tex_slat, resolution) diff --git a/backends/naf_attention.py b/backends/naf_attention.py new file mode 100644 index 0000000..b948984 --- /dev/null +++ b/backends/naf_attention.py @@ -0,0 +1,271 @@ +"""Memory-bounded, exact PyTorch implementation of NAF's 2-D attention. + +NAF normally calls NATTEN's CUDA CUTLASS kernel. The previous macOS port +replaced the whole learned upsampler with bilinear interpolation when that +kernel was unavailable, which changes the conditioning consumed by Pixal3D. + +This module preserves NATTEN's neighborhood definition and learned NAF model. +Only the attention evaluation is tiled by query rows so its temporary +``K x K`` key/value neighborhoods fit in unified memory. +""" + +from __future__ import annotations + +import math +import os +from types import MethodType +from typing import Any + +import torch + + +def _as_pair(value: int | tuple[int, int]) -> tuple[int, int]: + if isinstance(value, int): + return value, value + if len(value) != 2: + raise ValueError(f"Expected a 2-D value, got {value!r}") + return int(value[0]), int(value[1]) + + +def neighborhood_indices( + length: int, + kernel_size: int, + dilation: int, + *, + device: torch.device, +) -> torch.Tensor: + """Return NATTEN-compatible neighbor indices for every query coordinate. + + NATTEN shifts a boundary window inward instead of padding it. Dilation + partitions coordinates into independent modulo groups; each group receives + the same shifted ``kernel_size`` window. + """ + + if kernel_size <= 0 or kernel_size % 2 == 0: + raise ValueError("Only positive odd neighborhood sizes are supported") + if dilation <= 0: + raise ValueError("Dilation must be positive") + if kernel_size * dilation > length: + raise ValueError( + f"kernel_size * dilation ({kernel_size * dilation}) exceeds " + f"the input length ({length})" + ) + + coordinates = torch.arange(length, device=device, dtype=torch.long) + group_coordinate = torch.div(coordinates, dilation, rounding_mode="floor") + dilation_group = torch.remainder(coordinates, dilation) + # Number of valid coordinates in the query's dilation group. This is the + # integer form of NATTEN's qkv_shape_corrected boundary calculation. + group_length = torch.div( + length - 1 - dilation_group, + dilation, + rounding_mode="floor", + ) + 1 + + left = kernel_size // 2 + right = kernel_size - left - 1 + window_center = torch.minimum( + torch.maximum(group_coordinate, torch.full_like(group_coordinate, left)), + group_length - 1 - right, + ) + offsets = torch.arange( + -left, + right + 1, + device=device, + dtype=torch.long, + ) + return ( + (window_center[:, None] + offsets[None, :]) * dilation + + dilation_group[:, None] + ) + + +def chunked_na2d( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + kernel_size: int | tuple[int, int], + dilation: int | tuple[int, int] = 1, + scale: float | None = None, + chunk_rows: int | None = None, +) -> torch.Tensor: + """Evaluate NATTEN-compatible 2-D attention in bounded row chunks. + + Tensors use NATTEN's heads-last layout ``[B, H, W, heads, head_dim]``. + The function is intended for frozen inference and therefore deliberately + rejects autograd inputs. + """ + + if any(t.ndim != 5 for t in (query, key, value)): + raise ValueError("query, key and value must all have shape [B,H,W,heads,dim]") + if query.shape != key.shape: + raise ValueError( + "NAF query and key shapes must match, got " + f"{tuple(query.shape)} and {tuple(key.shape)}" + ) + if value.shape[:-1] != query.shape[:-1]: + raise ValueError( + "NAF value must match query through its head dimension, got " + f"{tuple(query.shape)} and {tuple(value.shape)}" + ) + if any(t.requires_grad for t in (query, key, value)): + raise RuntimeError("chunked_na2d is an inference-only implementation") + + kernel_y, kernel_x = _as_pair(kernel_size) + dilation_y, dilation_x = _as_pair(dilation) + batch, height, width, heads, head_dim = query.shape + value_head_dim = value.shape[-1] + neighborhood = kernel_y * kernel_x + scale = float(scale if scale is not None else head_dim**-0.5) + + y_indices = neighborhood_indices( + height, + kernel_y, + dilation_y, + device=query.device, + ) + x_indices = neighborhood_indices( + width, + kernel_x, + dilation_x, + device=query.device, + ) + + if chunk_rows is None: + requested = int(os.environ.get("PIXAL3D_NAF_CHUNK_ROWS", "0")) + if requested > 0: + chunk_rows = requested + else: + # Keep one gathered K/V neighborhood around 384 MiB in fp32. + # At NAF's 1024²/4-head/64-dim shape this resolves to four rows. + max_neighbor_values = 96 * 1024**2 + values_per_row = max( + 1, + batch + * width + * neighborhood + * heads + * max(head_dim, value_head_dim), + ) + chunk_rows = max(1, max_neighbor_values // values_per_row) + chunk_rows = max(1, min(height, int(chunk_rows))) + + output = torch.empty( + batch, + height, + width, + heads, + value_head_dim, + dtype=value.dtype, + device=value.device, + ) + for row_start in range(0, height, chunk_rows): + row_end = min(height, row_start + chunk_rows) + rows = row_end - row_start + y_chunk = y_indices[row_start:row_end] + + y_grid = y_chunk[:, None, :, None].expand( + rows, + width, + kernel_y, + kernel_x, + ) + x_grid = x_indices[None, :, None, :].expand( + rows, + width, + kernel_y, + kernel_x, + ) + y_flat = y_grid.reshape(rows, width, neighborhood) + x_flat = x_grid.reshape(rows, width, neighborhood) + + query_chunk = query[:, row_start:row_end] + key_neighborhood = key[:, y_flat, x_flat] + logits = torch.einsum( + "brwhd,brwkhd->brwhk", + query_chunk, + key_neighborhood, + ) + del key_neighborhood + weights = torch.softmax(logits * scale, dim=-1) + del logits + + value_neighborhood = value[:, y_flat, x_flat] + output[:, row_start:row_end] = torch.einsum( + "brwhk,brwkhd->brwhd", + weights, + value_neighborhood, + ) + del value_neighborhood, weights, y_grid, x_grid, y_flat, x_flat + + return output + + +def install_chunked_naf_attention(naf_model: Any) -> None: + """Replace only a loaded NAF model's CUDA-only attention operation.""" + + upsampler = getattr(naf_model, "upsampler", None) + if upsampler is None or not hasattr(upsampler, "_resize"): + raise TypeError("Loaded NAF model has an unsupported upsampler") + if getattr(upsampler, "_pixal3d_chunked_attention", False): + return + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + image=None, + return_weights: bool = False, + **kwargs, + ) -> torch.Tensor: + del image, kwargs + if return_weights: + raise NotImplementedError( + "Chunked NAF inference does not materialize attention weights" + ) + height, width = q.shape[-2:] + key_height, key_width = k.shape[-2:] + dilation = (height // key_height, width // key_width) + if ( + height % key_height + or width % key_width + or min(dilation) <= 0 + ): + raise ValueError( + "NAF target size must be an integer multiple of its feature map" + ) + + batch, query_channels, _, _ = q.shape + query_head_dim = query_channels // self.num_heads + query = ( + q.reshape( + batch, + self.num_heads, + query_head_dim, + height, + width, + ) + .permute(0, 3, 4, 1, 2) + .contiguous() + ) + key_resized = self._resize(k, size=(height, width), dtype=query.dtype) + value_resized = self._resize(v, size=(height, width), dtype=query.dtype) + result = chunked_na2d( + query, + key_resized, + value_resized, + kernel_size=self.kernel_size, + dilation=dilation, + scale=self.scale, + ) + output_channels = self.num_heads * result.shape[-1] + return ( + result.permute(0, 3, 4, 1, 2) + .reshape(batch, output_channels, height, width) + .contiguous() + ) + + upsampler.forward = MethodType(forward, upsampler) + upsampler._pixal3d_chunked_attention = True diff --git a/backends/sparse_grid_sample.py b/backends/sparse_grid_sample.py new file mode 100644 index 0000000..942a780 --- /dev/null +++ b/backends/sparse_grid_sample.py @@ -0,0 +1,122 @@ +"""Dense fallback for sparse 3D feature sampling. + +The sparse ``flex_gemm`` sampler treats integer sparse coordinates as voxel +indices whose centres live at ``coord + 0.5``. Trilinear interpolation also +ignores absent sparse neighbours and renormalizes the weights that remain. +Sampling a zero-filled dense feature volume alone does not preserve that +second property, so this fallback samples an occupancy volume in parallel and +uses it as the valid-weight denominator. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import torch.nn.functional as F + + +def dense_sparse_grid_sample_3d( + feats: torch.Tensor, + coords: torch.Tensor, + shape: Sequence[int], + grid: torch.Tensor, + mode: str = "trilinear", +) -> torch.Tensor: + """Sample sparse voxel features through a temporary dense volume. + + This is a compatibility fallback for + ``flex_gemm.ops.grid_sample.grid_sample_3d`` in the texture-baking path. + It intentionally returns ``(B * M, C)`` because that is the two-dimensional + shape consumed by ``o_voxel.postprocess.to_glb``. + + Args: + feats: Sparse features shaped ``(N, C)``. + coords: Sparse integer coordinates shaped ``(N, 4)`` as + ``(batch, x, y, z)``. + shape: Logical sparse tensor shape ``(B, C, D, H, W)``. TRELLIS uses + the three spatial entries for its x, y, and z axes respectively. + grid: Query points shaped ``(B, M, 3)`` in voxel-corner coordinates. + Therefore voxel index zero has its centre at query coordinate 0.5. + mode: Only ``"trilinear"`` is supported by this bake fallback. + + Returns: + Sampled features shaped ``(B * M, C)``. Queries with no occupied + sparse neighbour return exactly zero. + """ + + if mode != "trilinear": + raise ValueError( + "dense sparse grid sampling only supports mode='trilinear'; " + f"got {mode!r}" + ) + if feats.ndim != 2: + raise ValueError(f"feats must have shape (N, C), got {tuple(feats.shape)}") + if coords.ndim != 2 or coords.shape[1] != 4: + raise ValueError(f"coords must have shape (N, 4), got {tuple(coords.shape)}") + if grid.ndim != 3 or grid.shape[2] != 3: + raise ValueError(f"grid must have shape (B, M, 3), got {tuple(grid.shape)}") + if len(shape) != 5: + raise ValueError(f"shape must contain (B, C, D, H, W), got {tuple(shape)}") + if feats.shape[0] != coords.shape[0]: + raise ValueError("feats and coords must contain the same number of voxels") + + B, C, D, H, W = (int(value) for value in shape) + if grid.shape[0] != B: + raise ValueError(f"grid batch {grid.shape[0]} does not match shape batch {B}") + if feats.shape[1] != C: + raise ValueError(f"feature channels {feats.shape[1]} do not match shape channels {C}") + if min(D, H, W) < 2: + raise ValueError("dense sparse grid sampling requires spatial dimensions >= 2") + + device = feats.device + dense_features = torch.zeros( + (B, C, D, H, W), dtype=feats.dtype, device=device + ) + occupancy = torch.zeros((B, 1, D, H, W), dtype=feats.dtype, device=device) + + sparse_coords = coords.to(device=device, dtype=torch.long) + batch_idx = sparse_coords[:, 0] + coord_x = sparse_coords[:, 1] + coord_y = sparse_coords[:, 2] + coord_z = sparse_coords[:, 3] + dense_features[batch_idx, :, coord_x, coord_y, coord_z] = feats + occupancy[batch_idx, 0, coord_x, coord_y, coord_z] = 1 + + # PyTorch indexes dense samples at integer centres, whereas flex_gemm uses + # centres at sparse_coord + 0.5. Shift before converting to the normalized + # z/y/x order expected by five-dimensional F.grid_sample. + centred_grid = grid.to(device=device, dtype=feats.dtype) - 0.5 + normalized_grid = torch.stack( + [ + centred_grid[..., 2] / (W - 1) * 2 - 1, + centred_grid[..., 1] / (H - 1) * 2 - 1, + centred_grid[..., 0] / (D - 1) * 2 - 1, + ], + dim=-1, + ).reshape(B, 1, 1, -1, 3) + + sample_kwargs = { + "mode": "bilinear", + "align_corners": True, + "padding_mode": "zeros", + } + sampled_features = F.grid_sample(dense_features, normalized_grid, **sample_kwargs) + valid_weight = F.grid_sample(occupancy, normalized_grid, **sample_kwargs) + + # The sparse kernel divides by the sum of weights belonging to present, + # in-bounds neighbours. Keep unsupported samples exactly zero rather than + # introducing NaN/Inf through a zero denominator. + supported = valid_weight > 1e-12 + sampled_features = torch.where( + supported, + sampled_features / valid_weight.clamp_min(1e-12), + torch.zeros_like(sampled_features), + ) + + sample_count = grid.shape[1] + return ( + sampled_features.reshape(B, C, sample_count) + .permute(0, 2, 1) + .reshape(B * sample_count, C) + ) diff --git a/backends/stubs.py b/backends/stubs.py new file mode 100644 index 0000000..0aa9e15 --- /dev/null +++ b/backends/stubs.py @@ -0,0 +1,95 @@ +""" +Stub modules for CUDA-only libraries that TRELLIS.2 imports. + +These provide graceful error messages instead of ImportError crashes, +allowing the rest of the pipeline to run on MPS/CPU. + +Usage: + Call install_stubs(stubs_dir) to create the stub package structure, + or add the stubs directory to sys.path before importing TRELLIS.2. +""" + +import os + + +def install_stubs(stubs_dir): + """Create stub package files in the given directory.""" + os.makedirs(stubs_dir, exist_ok=True) + + # cumesh + _write(os.path.join(stubs_dir, "cumesh.py"), '''\ +"""Stub for cumesh — CUDA mesh operations (hole filling, simplification).""" + +class _Stub: + def __getattr__(self, name): + raise AttributeError(f"cumesh.{name} not available (CUDA required)") + +import sys +sys.modules[__name__] = _Stub() +''') + + # flex_gemm (top-level module + ops subpackage) + _write(os.path.join(stubs_dir, "flex_gemm.py"), '''\ +"""Stub for flex_gemm — CUDA sparse convolution kernels.""" + +class _Stub: + def __getattr__(self, name): + raise RuntimeError(f"flex_gemm.{name} requires CUDA.") + +import sys +sys.modules[__name__] = _Stub() +''') + + fg_ops = os.path.join(stubs_dir, "flex_gemm", "ops") + os.makedirs(fg_ops, exist_ok=True) + _write(os.path.join(stubs_dir, "flex_gemm", "__init__.py"), "pass\n") + _write(os.path.join(fg_ops, "__init__.py"), "pass\n") + _write(os.path.join(fg_ops, "grid_sample.py"), '''\ +def grid_sample_3d(*args, **kwargs): + raise RuntimeError("flex_gemm requires CUDA") +''') + + # nvdiffrast + nv_dir = os.path.join(stubs_dir, "nvdiffrast") + os.makedirs(nv_dir, exist_ok=True) + _write(os.path.join(stubs_dir, "nvdiffrast.py"), '"""Stub for nvdiffrast."""\npass\n') + _write(os.path.join(nv_dir, "__init__.py"), "pass\n") + _write(os.path.join(nv_dir, "torch.py"), '''\ +def RasterizeCudaContext(*args, **kwargs): + raise RuntimeError("nvdiffrast requires CUDA") +''') + + # o_voxel (with real mesh extraction in convert.py) + ov_dir = os.path.join(stubs_dir, "o_voxel") + os.makedirs(ov_dir, exist_ok=True) + _write(os.path.join(ov_dir, "__init__.py"), "pass\n") + _write(os.path.join(ov_dir, "io.py"), '''\ +def read(*args, **kwargs): + raise RuntimeError("o_voxel.io requires CUDA") + +def write(*args, **kwargs): + raise RuntimeError("o_voxel.io requires CUDA") + +def read_vxz(*args, **kwargs): + raise RuntimeError("o_voxel.io requires CUDA") +''') + _write(os.path.join(ov_dir, "rasterize.py"), '''\ +class VoxelRenderer: + def __init__(self, *args, **kwargs): + raise RuntimeError("o_voxel.rasterize requires CUDA") +''') + # Note: o_voxel/convert.py is provided by backends/mesh_extract.py + # and copied into place by the patch script. + + +def _write(path, content): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as f: + f.write(content) + + +if __name__ == "__main__": + import sys + target = sys.argv[1] if len(sys.argv) > 1 else "stubs" + install_stubs(target) + print(f"Stubs installed to {target}/") diff --git a/backends/texture_baker.py b/backends/texture_baker.py new file mode 100644 index 0000000..8e222a8 --- /dev/null +++ b/backends/texture_baker.py @@ -0,0 +1,326 @@ +""" +UV unwrap + texture baking for TRELLIS.2 meshes on Apple Silicon. + +Replaces nvdiffrast (CUDA-only) with: + - xatlas for UV unwrapping (C++ library, CPU) + - Vectorized numpy rasterizer for UV-space triangles + - scipy cKDTree for nearest-voxel lookup at native 512 resolution + - Inverse-distance weighting for trilinear-like interpolation + +Produces GLB files with PBR textures (base color, metallic, roughness). +""" + +import numpy as np +import time + + +def uv_unwrap(vertices, faces): + """ + Compute UV coordinates for a mesh using xatlas. + + Returns: + new_vertices: Remapped vertices + new_faces: Triangle indices into new_vertices + uvs: UV coordinates per new vertex, in [0, 1] + vmapping: Maps new vertex index -> original vertex index + """ + import xatlas + + v = np.ascontiguousarray(vertices.astype(np.float32)) + f = np.ascontiguousarray(faces.astype(np.uint32)) + + vmapping, indices, uvs = xatlas.parametrize(v, f) + + new_vertices = v[vmapping] + new_faces = indices.reshape(-1, 3) + + return new_vertices, new_faces, uvs, vmapping + + +def _rasterize_uv_triangles(vertices, faces, uvs, texture_size): + """ + Rasterize all triangles in UV space. For each texel, determine which + triangle covers it and compute 3D position via barycentric interpolation. + + Args: + vertices: [N, 3] mesh vertices + faces: [F, 3] triangle indices + uvs: [N, 2] UV coordinates in [0, 1] + texture_size: output texture resolution + + Returns: + positions: [H, W, 3] 3D position at each texel + mask: [H, W] bool mask of filled texels + """ + H = W = texture_size + positions = np.zeros((H, W, 3), dtype=np.float32) + mask = np.zeros((H, W), dtype=bool) + + n_faces = len(faces) + uv_scale = np.array([W - 1, H - 1], dtype=np.float32) + + for fi in range(n_faces): + if fi > 0 and fi % 100000 == 0: + print(f" Rasterizing: {fi:,}/{n_faces:,}") + + i0, i1, i2 = faces[fi] + uv0 = uvs[i0] * uv_scale + uv1 = uvs[i1] * uv_scale + uv2 = uvs[i2] * uv_scale + p0, p1, p2 = vertices[i0], vertices[i1], vertices[i2] + + min_x = max(int(np.floor(min(uv0[0], uv1[0], uv2[0]))), 0) + max_x = min(int(np.ceil(max(uv0[0], uv1[0], uv2[0]))), W - 1) + min_y = max(int(np.floor(min(uv0[1], uv1[1], uv2[1]))), 0) + max_y = min(int(np.ceil(max(uv0[1], uv1[1], uv2[1]))), H - 1) + + if max_x < min_x or max_y < min_y: + continue + + d00 = uv1[0] - uv0[0] + d01 = uv2[0] - uv0[0] + d10 = uv1[1] - uv0[1] + d11 = uv2[1] - uv0[1] + denom = d00 * d11 - d01 * d10 + if abs(denom) < 1e-10: + continue + inv_denom = 1.0 / denom + + px_range = np.arange(min_x, max_x + 1, dtype=np.float32) + py_range = np.arange(min_y, max_y + 1, dtype=np.float32) + if len(px_range) == 0 or len(py_range) == 0: + continue + + px_grid, py_grid = np.meshgrid(px_range, py_range) + dx = px_grid - uv0[0] + dy = py_grid - uv0[1] + + u = (dx * d11 - d01 * dy) * inv_denom + v = (d00 * dy - dx * d10) * inv_denom + w = 1.0 - u - v + + inside = (u >= -0.001) & (v >= -0.001) & (w >= -0.001) + if not inside.any(): + continue + + pos_3d = w[..., None] * p0 + u[..., None] * p1 + v[..., None] * p2 + + iy, ix = np.where(inside) + positions[py_range.astype(int)[iy], px_range.astype(int)[ix]] = pos_3d[iy, ix] + mask[py_range.astype(int)[iy], px_range.astype(int)[ix]] = True + + return positions, mask + + +def bake_texture(vertices, faces, uvs, voxel_coords, voxel_attrs, origin, voxel_size, + texture_size=2048, k_neighbors=8, **kwargs): + """ + Bake voxel attributes into a UV-mapped texture. + + Uses scipy cKDTree on sparse voxels for k-nearest-neighbor lookup + with inverse-distance weighting. Avoids dense 3D volume entirely, + preserving native voxel resolution without memory pressure. + + Pipeline: + 1. UV rasterize → 3D position per texel + 2. KDTree on sparse voxels + 3. For each texel: k-nearest voxels, inverse-distance-weighted average + 4. Gamma correct, fill holes, export + """ + from scipy.spatial import cKDTree + + H = W = texture_size + t0 = time.time() + + coords_np = voxel_coords.numpy() if hasattr(voxel_coords, 'numpy') else voxel_coords + attrs_np = voxel_attrs.numpy() if hasattr(voxel_attrs, 'numpy') else voxel_attrs + origin_np = origin.numpy() if hasattr(origin, 'numpy') else np.array(origin) + + C = attrs_np.shape[1] + n_voxels = len(coords_np) + print(f" Voxels: {n_voxels:,}, channels: {C}") + + # Voxel world-space positions (voxel centers) + voxel_world = coords_np.astype(np.float32) * voxel_size + origin_np + voxel_size * 0.5 + + # Build KDTree on voxel positions + print(f" Building KDTree...") + t_tree = time.time() + tree = cKDTree(voxel_world) + print(f" Tree built in {time.time() - t_tree:.1f}s") + + # Rasterize UV triangles to 3D positions + print(f" Rasterizing {len(faces):,} triangles into {texture_size}x{texture_size}...") + t_rast = time.time() + positions, mask = _rasterize_uv_triangles(vertices, faces, uvs, texture_size) + coverage = mask.sum() / (H * W) * 100 + print(f" Coverage: {coverage:.1f}%, rasterized in {time.time() - t_rast:.1f}s") + + # For each valid texel, find k nearest voxels + query_points = positions[mask] # [M, 3] + M = len(query_points) + print(f" Querying {M:,} texels, k={k_neighbors}...") + t_q = time.time() + distances, indices = tree.query(query_points, k=k_neighbors, workers=-1) + # distances: [M, k], indices: [M, k] + print(f" Query done in {time.time() - t_q:.1f}s") + + # Inverse-distance weighted average. Use distance threshold to skip far voxels. + # voxel_size is world units per voxel; 2x voxel_size = reasonable neighborhood + max_dist = voxel_size * 2.0 + print(f" Weighting colors (max_dist = {max_dist:.4f})...") + + # Weights: 1 / (d + eps), but zero weight for distances > max_dist + eps = voxel_size * 0.1 + weights = 1.0 / (distances + eps) + weights[distances > max_dist] = 0.0 + weights_sum = weights.sum(axis=1, keepdims=True) + + # Find texels with at least one nearby voxel + has_neighbor = (weights_sum > 0).squeeze() + + # Normalize weights + weights = np.where(weights_sum > 0, weights / np.maximum(weights_sum, 1e-10), 0.0) + + # Gather colors: attrs_np[indices] → [M, k, C] + neighbor_attrs = attrs_np[indices] # [M, k, C] + + # Weighted sum over k dimension + sampled = (neighbor_attrs * weights[..., None]).sum(axis=1) # [M, C] + + # Write texture + base_color = np.zeros((H, W, 3), dtype=np.float32) + metallic = np.zeros((H, W), dtype=np.float32) + roughness = np.ones((H, W), dtype=np.float32) + + ys, xs = np.where(mask) + valid = has_neighbor + base_color[ys[valid], xs[valid]] = np.clip(sampled[valid, 0:3], 0, 1) + if C > 3: + metallic[ys[valid], xs[valid]] = np.clip(sampled[valid, 3], 0, 1) + if C > 4: + roughness[ys[valid], xs[valid]] = np.clip(sampled[valid, 4], 0, 1) + + valid_mask = np.zeros((H, W), dtype=bool) + valid_mask[ys[valid], xs[valid]] = True + + # Fill holes via iterative dilation + from scipy.ndimage import binary_dilation, uniform_filter + current_mask = valid_mask.copy() + for _ in range(8): + dilated = binary_dilation(current_mask, iterations=1) + unfilled = dilated & ~current_mask + if not unfilled.any(): + break + for c in range(3): + channel = base_color[:, :, c] + blurred = uniform_filter(channel, size=3) + channel[unfilled] = blurred[unfilled] + current_mask = dilated + + # Gamma correction: linear -> sRGB + base_color = np.power(np.clip(base_color, 0, 1), 1.0 / 2.2) + + base_color_img = (base_color * 255).astype(np.uint8) + + # glTF metallic-roughness: R=0, G=roughness, B=metallic + mr_img = np.zeros((H, W, 3), dtype=np.uint8) + mr_img[:, :, 1] = (roughness * 255).astype(np.uint8) + mr_img[:, :, 2] = (metallic * 255).astype(np.uint8) + + total_coverage = current_mask.sum() / (H * W) * 100 + print(f" Final coverage: {total_coverage:.1f}%, total bake: {time.time() - t0:.1f}s") + + return base_color_img, mr_img, current_mask + + +def trellis_to_gltf_coordinates(vertices, uvs): + """Convert TRELLIS Z-up coordinates and xatlas UVs for glTF export.""" + + converted_vertices = np.asarray(vertices).copy() + converted_vertices[:, 1] = np.asarray(vertices)[:, 2] + converted_vertices[:, 2] = -np.asarray(vertices)[:, 1] + converted_uvs = np.asarray(uvs).copy() + converted_uvs[:, 1] = 1.0 - converted_uvs[:, 1] + return converted_vertices, converted_uvs + + +def export_glb_with_texture( + vertices, + faces, + uvs, + base_color_img, + mr_img=None, + output_path="output.glb", + *, + trellis_z_up=False, + reference_vertices=None, + reference_faces=None, + return_normal_report=False, +): + """Export a UV-mapped PBR mesh while preserving indexed shading normals. + + ``vertices`` and ``faces`` describe the UV-split mesh. Supplying both + ``reference_vertices`` and ``reference_faces`` selects the strict normal + transfer used by the geometry-preserving Metal path: normals are computed + on that indexed, pre-UV mesh and copied to every compatible UV duplicate. + Existing callers which omit the reference keep the historical post-UV + recomputation, and the function still returns ``output_path`` unless + ``return_normal_report`` is requested. + """ + import trimesh + from PIL import Image + + if (reference_vertices is None) != (reference_faces is None): + raise ValueError( + "reference_vertices and reference_faces must be supplied together" + ) + + if trellis_z_up: + vertices, uvs = trellis_to_gltf_coordinates(vertices, uvs) + mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) + + base_color_pil = Image.fromarray(base_color_img) + + material = trimesh.visual.material.PBRMaterial( + baseColorTexture=base_color_pil, + metallicFactor=0.0, + roughnessFactor=0.8, + ) + + if mr_img is not None: + mr_pil = Image.fromarray(mr_img) + material.metallicRoughnessTexture = mr_pil + + mesh.visual = trimesh.visual.TextureVisuals( + uv=uvs, + material=material, + ) + + # Trimesh does not necessarily emit NORMAL unless it has been explicitly + # populated. When the indexed proxy is available, compute normals before + # UV duplication and transfer them without welding positions. The legacy + # recompute remains only for API-compatible standalone uses which do not + # provide that proxy. + from backends.export_normals import ( + recompute_asset_vertex_normals, + transfer_reference_vertex_normals, + trellis_z_up_to_gltf_y_up_transform, + ) + + if reference_vertices is not None: + reference_to_asset = ( + trellis_z_up_to_gltf_y_up_transform() if trellis_z_up else None + ) + normal_report = transfer_reference_vertex_normals( + mesh, + reference_vertices, + reference_faces, + reference_to_asset=reference_to_asset, + ) + else: + normal_report = recompute_asset_vertex_normals(mesh) + mesh.export(output_path) + if return_normal_report: + return output_path, normal_report + return output_path diff --git a/inference.py b/inference.py index 4fe1c44..1e942e5 100644 --- a/inference.py +++ b/inference.py @@ -2,19 +2,33 @@ import argparse import math import time -import torch import numpy as np import cv2 from PIL import Image os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" -os.environ.setdefault("ATTN_BACKEND", "flash_attn") +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +os.environ.setdefault("ATTN_BACKEND", "sdpa") +# PyTorch's MPS SDPA remains substantially faster for Pixal3D's ~40k-token +# HR sequences. The validated flex_gemm kernel stays available as an opt-in +# backend, but regresses badly once attention becomes this long. +os.environ.setdefault("SPARSE_ATTN_BACKEND", "sdpa") +os.environ.setdefault("SPARSE_CONV_BACKEND", "flex_gemm") os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json') -os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1' +from macos_compat import configure + +DEVICE = configure() +import torch + +from backends.cuda_parity_export import to_glb_cuda_parity +from backends.decoded_checkpoint import ( + load_decoded_checkpoint, + save_decoded_checkpoint, +) +from backends.memory import release_accelerator_memory from pixal3d.pipelines import Pixal3DImageTo3DPipeline -import o_voxel +from o_voxel import postprocess_cpu # ============================================================================ # Constants & Defaults @@ -56,14 +70,14 @@ # Model Loading # ============================================================================ -def build_image_cond_model(config: dict): +def build_image_cond_model(config: dict, shared_model=None): from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor - model = DinoV3ProjFeatureExtractor(**config) + model = DinoV3ProjFeatureExtractor(**config, shared_model=shared_model) model.eval() return model -def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME): +def load_moge_model(device=DEVICE, model_name=MOGE_MODEL_NAME): from moge.model.v2 import MoGeModel moge_model = MoGeModel.from_pretrained(model_name) moge_model = moge_model.to(device) @@ -71,42 +85,59 @@ def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME): return moge_model -def init_pipeline(model_path=MODEL_PATH, device="cuda", low_vram=False): +def init_pipeline(model_path=MODEL_PATH, device=DEVICE, low_vram=True): print(f"[Pipeline] Loading from {model_path}...") pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path) print("[ImageCond] Building DinoV3ProjFeatureExtractor models...") pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"]) - pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"]) - pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"]) - pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"]) + shared_dino = pipeline.image_cond_model_ss.model + pipeline.image_cond_model_shape_512 = build_image_cond_model( + IMAGE_COND_CONFIGS["shape_512"], shared_model=shared_dino + ) + pipeline.image_cond_model_shape_1024 = build_image_cond_model( + IMAGE_COND_CONFIGS["shape_1024"], shared_model=shared_dino + ) + pipeline.image_cond_model_tex_1024 = build_image_cond_model( + IMAGE_COND_CONFIGS["tex_1024"], shared_model=shared_dino + ) if low_vram: # Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage. # Peak VRAM = one flow model + one DinoV3, not all ~18 GB at once. print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...") + shared_naf = None for attr in ['image_cond_model_ss', 'image_cond_model_shape_512', 'image_cond_model_shape_1024', 'image_cond_model_tex_1024']: m = getattr(pipeline, attr, None) if m is not None and getattr(m, 'use_naf_upsample', False): - m._load_naf() + if shared_naf is None: + m._load_naf() + shared_naf = m.naf_model + else: + m.naf_model = shared_naf pipeline._device = torch.device(device) pipeline.low_vram = True print("[Pipeline] Low-VRAM mode enabled.") else: # Standard mode: all models loaded to GPU at once (faster, needs more VRAM). pipeline.low_vram = False - pipeline.cuda() - pipeline.image_cond_model_ss.cuda() - pipeline.image_cond_model_shape_512.cuda() - pipeline.image_cond_model_shape_1024.cuda() - pipeline.image_cond_model_tex_1024.cuda() + pipeline.to(device) + pipeline.image_cond_model_ss.to(device) + pipeline.image_cond_model_shape_512.to(device) + pipeline.image_cond_model_shape_1024.to(device) + pipeline.image_cond_model_tex_1024.to(device) print("[NAF] Pre-loading NAF upsampler model...") + shared_naf = None for attr in ['image_cond_model_ss', 'image_cond_model_shape_512', 'image_cond_model_shape_1024', 'image_cond_model_tex_1024']: m = getattr(pipeline, attr, None) if m is not None and getattr(m, 'use_naf_upsample', False): - m._load_naf() + if shared_naf is None: + m._load_naf() + shared_naf = m.naf_model + else: + m.naf_model = shared_naf print("[Pipeline] Standard mode (all models on GPU).") return pipeline @@ -134,7 +165,7 @@ def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, imag return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)} -def get_camera_params_wild_moge(image_path, moge_model, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512): +def get_camera_params_wild_moge(image_path, moge_model, device=DEVICE, mesh_scale=1.0, extend_pixel=0, image_resolution=512): pil_image = Image.open(image_path).convert("RGB") width, height = pil_image.size image_np = np.array(pil_image).astype(np.float32) / 255.0 @@ -159,7 +190,7 @@ def get_camera_params_wild_moge(image_path, moge_model, device="cuda", mesh_scal # ============================================================================ def run_inference( - image_path: str, + image_path: str | None, output_path: str, seed: int = 42, ss_guidance_strength: float = 7.5, @@ -182,91 +213,237 @@ def run_inference( manual_fov: float = -1.0, low_vram: bool = False, resolution: int = -1, + decimation_target: int | None = None, + texture_size: int | None = None, + export_profile: str = "cuda-parity", + decoded_checkpoint: str | None = None, + checkpoint_output: str | None = None, + save_decoded: bool | None = None, + bvh_chunk_size: int = 262_144, + grid_chunk_size: int = 262_144, + source_face_chunk_size: int = 250_000, + remesh_resolution: int | None = None, ): - # Load models - pipeline = init_pipeline(model_path, low_vram=low_vram) - - # Preprocess image first — rembg loads to GPU for this call, then offloads. - # MoGe is loaded afterwards so both never occupy VRAM at the same time. - print(f"[Inference] Processing image: {image_path}") - img = Image.open(image_path) - image_preprocessed = pipeline.preprocess_image(img) - - # Save preprocessed image for MoGe - tmp_path = os.path.join(os.path.dirname(os.path.abspath(output_path)), f"_tmp_preprocessed_{int(time.time()*1000)}.png") - image_preprocessed.save(tmp_path) - - # Camera estimation - if manual_fov > 0: - # Use manually specified FOV (in radians) - camera_angle_x = float(manual_fov) - grid_point = torch.tensor([-1.0, 0.0, 0.0]) - distance = distance_from_fov( - camera_angle_x, grid_point, - torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]), - mesh_scale, image_resolution - )["distance_from_x"] - camera_params = {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale} - print(f"[Inference] Using manual FOV: {math.degrees(manual_fov):.2f}° ({manual_fov:.4f} rad), distance={distance:.4f}") + if export_profile not in {"portable", "cuda-parity"}: + raise ValueError( + f"Unknown export profile {export_profile!r}; " + "expected 'portable' or 'cuda-parity'" + ) + if decoded_checkpoint is None and not image_path: + raise ValueError("image_path is required unless decoded_checkpoint is used") + if save_decoded is None: + save_decoded = export_profile == "cuda-parity" + + total_started = time.perf_counter() + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + + if decoded_checkpoint: + print(f"[Checkpoint] Loading decoded tensors: {decoded_checkpoint}") + mesh, res, checkpoint_metadata = load_decoded_checkpoint( + decoded_checkpoint + ) + generation_seconds = checkpoint_metadata.get("generation_seconds") + print( + f"[Checkpoint] Loaded {len(mesh.vertices):,} vertices, " + f"{len(mesh.faces):,} faces at resolution {res}." + ) else: - print("[MoGe-2] Loading model for camera estimation...") - moge_model = load_moge_model(device="cuda") - print("[Inference] Estimating camera parameters...") - camera_params = get_camera_params_wild_moge( - tmp_path, moge_model, device="cuda", - mesh_scale=mesh_scale, extend_pixel=extend_pixel, - image_resolution=image_resolution, + generation_started = time.perf_counter() + pipeline = init_pipeline(model_path, low_vram=low_vram) + + # Preprocess first. The background-removal network is one-shot and is + # discarded before MoGe or any flow model enters unified GPU memory. + assert image_path is not None + print(f"[Inference] Processing image: {image_path}") + with Image.open(image_path) as img: + image_preprocessed = pipeline.preprocess_image(img) + pipeline.rembg_model = None + release_accelerator_memory( + "background-removal model released", + verbose=True, ) - print(f" camera_angle_x={camera_params['camera_angle_x']:.4f}, distance={camera_params['distance']:.4f}") - # MoGe is only needed for camera estimation; free its VRAM for inference. - moge_model.cpu() - del moge_model - torch.cuda.empty_cache() - os.remove(tmp_path) - - # Run pipeline - print("[Inference] Running 3D generation pipeline...") - torch.manual_seed(seed) - - ss_sampler_override = { - "steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength, - "guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t, - } - shape_sampler_override = { - "steps": shape_slat_sampling_steps, "guidance_strength": shape_slat_guidance_strength, - "guidance_rescale": shape_slat_guidance_rescale, "rescale_t": shape_slat_rescale_t, - } - tex_sampler_override = { - "steps": tex_slat_sampling_steps, "guidance_strength": tex_slat_guidance_strength, - "guidance_rescale": tex_slat_guidance_rescale, "rescale_t": tex_slat_rescale_t, - } - - pipeline_type = f"{resolution if resolution > 0 else (1024 if low_vram else 1536)}_cascade" - print(f"[Inference] Using pipeline_type={pipeline_type}") - mesh_list, (shape_slat, tex_slat, res) = pipeline.run( - image_preprocessed, - camera_params=camera_params, - seed=seed, - sparse_structure_sampler_params=ss_sampler_override, - shape_slat_sampler_params=shape_sampler_override, - tex_slat_sampler_params=tex_sampler_override, - preprocess_image=False, - return_latent=True, - pipeline_type=pipeline_type, - max_num_tokens=max_num_tokens, - ) - mesh = mesh_list[0] + tmp_path = os.path.join( + os.path.dirname(os.path.abspath(output_path)), + f"_tmp_preprocessed_{int(time.time() * 1000)}.png", + ) + image_preprocessed.save(tmp_path) + try: + if manual_fov > 0: + camera_angle_x = float(manual_fov) + grid_point = torch.tensor([-1.0, 0.0, 0.0]) + distance = distance_from_fov( + camera_angle_x, + grid_point, + torch.tensor( + [ + 0 - extend_pixel, + image_resolution - 1 + extend_pixel, + ] + ), + mesh_scale, + image_resolution, + )["distance_from_x"] + camera_params = { + "camera_angle_x": camera_angle_x, + "distance": distance, + "mesh_scale": mesh_scale, + } + print( + "[Inference] Using manual FOV: " + f"{math.degrees(manual_fov):.2f}° " + f"({manual_fov:.4f} rad), distance={distance:.4f}" + ) + else: + print("[MoGe-2] Loading model for camera estimation...") + moge_model = load_moge_model(device=DEVICE) + print("[Inference] Estimating camera parameters...") + camera_params = get_camera_params_wild_moge( + tmp_path, + moge_model, + device=DEVICE, + mesh_scale=mesh_scale, + extend_pixel=extend_pixel, + image_resolution=image_resolution, + ) + print( + f" camera_angle_x={camera_params['camera_angle_x']:.4f}, " + f"distance={camera_params['distance']:.4f}" + ) + moge_model.cpu() + del moge_model + release_accelerator_memory( + "MoGe camera model released", + verbose=True, + ) + finally: + if os.path.exists(tmp_path): + os.remove(tmp_path) + + print("[Inference] Running 3D generation pipeline...") + torch.manual_seed(seed) + ss_sampler_override = { + "steps": ss_sampling_steps, + "guidance_strength": ss_guidance_strength, + "guidance_rescale": ss_guidance_rescale, + "rescale_t": ss_rescale_t, + } + shape_sampler_override = { + "steps": shape_slat_sampling_steps, + "guidance_strength": shape_slat_guidance_strength, + "guidance_rescale": shape_slat_guidance_rescale, + "rescale_t": shape_slat_rescale_t, + } + tex_sampler_override = { + "steps": tex_slat_sampling_steps, + "guidance_strength": tex_slat_guidance_strength, + "guidance_rescale": tex_slat_guidance_rescale, + "rescale_t": tex_slat_rescale_t, + } + + requested_resolution = ( + resolution + if resolution > 0 + else (1024 if low_vram else 1536) + ) + pipeline_type = f"{requested_resolution}_cascade" + print(f"[Inference] Using pipeline_type={pipeline_type}") + mesh_list = pipeline.run( + image_preprocessed, + camera_params=camera_params, + seed=seed, + sparse_structure_sampler_params=ss_sampler_override, + shape_slat_sampler_params=shape_sampler_override, + tex_slat_sampler_params=tex_sampler_override, + preprocess_image=False, + return_latent=False, + pipeline_type=pipeline_type, + max_num_tokens=max_num_tokens, + release_models=True, + output_device="cpu", + ) + mesh = mesh_list[0] + res = round(1 / float(mesh.voxel_size)) + generation_seconds = time.perf_counter() - generation_started + + if save_decoded: + if checkpoint_output is None: + checkpoint_output = ( + os.path.splitext(os.path.abspath(output_path))[0] + + ".decoded.pt" + ) + save_decoded_checkpoint( + checkpoint_output, + mesh, + resolution=res, + metadata={ + "camera_params": camera_params, + "generation_seconds": generation_seconds, + "image_path": os.path.abspath(image_path), + "pipeline_type": pipeline_type, + "seed": seed, + }, + ) + print(f"[Checkpoint] Saved decoded tensors: {checkpoint_output}") + + del mesh_list, pipeline, image_preprocessed + release_accelerator_memory( + "all neural models released before export", + verbose=True, + ) - # Extract GLB - print("[Inference] Extracting GLB...") - glb = o_voxel.postprocess.to_glb( - vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs, - coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout, - grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], - decimation_target=1000000, texture_size=4096, - remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True, - ) + attr_layout = dict(mesh.layout) + export_started = time.perf_counter() + if export_profile == "cuda-parity": + decimation_target = decimation_target or 1_000_000 + texture_size = texture_size or 4096 + effective_remesh_resolution = remesh_resolution or min(512, res) + print( + "[Export] CUDA-parity profile: " + f"remesh={effective_remesh_resolution}, " + f"faces={decimation_target:,}, texture={texture_size}², " + f"source BVH faces/chunk={source_face_chunk_size:,}, " + f"BVH chunks={bvh_chunk_size:,}, " + f"volume chunks={grid_chunk_size:,}." + ) + glb = to_glb_cuda_parity( + vertices=mesh.vertices, + faces=mesh.faces, + attr_volume=mesh.attrs, + coords=mesh.coords, + attr_layout=attr_layout, + resolution=res, + decimation_target=decimation_target, + texture_size=texture_size, + bvh_chunk_size=bvh_chunk_size, + grid_chunk_size=grid_chunk_size, + source_face_chunk_size=source_face_chunk_size, + remesh_resolution=effective_remesh_resolution, + verbose=True, + use_tqdm=True, + ) + else: + decimation_target = decimation_target or 50_000 + texture_size = texture_size or 256 + print( + "[Export] Portable profile: " + f"faces={decimation_target:,}, texture={texture_size}²." + ) + glb = postprocess_cpu.to_glb( + vertices=mesh.vertices, + faces=mesh.faces, + attr_volume=mesh.attrs, + coords=mesh.coords, + attr_layout=attr_layout, + grid_size=res, + aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + decimation_target=decimation_target, + texture_size=texture_size, + remesh=False, + remesh_band=1, + remesh_project=0, + use_tqdm=True, + ) # Apply rotation rot = np.array([ @@ -281,11 +458,24 @@ def run_inference( os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) glb.export(output_path, extension_webp=True) print(f"[Done] GLB saved to: {output_path}") + export_seconds = time.perf_counter() - export_started + total_seconds = time.perf_counter() - total_started + if generation_seconds is not None: + print(f"[Timing] Generation: {generation_seconds:.2f} s") + print( + f"[Timing] Export: {export_seconds:.2f} s; " + f"this invocation: {total_seconds:.2f} s" + ) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Pixal3D Inference: Image to GLB") - parser.add_argument("--image", type=str, required=True, help="Path to input image") + parser.add_argument( + "--image", + type=str, + default=None, + help="Path to input image (not needed with --decoded-checkpoint).", + ) parser.add_argument("--output", type=str, default="./output.glb", help="Output GLB file path") parser.add_argument("--seed", type=int, default=42, help="Random seed") parser.add_argument("--fov", type=float, default=-1.0, @@ -293,11 +483,87 @@ def run_inference( "If not set, FOV is auto-estimated via MoGe-2. " "Try 0.2 rad if you notice distortion.") parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Model path or HuggingFace repo") - parser.add_argument("--low_vram", action="store_true", + parser.add_argument("--low_vram", action="store_true", default=True, help="Enable low-VRAM mode: models stay on CPU and are loaded to GPU on-demand per stage. " "Reduces peak VRAM from ~18GB to ~10-12GB at the cost of slower inference.") - parser.add_argument("--resolution", type=int, default=-1, + parser.add_argument("--standard", action="store_false", dest="low_vram", + help="Keep all Pixal3D stages resident (not recommended on 36GB unified memory).") + parser.add_argument("--resolution", type=int, default=-1, choices=[1024, 1536], help="Pipeline resolution (1024 or 1536). Default: 1024 if --low_vram, else 1536.") + parser.add_argument( + "--export-profile", + choices=["cuda-parity", "portable"], + default="cuda-parity", + help=( + "GLB profile. cuda-parity uses native Metal remeshing, one million " + "faces and 4096px textures; portable keeps the former lightweight " + "fallback." + ), + ) + parser.add_argument( + "--decimation-target", + type=int, + default=None, + help="Override target face count for the selected export profile.", + ) + parser.add_argument( + "--texture-size", + type=int, + default=None, + help="Override square PBR texture size for the selected export profile.", + ) + parser.add_argument( + "--decoded-checkpoint", + type=str, + default=None, + help=( + "Resume directly from a .decoded.pt checkpoint and skip all neural " + "generation stages." + ), + ) + parser.add_argument( + "--checkpoint-output", + type=str, + default=None, + help="Optional path for the decoded checkpoint saved before export.", + ) + parser.add_argument( + "--no-save-decoded", + action="store_false", + dest="save_decoded", + default=None, + help="Do not save a resumable decoded checkpoint before parity export.", + ) + parser.add_argument( + "--bvh-chunk-size", + type=int, + default=262_144, + help="Closest-point queries per native Metal batch (default: 262144).", + ) + parser.add_argument( + "--grid-chunk-size", + type=int, + default=262_144, + help="Texture-volume samples per Metal batch (default: 262144).", + ) + parser.add_argument( + "--source-face-chunk-size", + type=int, + default=250_000, + help=( + "Decoded source faces per accurate Metal BVH " + "(default: 250000)." + ), + ) + parser.add_argument( + "--remesh-resolution", + type=int, + default=0, + help=( + "Dual-contouring grid size. 0 selects the validated 512-cell " + "36GB profile; higher values require substantially more memory." + ), + ) args = parser.parse_args() @@ -309,4 +575,14 @@ def run_inference( model_path=args.model_path, low_vram=args.low_vram, resolution=args.resolution, + decimation_target=args.decimation_target, + texture_size=args.texture_size, + export_profile=args.export_profile, + decoded_checkpoint=args.decoded_checkpoint, + checkpoint_output=args.checkpoint_output, + save_decoded=args.save_decoded, + bvh_chunk_size=args.bvh_chunk_size, + grid_chunk_size=args.grid_chunk_size, + source_face_chunk_size=args.source_face_chunk_size, + remesh_resolution=args.remesh_resolution or None, ) diff --git a/macos_compat.py b/macos_compat.py new file mode 100644 index 0000000..28beba4 --- /dev/null +++ b/macos_compat.py @@ -0,0 +1,44 @@ +"""Runtime compatibility for CUDA calls made by the upstream Pixal3D code.""" + +from __future__ import annotations + +import os + + +def configure() -> str: + """Route legacy ``.cuda()`` calls to MPS and make cache calls harmless.""" + + os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") + os.environ.setdefault("ATTN_BACKEND", "sdpa") + os.environ.setdefault("SPARSE_ATTN_BACKEND", "sdpa") + os.environ.setdefault("SPARSE_CONV_BACKEND", "flex_gemm") + + import torch + + device_name = os.environ.get("PIXAL3D_DEVICE", "mps") + if device_name == "mps" and not torch.backends.mps.is_available(): + device_name = "cpu" + device = torch.device(device_name) + + def tensor_cuda(self, device=None, non_blocking=False, memory_format=torch.preserve_format): + target = device if device is not None and str(device) != "cuda" else device_name + return self.to(target, non_blocking=non_blocking, memory_format=memory_format) + + def module_cuda(self, device=None): + target = device if device is not None and str(device) != "cuda" else device_name + return self.to(target) + + torch.Tensor.cuda = tensor_cuda + torch.nn.Module.cuda = module_cuda + + def empty_cache() -> None: + if device_name == "mps": + torch.mps.empty_cache() + + def synchronize(device=None) -> None: + if device_name == "mps": + torch.mps.synchronize() + + torch.cuda.empty_cache = empty_cache + torch.cuda.synchronize = synchronize + return str(device) diff --git a/patches/mps_compat.py b/patches/mps_compat.py new file mode 100644 index 0000000..469d957 --- /dev/null +++ b/patches/mps_compat.py @@ -0,0 +1,413 @@ +"""Apply the small source changes needed by Pixal3D on Apple Silicon. + +The upstream project assumes CUDA. The actual Metal kernels are installed by +``setup_macos.sh``; this script wires them into Pixal3D and installs the +portable mesh-extraction fallback used by the TRELLIS.2 macOS port. +""" + +from __future__ import annotations + +import os +import re +import shutil +from pathlib import Path + + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def read(path: str) -> str: + with open(path, encoding="utf-8") as handle: + return handle.read() + + +def write(path: str, text: str) -> None: + with open(path, "w", encoding="utf-8") as handle: + handle.write(text) + print(f" patched {os.path.relpath(path, ROOT)}") + + +def replace_once(path: str, old: str, new: str, label: str) -> None: + text = read(path) + if new in text: + print(f" already patched {label}") + return + if old not in text: + raise RuntimeError(f"Could not find patch anchor for {label}: {path}") + write(path, text.replace(old, new, 1)) + + +def patch_pipeline_base() -> None: + path = os.path.join(ROOT, "pixal3d/pipelines/base.py") + replace_once( + path, + ' self.to(torch.device("cuda"))', + ' self.to(torch.device("mps") if torch.backends.mps.is_available() else torch.device("cuda"))', + "Pipeline.cuda()", + ) + + +def patch_birefnet() -> None: + path = os.path.join(ROOT, "pixal3d/pipelines/rembg/BiRefNet.py") + text = read(path) + if "def device(self)" not in text: + text = text.replace( + " def to(self, device: str):\n self.model.to(device)\n\n def cuda(self):", + " @property\n def device(self):\n return next(self.model.parameters()).device\n\n def to(self, device: str):\n self.model.to(device)\n return self\n\n def cuda(self):", + 1, + ) + text = text.replace('.unsqueeze(0).to("cuda")', ".unsqueeze(0).to(self.device)") + write(path, text) + + +def patch_image_extractors() -> None: + for relative in ( + "pixal3d/modules/image_feature_extractor.py", + ): + path = os.path.join(ROOT, relative) + text = read(path) + text = text.replace(".cuda()", ".to(self.device)") + if relative.endswith("image_feature_extractor.py"): + property_block = ( + " @property\n" + " def device(self):\n" + " return next(self.model.parameters()).device\n\n" + ) + first_anchor = "class DinoV2FeatureExtractor:" + second_anchor = "class DinoV3FeatureExtractor:" + for anchor in (first_anchor, second_anchor): + start = text.index(anchor) + end = text.find("\nclass ", start + len(anchor)) + if end == -1: + end = len(text) + block = text[start:end] + if "def device(self)" not in block: + block = block.replace( + " def to(self, device):\n", + property_block + " def to(self, device):\n", + 1, + ) + text = text[:start] + block + text[end:] + write(path, text) + + +def patch_varlen_reduce() -> None: + path = os.path.join(ROOT, "pixal3d/modules/sparse/basic.py") + text = read(path) + marker = "pixal3d-macos: MPS segment reduce" + if marker in text: + print(" already patched pixal3d/modules/sparse/basic.py") + return + old = " red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen)\n return red" + new = """ # pixal3d-macos: MPS segment reduce. The layout is authoritative; + # cached lengths can describe a previous cascade scale. + lengths = self.seqlen + if int(lengths.sum().item()) != red.shape[0]: + lengths = torch.tensor( + [s.stop - s.start for s in self.layout], + dtype=torch.long, + device=red.device, + ) + if int(lengths.sum().item()) != red.shape[0]: + raise RuntimeError("Sparse VarLenTensor has inconsistent segment lengths") + if red.device.type == 'mps': + return torch.segment_reduce( + red.cpu(), reduce=op, lengths=lengths.cpu() + ).to(red.device) + return torch.segment_reduce(red, reduce=op, lengths=lengths)""" + replace_once(path, old, new, marker) + + +def patch_fdg_vae() -> None: + path = os.path.join(ROOT, "pixal3d/models/sc_vaes/fdg_vae.py") + old = "from o_voxel.convert import flexible_dual_grid_to_mesh\n" + new = """# The Metal o_voxel converter is not reliable for decoder output on every +# macOS/PyTorch combination. Prefer the portable implementation shipped in +# backends/mesh_extract.py; the Metal postprocess module remains available for +# textured GLB export. +import sys as _sys +_stubs = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'stubs') +if _stubs not in _sys.path: + _sys.path.append(_stubs) +from o_voxel_override_convert import flexible_dual_grid_to_mesh +""" + replace_once(path, old, new, "pure-Python dual-grid mesh extraction") + + +def install_backends() -> None: + source = os.path.join(ROOT, "backends/conv_none.py") + target = os.path.join(ROOT, "pixal3d/modules/sparse/conv/conv_none.py") + if not os.path.exists(target): + shutil.copy2(source, target) + print(" installed pixal3d/modules/sparse/conv/conv_none.py") + + stubs = os.path.join(ROOT, "stubs") + shutil.copy2( + os.path.join(ROOT, "backends/mesh_extract.py"), + os.path.join(stubs, "o_voxel_override_convert.py"), + ) + shutil.copy2( + os.path.join(ROOT, "backends/mesh_extract.py"), + os.path.join(stubs, "o_voxel/convert.py"), + ) + print(" installed portable o_voxel converter") + + +def patch_o_voxel() -> None: + """Install the tracked Metal post-processing compatibility changes.""" + try: + import o_voxel.postprocess as postprocess + except ImportError: + return + path = Path(postprocess.__file__) + text = path.read_text(encoding="utf-8") + call = "mesh.fill_holes(max_hole_perimeter=3e-2)" + # Normalize both a pristine install and a previous partially patched one, + # preserving the surrounding block indentation. + text = re.sub( + rf"(?m)^([ ]*)if _BACKEND != 'metal':\n\1 if _BACKEND != 'metal':\n\1 {re.escape(call)}$", + rf"\1if _BACKEND != 'metal':\n\1 {call}", + text, + ) + if "if _BACKEND != 'metal':" not in text: + text = re.sub( + rf"(?m)^([ ]*){re.escape(call)}$", + rf"\1if _BACKEND != 'metal':\n\1 {call}", + text, + ) + + if "reproject_texture_to_source: bool = True" not in text: + anchors = ( + ( + " verbose: bool = False,\n" + " use_tqdm: bool = False,\n" + "):", + " verbose: bool = False,\n" + " use_tqdm: bool = False,\n" + " reproject_texture_to_source: bool = True,\n" + "):", + "texture projection argument", + ), + ( + " use_tqdm: whether to use tqdm to display progress bar\n", + " use_tqdm: whether to use tqdm to display progress bar\n" + " reproject_texture_to_source: project UV texels back to " + "the input mesh\n" + " before sampling attributes. Disable only when " + "post-processing is\n" + " guaranteed to preserve the input geometry exactly.\n", + "texture projection documentation", + ), + ( + " # Build BVH for the current mesh to guide remeshing\n" + " if use_tqdm:\n" + ' pbar.set_description("Building BVH")\n' + " if verbose:\n" + ' print(f"Building BVH for current mesh...", end=\'\', flush=True)\n' + " bvh = _BVH(vertices, faces)\n", + " # A BVH is needed only when remeshing or projecting UV " + "texels back to the\n" + " # source surface. Preserve mode does neither.\n" + " needs_bvh = remesh or reproject_texture_to_source\n" + " if use_tqdm:\n" + ' pbar.set_description("Building BVH" if needs_bvh else "Skipping BVH")\n' + " if verbose:\n" + ' action = "Building" if needs_bvh else "Skipping"\n' + ' print(f"{action} BVH for current mesh...", end=\'\', flush=True)\n' + " bvh = _BVH(vertices, faces) if needs_bvh else None\n", + "optional source BVH", + ), + ( + " # Map these positions back to the *original* high-res mesh " + "to get accurate attributes\n" + " # This corrects geometric errors introduced by " + "simplification/remeshing\n" + " _, face_id, uvw = bvh.unsigned_distance(valid_pos, " + "return_uvw=True)\n" + " orig_tri_verts = vertices[faces[face_id.long()]] # " + "(N_new, 3, 3)\n" + " valid_pos = (orig_tri_verts * " + "uvw.unsqueeze(-1)).sum(dim=1)\n", + " # Reproject only when geometry processing has moved the " + "surface.\n" + " if reproject_texture_to_source:\n" + " _, face_id, uvw = bvh.unsigned_distance(\n" + " valid_pos, return_uvw=True\n" + " )\n" + " orig_tri_verts = vertices[faces[face_id.long()]]\n" + " valid_pos = (orig_tri_verts * " + "uvw.unsqueeze(-1)).sum(dim=1)\n", + "optional texture reprojection", + ), + ) + for old, new, label in anchors: + if old not in text: + raise RuntimeError( + f"Could not find o_voxel patch anchor for {label}: {path}" + ) + text = text.replace(old, new, 1) + + if "texture_sample_vertices: Optional[torch.Tensor] = None" not in text: + anchors = ( + ( + " reproject_texture_to_source: bool = True,\n" + "):", + " reproject_texture_to_source: bool = True,\n" + " texture_sample_vertices: Optional[torch.Tensor] = None,\n" + "):", + "texture sampling vertices argument", + ), + ( + " guaranteed to preserve the input geometry exactly.\n", + " guaranteed to preserve the input geometry exactly.\n" + " texture_sample_vertices: optional positions aligned " + "with ``vertices``\n" + " that are used only for volume sampling.\n", + "texture sampling vertices documentation", + ), + ( + " vertices = vertices.to(device)\n" + " faces = faces.to(device)\n", + " vertices = vertices.to(device)\n" + " faces = faces.to(device)\n" + " if texture_sample_vertices is not None:\n" + " if texture_sample_vertices.shape != vertices.shape:\n" + " raise ValueError(\n" + ' "texture_sample_vertices must have the same ' + 'shape as vertices"\n' + " )\n" + " texture_sample_vertices = " + "texture_sample_vertices.to(device)\n", + "texture sampling vertices validation", + ), + ( + " out_vmaps = out_vmaps.to(device)\n" + " mesh.compute_vertex_normals()\n", + " out_vmaps = out_vmaps.to(device)\n" + " texture_out_vertices = (\n" + " out_vertices\n" + " if texture_sample_vertices is None\n" + " else texture_sample_vertices[out_vmaps]\n" + " )\n" + " mesh.compute_vertex_normals()\n", + "UV texture sampling vertices", + ), + ( + " pos = dr.interpolate(out_vertices.unsqueeze(0), rast, " + "out_faces)[0][0]\n", + " pos = dr.interpolate(\n" + " texture_out_vertices.unsqueeze(0),\n" + " rast,\n" + " out_faces,\n" + " )[0][0]\n", + "texture position interpolation", + ), + ) + for old, new, label in anchors: + if old not in text: + raise RuntimeError( + f"Could not find o_voxel patch anchor for {label}: {path}" + ) + text = text.replace(old, new, 1) + + if "texture_fallback_projector: Optional[" not in text: + anchors = ( + ( + " texture_sample_vertices: Optional[torch.Tensor] = None,\n" + "):", + " texture_sample_vertices: Optional[torch.Tensor] = None,\n" + " texture_fallback_projector: Optional[\n" + " Callable[[torch.Tensor], torch.Tensor]\n" + " ] = None,\n" + "):", + "texture fallback projector argument", + ), + ( + " that are used only for volume sampling.\n", + " that are used only for volume sampling.\n" + " texture_fallback_projector: optional callable that " + "projects only\n" + " texels whose first sparse-volume sample has " + "invalid alpha.\n", + "texture fallback projector documentation", + ), + ( + " attrs = torch.zeros(texture_size, texture_size, " + "attr_volume.shape[1], device=device)\n" + " attrs[mask] = _grid_sample_3d(\n" + " attr_volume,\n" + " torch.cat([torch.zeros_like(coords[:, :1]), coords], " + "dim=-1),\n" + " shape=torch.Size([1, attr_volume.shape[1], " + "*grid_size.tolist()]),\n" + " grid=((valid_pos - aabb[0]) / " + "voxel_size).reshape(1, -1, 3),\n" + " mode='trilinear',\n" + " )\n", + " attrs = torch.zeros(texture_size, texture_size, " + "attr_volume.shape[1], device=device)\n" + " sampled_attrs = _grid_sample_3d(\n" + " attr_volume,\n" + " torch.cat([torch.zeros_like(coords[:, :1]), coords], " + "dim=-1),\n" + " shape=torch.Size([1, attr_volume.shape[1], " + "*grid_size.tolist()]),\n" + " grid=((valid_pos - aabb[0]) / " + "voxel_size).reshape(1, -1, 3),\n" + " mode='trilinear',\n" + " )\n" + " if texture_fallback_projector is not None:\n" + " alpha_values = sampled_attrs[..., " + "attr_layout['alpha']]\n" + " needs_fallback = alpha_values.amin(dim=-1) < " + "(250.0 / 255.0)\n" + " if needs_fallback.any():\n" + " corrected_pos = texture_fallback_projector(\n" + " valid_pos[needs_fallback]\n" + " )\n" + " sampled_attrs[needs_fallback] = " + "_grid_sample_3d(\n" + " attr_volume,\n" + " torch.cat(\n" + " [torch.zeros_like(coords[:, :1]), " + "coords], dim=-1\n" + " ),\n" + " shape=torch.Size(\n" + " [1, attr_volume.shape[1], " + "*grid_size.tolist()]\n" + " ),\n" + " grid=((corrected_pos - aabb[0]) / " + "voxel_size).reshape(\n" + " 1, -1, 3\n" + " ),\n" + " mode='trilinear',\n" + " )\n" + " attrs[mask] = sampled_attrs\n", + "targeted texture fallback", + ), + ) + for old, new, label in anchors: + if old not in text: + raise RuntimeError( + f"Could not find o_voxel patch anchor for {label}: {path}" + ) + text = text.replace(old, new, 1) + + path.write_text(text, encoding="utf-8") + print(" patched o_voxel Metal post-processing path") + + +def main() -> None: + print("Applying Pixal3D macOS/MPS compatibility patches...") + patch_pipeline_base() + patch_birefnet() + patch_image_extractors() + patch_varlen_reduce() + patch_fdg_vae() + install_backends() + patch_o_voxel() + print("All Pixal3D macOS patches applied.") + + +if __name__ == "__main__": + main() diff --git a/patches/o_voxel_preserve_texture.patch b/patches/o_voxel_preserve_texture.patch new file mode 100644 index 0000000..2fd50cd --- /dev/null +++ b/patches/o_voxel_preserve_texture.patch @@ -0,0 +1,44 @@ +diff --git a/o-voxel/o_voxel/postprocess.py b/o-voxel/o_voxel/postprocess.py +index 1f16265..b432f15 100644 +--- a/o-voxel/o_voxel/postprocess.py ++++ b/o-voxel/o_voxel/postprocess.py +@@ -93,0 +94 @@ def to_glb( ++ reproject_texture_to_source: bool = True, +@@ -116,0 +118,3 @@ def to_glb( ++ reproject_texture_to_source: project UV texels back to the input mesh ++ before sampling attributes. Disable only when post-processing is ++ guaranteed to preserve the input geometry exactly. +@@ -198 +202,4 @@ def to_glb( +- # Build BVH for the current mesh to guide remeshing ++ # A BVH is needed only when remeshing or projecting UV texels back to the ++ # source surface. Preserve mode does neither, so avoid a redundant and ++ # potentially fragile build on dense/non-manifold geometry. ++ needs_bvh = remesh or reproject_texture_to_source +@@ -200 +207 @@ def to_glb( +- pbar.set_description("Building BVH") ++ pbar.set_description("Building BVH" if needs_bvh else "Skipping BVH") +@@ -202,2 +209,3 @@ def to_glb( +- print(f"Building BVH for current mesh...", end='', flush=True) +- bvh = _BVH(vertices, faces) ++ action = "Building" if needs_bvh else "Skipping" ++ print(f"{action} BVH for current mesh...", end='', flush=True) ++ bvh = _BVH(vertices, faces) if needs_bvh else None +@@ -341,5 +349,9 @@ def to_glb( +- # Map these positions back to the *original* high-res mesh to get accurate attributes +- # This corrects geometric errors introduced by simplification/remeshing +- _, face_id, uvw = bvh.unsigned_distance(valid_pos, return_uvw=True) +- orig_tri_verts = vertices[faces[face_id.long()]] # (N_new, 3, 3) +- valid_pos = (orig_tri_verts * uvw.unsqueeze(-1)).sum(dim=1) ++ # Map these positions back to the *original* high-res mesh to get accurate ++ # attributes after simplification/remeshing. Geometry-preserving callers ++ # should skip this projection: valid_pos already lies on the immutable ++ # input triangles, while a second closest-point query can jump between ++ # nearby or overlapping non-manifold sheets. ++ if reproject_texture_to_source: ++ _, face_id, uvw = bvh.unsigned_distance(valid_pos, return_uvw=True) ++ orig_tri_verts = vertices[faces[face_id.long()]] # (N_new, 3, 3) ++ valid_pos = (orig_tri_verts * uvw.unsqueeze(-1)).sum(dim=1) +@@ -427 +439 @@ def to_glb( +- return textured_mesh +\ No newline at end of file ++ return textured_mesh diff --git a/pixal3d/models/sc_vaes/fdg_vae.py b/pixal3d/models/sc_vaes/fdg_vae.py index 0209e69..3bb571c 100644 --- a/pixal3d/models/sc_vaes/fdg_vae.py +++ b/pixal3d/models/sc_vaes/fdg_vae.py @@ -1,4 +1,5 @@ from typing import * +import os import torch import torch.nn as nn import torch.nn.functional as F @@ -17,7 +18,15 @@ SparseUnetVaeDecoder, ) from ...representations import Mesh -from o_voxel.convert import flexible_dual_grid_to_mesh +# The Metal o_voxel converter is not reliable for decoder output on every +# macOS/PyTorch combination. Prefer the portable implementation shipped in +# backends/mesh_extract.py; the Metal postprocess module remains available for +# textured GLB export. +import sys as _sys +_stubs = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'stubs') +if _stubs not in _sys.path: + _sys.path.append(_stubs) +from o_voxel_override_convert import flexible_dual_grid_to_mesh class FlexiDualGridVaeEncoder(SparseUnetVaeEncoder): diff --git a/pixal3d/modules/image_feature_extractor.py b/pixal3d/modules/image_feature_extractor.py index c3cb515..9da3f4e 100644 --- a/pixal3d/modules/image_feature_extractor.py +++ b/pixal3d/modules/image_feature_extractor.py @@ -19,11 +19,15 @@ def __init__(self, model_name: str): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) + @property + def device(self): + return next(self.model.parameters()).device + def to(self, device): self.model.to(device) def cuda(self): - self.model.cuda() + self.model.to(self.device) def cpu(self): self.model.cpu() @@ -46,11 +50,11 @@ def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tenso image = [i.resize((518, 518), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") - image = self.transform(image).cuda() + image = self.transform(image).to(self.device) features = self.model(image, is_training=True)['x_prenorm'] patchtokens = F.layer_norm(features, features.shape[-1:]) return patchtokens @@ -69,11 +73,15 @@ def __init__(self, model_name: str, image_size=512): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) + @property + def device(self): + return next(self.model.parameters()).device + def to(self, device): self.model.to(device) def cuda(self): - self.model.cuda() + self.model.to(self.device) def cpu(self): self.model.cpu() @@ -109,10 +117,10 @@ def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tenso image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") - image = self.transform(image).cuda() + image = self.transform(image).to(self.device) features = self.extract_features(image) return features diff --git a/pixal3d/modules/sparse/attention/full_attn.py b/pixal3d/modules/sparse/attention/full_attn.py index 363c048..71cab11 100644 --- a/pixal3d/modules/sparse/attention/full_attn.py +++ b/pixal3d/modules/sparse/attention/full_attn.py @@ -229,6 +229,34 @@ def sparse_scaled_dot_product_attention(*args, **kwargs): max_q_seqlen = max(q_seqlen) max_kv_seqlen = max(kv_seqlen) out, _ = flash_attn_4_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_q_seqlen, max_kv_seqlen) + elif config.ATTN == 'flex_gemm_sparse_attn': + # Metal flash-attention-v2 kernel from mtlgemm. Its packed variable + # length interface is equivalent to flash_attn's CUDA interface and + # avoids padding every sparse window into a dense SDPA batch. + if num_all_args == 1: + q, k, v = qkv.unbind(dim=1) + elif num_all_args == 2: + k, v = kv.unbind(dim=1) + + import math + import flex_gemm + + cu_seqlens_q = torch.cat( + [torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)] + ).int().to(device) + cu_seqlens_kv = torch.cat( + [torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)] + ).int().to(device) + out = flex_gemm.kernels.cuda.sparse_attention_fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + cu_seqlens_q, + cu_seqlens_kv, + max(q_seqlen), + max(kv_seqlen), + 1.0 / math.sqrt(q.shape[-1]), + ) elif config.ATTN == 'sdpa': from torch.nn.functional import scaled_dot_product_attention as _sdpa if num_all_args == 1: diff --git a/pixal3d/modules/sparse/basic.py b/pixal3d/modules/sparse/basic.py index 880973b..0723d24 100644 --- a/pixal3d/modules/sparse/basic.py +++ b/pixal3d/modules/sparse/basic.py @@ -280,8 +280,22 @@ def reduce(self, op: str, dim: Optional[Union[int, Tuple[int,...]]] = None, keep if dim is None or 0 in dim: return red - red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen) - return red + # pixal3d-macos: MPS segment reduce. The layout is authoritative; + # cached lengths can describe a previous cascade scale. + lengths = self.seqlen + if int(lengths.sum().item()) != red.shape[0]: + lengths = torch.tensor( + [s.stop - s.start for s in self.layout], + dtype=torch.long, + device=red.device, + ) + if int(lengths.sum().item()) != red.shape[0]: + raise RuntimeError("Sparse VarLenTensor has inconsistent segment lengths") + if red.device.type == 'mps': + return torch.segment_reduce( + red.cpu(), reduce=op, lengths=lengths.cpu() + ).to(red.device) + return torch.segment_reduce(red, reduce=op, lengths=lengths) def mean(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor: return self.reduce(op='mean', dim=dim, keepdim=keepdim) diff --git a/pixal3d/modules/sparse/config.py b/pixal3d/modules/sparse/config.py index 25610e5..18df5ac 100644 --- a/pixal3d/modules/sparse/config.py +++ b/pixal3d/modules/sparse/config.py @@ -21,7 +21,14 @@ def __from_env(): CONV = env_sparse_conv_backend if env_sparse_debug is not None: DEBUG = env_sparse_debug == '1' - if env_sparse_attn_backend is not None and env_sparse_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'flash_attn_4', 'sdpa']: + if env_sparse_attn_backend is not None and env_sparse_attn_backend in [ + 'xformers', + 'flash_attn', + 'flash_attn_3', + 'flash_attn_4', + 'sdpa', + 'flex_gemm_sparse_attn', + ]: ATTN = env_sparse_attn_backend print(f"[SPARSE] Conv backend: {CONV}; Attention backend: {ATTN}") @@ -38,6 +45,13 @@ def set_debug(debug: bool): global DEBUG DEBUG = debug -def set_attn_backend(backend: Literal['xformers', 'flash_attn', 'flash_attn_3', 'flash_attn_4', 'sdpa']): +def set_attn_backend(backend: Literal[ + 'xformers', + 'flash_attn', + 'flash_attn_3', + 'flash_attn_4', + 'sdpa', + 'flex_gemm_sparse_attn', +]): global ATTN ATTN = backend diff --git a/pixal3d/modules/sparse/conv/conv_none.py b/pixal3d/modules/sparse/conv/conv_none.py new file mode 100644 index 0000000..766b372 --- /dev/null +++ b/pixal3d/modules/sparse/conv/conv_none.py @@ -0,0 +1,133 @@ +""" +Pure-PyTorch sparse 3D convolution backend. + +Implements submanifold sparse convolution by gathering neighbor features, +applying convolution weights via matrix multiply, and scatter-adding results. +No CUDA extensions needed — works on MPS and CPU. + +Slower than flex_gemm/spconv but fully portable. +""" + +import math +import torch +import torch.nn as nn +from .. import SparseTensor + + +def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None): + assert stride == 1 and (padding is None), \ + "Naive implementation only supports submanifold sparse convolution (stride=1, padding=None)" + + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = tuple(kernel_size) if isinstance(kernel_size, (list, tuple)) else (kernel_size,) * 3 + self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride,) * 3 + self.dilation = tuple(dilation) if isinstance(dilation, (list, tuple)) else (dilation,) * 3 + + self.weight = nn.Parameter(torch.empty((out_channels, in_channels, *self.kernel_size))) + if bias: + self.bias = nn.Parameter(torch.empty(out_channels)) + else: + self.register_parameter("bias", None) + + torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(self.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + torch.nn.init.uniform_(self.bias, -bound, bound) + + # Match flex_gemm weight layout: (Co, Ci, Kd, Kh, Kw) -> (Co, Kd, Kh, Kw, Ci) + self.weight = nn.Parameter(self.weight.permute(0, 2, 3, 4, 1).contiguous()) + + +def sparse_conv3d_forward(self, x: SparseTensor) -> SparseTensor: + """ + Submanifold sparse 3D convolution via gather-scatter. + + For each active voxel, gather features from its kernel-sized neighborhood + (only where other active voxels exist), multiply by the corresponding + kernel weight, and scatter-add results back. + """ + Co, Kd, Kh, Kw, Ci = self.weight.shape + device = x.feats.device + dtype = x.feats.dtype + + coords = x.coords # [N, 4] (batch_idx, z, y, x) + feats = x.feats # [N, Ci] + N = coords.shape[0] + + # Build neighbor index cache (reused across forward passes for same coords) + cache_key = f'SubMConv3d_naive_neighbor_{Kw}x{Kh}x{Kd}_dilation{self.dilation}' + neighbor_cache = x.get_spatial_cache(cache_key) + + if neighbor_cache is None: + # Build spatial hash: coord tuple -> voxel index + coord_to_idx = {} + coords_cpu = coords.cpu() + for i in range(N): + key = tuple(coords_cpu[i].tolist()) + coord_to_idx[key] = i + + # For each kernel position, find (source, target) voxel pairs + dz, dy, dx = self.dilation + src_indices = [] + tgt_indices = [] + kernel_indices = [] + + for kz in range(Kd): + for ky in range(Kh): + for kx in range(Kw): + oz = (kz - Kd // 2) * dz + oy = (ky - Kh // 2) * dy + ox = (kx - Kw // 2) * dx + k_idx = kz * Kh * Kw + ky * Kw + kx + + for i in range(N): + b, z, y, xc = coords_cpu[i].tolist() + neighbor_key = (b, z + oz, y + oy, xc + ox) + if neighbor_key in coord_to_idx: + j = coord_to_idx[neighbor_key] + src_indices.append(j) + tgt_indices.append(i) + kernel_indices.append(k_idx) + + neighbor_cache = ( + torch.tensor(src_indices, dtype=torch.long, device=device), + torch.tensor(tgt_indices, dtype=torch.long, device=device), + torch.tensor(kernel_indices, dtype=torch.long, device=device), + ) + x.register_spatial_cache(cache_key, neighbor_cache) + + src_idx, tgt_idx, k_idx = neighbor_cache + + # Reshape weight: (Co, Kd, Kh, Kw, Ci) -> (K, Ci, Co) + K_total = Kd * Kh * Kw + w = self.weight.reshape(Co, K_total, Ci).permute(1, 2, 0) # (K, Ci, Co) + + out = torch.zeros(N, Co, device=device, dtype=dtype) + + if len(src_idx) > 0: + # Process each kernel position to keep memory bounded + for k in range(K_total): + mask = (k_idx == k) + if not mask.any(): + continue + s_idx = src_idx[mask] + t_idx = tgt_idx[mask] + src_f = feats[s_idx] # [E_k, Ci] + edge_out = src_f @ w[k] # [E_k, Co] + out.scatter_add_(0, t_idx.unsqueeze(1).expand(-1, Co), edge_out) + + if self.bias is not None: + out = out + self.bias + + return x.replace(out) + + +def sparse_inverse_conv3d_init(self, *args, **kwargs): + raise NotImplementedError("SparseInverseConv3d with naive backend is not implemented") + + +def sparse_inverse_conv3d_forward(self, x: SparseTensor) -> SparseTensor: + raise NotImplementedError("SparseInverseConv3d with naive backend is not implemented") diff --git a/pixal3d/pipelines/base.py b/pixal3d/pipelines/base.py index 331e1ed..748336b 100644 --- a/pixal3d/pipelines/base.py +++ b/pixal3d/pipelines/base.py @@ -66,7 +66,7 @@ def to(self, device: torch.device) -> None: model.to(device) def cuda(self) -> None: - self.to(torch.device("cuda")) + self.to(torch.device("mps") if torch.backends.mps.is_available() else torch.device("cuda")) def cpu(self) -> None: self.to(torch.device("cpu")) \ No newline at end of file diff --git a/pixal3d/pipelines/pixal3d_image_to_3d.py b/pixal3d/pipelines/pixal3d_image_to_3d.py index 7129482..0ecad92 100644 --- a/pixal3d/pipelines/pixal3d_image_to_3d.py +++ b/pixal3d/pipelines/pixal3d_image_to_3d.py @@ -8,6 +8,7 @@ from ..modules.sparse import SparseTensor from ..modules import image_feature_extractor from ..representations import Mesh, MeshWithVoxel +from backends.memory import drop_model, release_accelerator_memory class Pixal3DImageTo3DPipeline(Pipeline): @@ -41,7 +42,6 @@ class Pixal3DImageTo3DPipeline(Pipeline): 'shape_slat_flow_model_512', 'shape_slat_flow_model_1024', 'shape_slat_decoder', - 'tex_slat_flow_model_512', 'tex_slat_flow_model_1024', 'tex_slat_decoder', ] @@ -577,6 +577,9 @@ def decode_latent( shape_slat: SparseTensor, tex_slat: SparseTensor, resolution: int, + *, + release_models: bool = False, + output_device: Optional[Union[str, torch.device]] = None, ) -> List[MeshWithVoxel]: """ Decode the latent codes. @@ -587,22 +590,51 @@ def decode_latent( resolution (int): The resolution of the output. """ meshes, subs = self.decode_shape_slat(shape_slat, resolution) + if release_models: + drop_model(self, 'shape_slat_decoder') + release_accelerator_memory( + "shape decoder released", + verbose=True, + ) + tex_voxels = self.decode_tex_slat(tex_slat, subs) + if release_models: + drop_model(self, 'tex_slat_decoder') + release_accelerator_memory( + "texture decoder released", + verbose=True, + ) + out_mesh = [] torch.cuda.synchronize() for m, v in zip(meshes, tex_voxels): m.fill_holes() + vertices = m.vertices + faces = m.faces + coords = v.coords[:, 1:] + attrs = v.feats + if output_device is not None: + vertices = vertices.to(output_device) + faces = faces.to(output_device) + coords = coords.to(output_device) + attrs = attrs.to(output_device) out_mesh.append( MeshWithVoxel( - m.vertices, m.faces, + vertices, faces, origin = [-0.5, -0.5, -0.5], voxel_size = 1 / resolution, - coords = v.coords[:, 1:], - attrs = v.feats, + coords = coords, + attrs = attrs, voxel_shape = torch.Size([*v.shape, *v.spatial_shape]), layout=self.pbr_attr_layout ) ) + if output_device is not None: + del meshes, tex_voxels, subs + release_accelerator_memory( + f"decoded output moved to {output_device}", + verbose=release_models, + ) return out_mesh @torch.no_grad() @@ -619,6 +651,8 @@ def run( return_latent: bool = False, pipeline_type: Optional[str] = None, max_num_tokens: int = 49152, + release_models: bool = False, + output_device: Optional[Union[str, torch.device]] = None, ) -> List[MeshWithVoxel]: """ Run the Pixal3D pipeline (proj mode, cascade). @@ -638,6 +672,11 @@ def run( return_latent (bool): Whether to return the latent codes. pipeline_type (str): The type of the pipeline. Options: '1024_cascade', '1536_cascade'. max_num_tokens (int): The maximum number of tokens to use. + release_models (bool): Permanently discard one-shot models after + their final stage. Intended for single-image inference. + output_device: Move decoded mesh and voxel attributes to this + device before returning. ``"cpu"`` frees unified GPU memory + before native remeshing and texture baking. """ # Check pipeline type pipeline_type = pipeline_type or self.default_pipeline_type @@ -677,13 +716,21 @@ def run( distance=distance, mesh_scale=mesh_scale, ) + if release_models: + self.image_cond_model_ss = None ss_res = 32 coords = self.sample_sparse_structure( cond_ss, ss_res, num_samples, sparse_structure_sampler_params ) del cond_ss - torch.cuda.empty_cache() + if release_models: + drop_model(self, 'sparse_structure_flow_model') + drop_model(self, 'sparse_structure_decoder') + release_accelerator_memory( + "sparse-structure stage released", + verbose=release_models, + ) # ---- Stage 2: Shape LR 512 (proj) ---- cond_shape_lr = self.get_proj_cond_shape( @@ -692,12 +739,19 @@ def run( distance=distance, mesh_scale=mesh_scale, ) + if release_models: + self.image_cond_model_shape_512 = None lr_slat = self.sample_shape_slat( cond_shape_lr, self.models['shape_slat_flow_model_512'], coords, shape_slat_sampler_params ) del cond_shape_lr - torch.cuda.empty_cache() + if release_models: + drop_model(self, 'shape_slat_flow_model_512') + release_accelerator_memory( + "low-resolution shape stage released", + verbose=release_models, + ) # ---- Stage 3a: Upsample LR → HR ---- if self.low_vram: @@ -724,7 +778,10 @@ def run( actual_grid_res = actual_hr_resolution // 16 del lr_slat, hr_coords, quant_coords - torch.cuda.empty_cache() + release_accelerator_memory( + "shape-coordinate upsample released", + verbose=release_models, + ) # ---- Stage 3b: Shape HR (proj) ---- cond_shape_hr = self.get_proj_cond_shape( @@ -734,6 +791,8 @@ def run( mesh_scale=mesh_scale, grid_resolution_override=actual_grid_res, ) + if release_models: + self.image_cond_model_shape_1024 = None noise_hr = SparseTensor( feats=torch.randn(hr_coords_unique.shape[0], self.models['shape_slat_flow_model_1024'].in_channels).to(self.device), coords=hr_coords_unique, @@ -755,8 +814,21 @@ def run( std = torch.tensor(self.shape_slat_normalization['std'])[None].to(hr_slat.device) mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(hr_slat.device) shape_slat = hr_slat * std + mean - del cond_shape_hr, noise_hr, hr_slat, hr_coords_unique - torch.cuda.empty_cache() + del ( + cond_shape_hr, + noise_hr, + hr_slat, + hr_coords_unique, + flow_model_hr, + std, + mean, + ) + if release_models: + drop_model(self, 'shape_slat_flow_model_1024') + release_accelerator_memory( + "high-resolution shape stage released", + verbose=release_models, + ) # ---- Stage 4: Texture (proj) ---- tex_grid_res = actual_hr_resolution // 16 @@ -767,17 +839,35 @@ def run( mesh_scale=mesh_scale, grid_resolution_override=tex_grid_res, ) + if release_models: + self.image_cond_model_tex_1024 = None tex_slat = self.sample_tex_slat( cond_tex, self.models['tex_slat_flow_model_1024'], shape_slat, tex_slat_sampler_params ) del cond_tex - torch.cuda.empty_cache() + if release_models: + drop_model(self, 'tex_slat_flow_model_1024') + release_accelerator_memory( + "texture-flow stage released", + verbose=release_models, + ) # ---- Stage 5: Decode ---- res = actual_hr_resolution - out_mesh = self.decode_latent(shape_slat, tex_slat, res) + out_mesh = self.decode_latent( + shape_slat, + tex_slat, + res, + release_models=release_models, + output_device=output_device, + ) if return_latent: return out_mesh, (shape_slat, tex_slat, res) else: + del shape_slat, tex_slat + release_accelerator_memory( + "latent tensors released", + verbose=release_models, + ) return out_mesh diff --git a/pixal3d/pipelines/rembg/BiRefNet.py b/pixal3d/pipelines/rembg/BiRefNet.py index c71a992..af91c7d 100644 --- a/pixal3d/pipelines/rembg/BiRefNet.py +++ b/pixal3d/pipelines/rembg/BiRefNet.py @@ -19,8 +19,13 @@ def __init__(self, model_name: str = "ZhengPeng7/BiRefNet"): ] ) + @property + def device(self): + return next(self.model.parameters()).device + def to(self, device: str): self.model.to(device) + return self def cuda(self): self.model.cuda() @@ -30,7 +35,7 @@ def cpu(self): def __call__(self, image: Image.Image) -> Image.Image: image_size = image.size - input_images = self.transform_image(image).unsqueeze(0).to("cuda") + input_images = self.transform_image(image).unsqueeze(0).to(self.device) # Prediction with torch.no_grad(): preds = self.model(input_images)[-1].sigmoid().cpu() diff --git a/pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py b/pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py index ccafe87..66f2f01 100644 --- a/pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py +++ b/pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py @@ -368,6 +368,7 @@ def __init__( grid_resolution: int = 16, use_naf_upsample: bool = False, naf_target_size: Optional[List[int]] = None, + shared_model: Optional[nn.Module] = None, ): super().__init__() self.model_name = model_name @@ -382,7 +383,11 @@ def __init__( self.naf_target_size = tuple(naf_target_size) # Load DINOv3 model (frozen, no trainable params in this module) - self.model = DINOv3ViTModel.from_pretrained(model_name) + self.model = ( + shared_model + if shared_model is not None + else DINOv3ViTModel.from_pretrained(model_name) + ) self.model.eval() self.model.requires_grad_(False) @@ -420,6 +425,13 @@ def _load_naf(self): self.naf_model = torch.hub.load( "valeoai/NAF", "naf", pretrained=True, device=device, trust_repo=True ) + # Preserve the learned NAF upsampler on Apple Silicon. Only its + # CUDA-only NATTEN operation is replaced by an equivalent, + # row-chunked PyTorch/MPS implementation. + if device.type == "mps" or not torch.cuda.is_available(): + from backends.naf_attention import install_chunked_naf_attention + + install_chunked_naf_attention(self.naf_model) self.naf_model.eval() self.naf_model.requires_grad_(False) @@ -431,12 +443,16 @@ def to(self, device): self.naf_model.to(device) return self + @property + def device(self): + return next(self.parameters()).device + def cuda(self): - super().cuda() - self.model.cuda() - self.proj_grid.cuda() + super().to(self.device) + self.model.to(self.device) + self.proj_grid.to(self.device) if self.naf_model is not None: - self.naf_model.cuda() + self.naf_model.to(self.device) return self def cpu(self): @@ -493,7 +509,7 @@ def forward( image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") @@ -692,12 +708,16 @@ def to(self, device): self._vae.to(device) return self + @property + def device(self): + return next(self.parameters()).device + def cuda(self): - super().cuda() - self.dino_model.cuda() - self.proj_grid.cuda() + super().to(self.device) + self.dino_model.to(self.device) + self.proj_grid.to(self.device) if self._vae is not None: - self._vae.cuda() + self._vae.to(self.device) return self def cpu(self): @@ -756,7 +776,7 @@ def forward( image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") @@ -836,7 +856,7 @@ def _init_image_cond_model(self): from . import image_conditioned self.image_cond_model = getattr(image_conditioned, model_name)(**model_args) - self.image_cond_model.cuda() + self.image_cond_model.to(self.device) # Expose proj_channels for denoiser to know the correct proj_in_channels if hasattr(self.image_cond_model, 'proj_channels'): diff --git a/requirements-macos.txt b/requirements-macos.txt new file mode 100644 index 0000000..870d003 --- /dev/null +++ b/requirements-macos.txt @@ -0,0 +1,3 @@ +# CUDA-free additions required by the Apple/Metal o_voxel postprocess path. +xatlas +fast-simplification diff --git a/runpod/README.md b/runpod/README.md new file mode 100644 index 0000000..59841b3 --- /dev/null +++ b/runpod/README.md @@ -0,0 +1,30 @@ +# Artefacts RunPod Pixal3D + +Sauvegarde ciblée de l’environnement CUDA utilisé pour Pixal3D sur RunPod. + +Les archives binaires sont volontairement conservées uniquement en local : +elles sont volumineuses et ne sont pas nécessaires au portage macOS/MPS. + +## Contenu + +- `pixal3d-cuda-binaries.tar.gz` — 669 Mo (local uniquement) : NATTEN, FlashAttention, `o_voxel`, `cumesh` et `flex_gemm`. +- `pixal3d-cuda-extras.tar.gz` — 3,9 Mo (local uniquement) : `nvdiffrast` et `nvdiffrec_render`. +- `pixal3d-python-lock.txt` — versions exactes observées dans `/workspace/venv`. + +Les modèles Hugging Face, les caches et le dépôt Pixal3D ne sont pas inclus. + +## Compatibilité + +Les archives ciblent l’environnement Linux suivant : Python 3.11, PyTorch 2.6.0 + CUDA 12.4. Elles ne sont pas utilisables directement dans l’environnement macOS/MPS. + +## Restauration + +Sur un nouveau pod compatible, installer d’abord Python et les dépendances générales, puis extraire les extensions dans l’environnement virtuel : + +```bash +SITE=/workspace/venv/lib/python3.11/site-packages +tar -xzf pixal3d-cuda-binaries.tar.gz -C "$SITE" +tar -xzf pixal3d-cuda-extras.tar.gz -C "$SITE" +``` + +Le fichier de lock contient quelques références locales au pod (`file:///tmp/...` et `file:///workspace/...`) ; il sert de référence et de contrôle, mais ne doit pas être utilisé tel quel comme unique commande `pip install -r`. diff --git a/runpod/pixal3d-python-lock.txt b/runpod/pixal3d-python-lock.txt new file mode 100644 index 0000000..ab639b4 --- /dev/null +++ b/runpod/pixal3d-python-lock.txt @@ -0,0 +1,76 @@ +accelerate==1.13.0 +anyio==4.14.2 +certifi==2026.7.22 +charset-normalizer==3.4.9 +click==8.4.2 +cumesh @ git+https://github.com/JeffreyXiang/CuMesh.git@12289e1062f0603f2f0d0771b02e1395d247f26f +diffusers==0.37.1 +easydict==1.13 +einops==0.8.2 +filelock==3.29.0 +flash-attn==2.7.3 +flex_gemm @ git+https://github.com/JeffreyXiang/FlexGEMM.git@6dd94a859c26ee8246888502eada3dd8ad85532e +fsspec==2026.4.0 +glcontext==3.0.0 +h11==0.16.0 +hf-xet==1.5.2 +httpcore==1.0.9 +httpx==0.28.1 +huggingface_hub==0.36.2 +idna==3.18 +ImageIO==2.37.2 +imageio-ffmpeg==0.6.0 +importlib_metadata==9.0.0 +Jinja2==3.1.6 +kornia==0.8.2 +kornia_rs==0.1.14 +MarkupSafe==3.0.3 +moderngl==5.12.0 +moge @ git+https://github.com/microsoft/MoGe.git@925b8ed835a7a9cdb7578ba15c658a0afc969030 +mpmath==1.3.0 +natten==0.17.5+torch260cu124 +networkx==3.6.1 +ninja==1.13.0 +numpy==2.2.6 +nvdiffrast @ file:///tmp/extensions/nvdiffrast +nvdiffrec_render @ file:///tmp/extensions/nvdiffrec +nvidia-cublas-cu12==12.4.5.8 +nvidia-cuda-cupti-cu12==12.4.127 +nvidia-cuda-nvrtc-cu12==12.4.127 +nvidia-cuda-runtime-cu12==12.4.127 +nvidia-cudnn-cu12==9.1.0.70 +nvidia-cufft-cu12==11.2.1.3 +nvidia-curand-cu12==10.3.5.147 +nvidia-cusolver-cu12==11.6.1.9 +nvidia-cusparse-cu12==12.3.1.170 +nvidia-cusparselt-cu12==0.6.2 +nvidia-nccl-cu12==2.21.5 +nvidia-nvjitlink-cu12==12.4.127 +nvidia-nvtx-cu12==12.4.127 +o_voxel @ file:///workspace/TRELLIS.2/o-voxel +opencv-python-headless==4.12.0.88 +packaging==26.2 +pillow==12.0.0 +pipeline @ git+https://github.com/EasternJournalist/pipeline.git@866f059d2a05cde05e4a52211ec5051fd5f276d6 +plyfile==1.1.3 +psutil==7.2.2 +pybind11==3.0.4 +PyYAML==6.0.3 +regex==2026.7.19 +requests==2.34.2 +safetensors==0.8.0 +scipy==1.17.1 +sympy==1.13.1 +timm==1.0.22 +tokenizers==0.22.2 +torch==2.6.0+cu124 +torchvision==0.21.0+cu124 +tqdm==4.67.1 +transformers==4.57.3 +trimesh==4.10.1 +triton==3.2.0 +typing_extensions==4.15.0 +urllib3==2.7.0 +utils3d @ https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl#sha256=ff63440827d6933807dd06c8a5a2db7e51fd5f33c7f3dddcc766a80e0f419252 +zipp==4.1.0 +zstandard==0.25.0 diff --git a/scripts/compare_glb_quality.py b/scripts/compare_glb_quality.py new file mode 100644 index 0000000..92858b6 --- /dev/null +++ b/scripts/compare_glb_quality.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Compare structural GLB quality against the official CUDA reference.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +import trimesh + +from backends.export_validation import weld_exact_float32 +from backends.mesh_postprocess import mesh_metrics + + +def _single_geometry(path: str | Path): + loaded = trimesh.load(path, force="scene", process=False) + geometries = list(loaded.geometry.values()) + if len(geometries) != 1: + raise ValueError(f"{path} contains {len(geometries)} geometries") + return geometries[0] + + +def _texture_report(material: Any) -> dict[str, Any]: + texture = getattr(material, "baseColorTexture", None) + if texture is None: + return {"present": False} + pixels = np.asarray(texture) + report: dict[str, Any] = { + "present": True, + "shape": list(pixels.shape), + } + if pixels.ndim == 3 and pixels.shape[-1] >= 4: + alpha = pixels[..., 3] + report["alpha"] = { + "min": int(alpha.min()), + "median": float(np.median(alpha)), + "p01": float(np.percentile(alpha, 1)), + "below_250_share": float(np.mean(alpha < 250)), + } + return report + + +def _surface_area( + vertices: np.ndarray, + faces: np.ndarray, + *, + chunk_size: int = 250_000, +) -> float: + """Measure area without materializing every triangle at once.""" + + total = 0.0 + for start in range(0, len(faces), chunk_size): + triangles = vertices[faces[start : start + chunk_size]] + cross = np.cross( + triangles[:, 1] - triangles[:, 0], + triangles[:, 2] - triangles[:, 0], + ) + total += float(np.linalg.norm(cross, axis=1).sum()) * 0.5 + return total + + +def inspect(path: str | Path) -> dict[str, Any]: + geometry = _single_geometry(path) + vertices = np.asarray(geometry.vertices, dtype=np.float32) + faces = np.asarray(geometry.faces, dtype=np.int64) + welded_vertices, welded_faces = weld_exact_float32(vertices, faces) + material = getattr(geometry.visual, "material", None) + topology = mesh_metrics(welded_vertices, welded_faces) + topology["surface_area"] = _surface_area( + welded_vertices, + welded_faces, + ) + return { + "path": str(Path(path).resolve()), + "bytes": Path(path).stat().st_size, + "exported_vertices": int(len(vertices)), + "welded_vertices": int(len(welded_vertices)), + "material": { + "alpha_mode": getattr(material, "alphaMode", None), + "double_sided": getattr(material, "doubleSided", None), + "base_color": _texture_report(material), + }, + "topology": topology, + } + + +def compare(reference: dict[str, Any], candidate: dict[str, Any]) -> dict[str, Any]: + ref_topology = reference["topology"] + candidate_topology = candidate["topology"] + texture = candidate["material"]["base_color"] + alpha = texture.get("alpha", {}) + texture_shape = texture.get("shape", [0, 0]) + + checks = { + "opaque_material": candidate["material"]["alpha_mode"] == "OPAQUE", + "single_sided_material": candidate["material"]["double_sided"] is False, + "texture_4096": min(texture_shape[:2], default=0) >= 4096, + "mostly_opaque_texture": ( + alpha.get("median", 0) >= 250 + and alpha.get("below_250_share", 1) <= 0.01 + ), + "million_face_profile": 850_000 + <= candidate_topology["faces"] + <= 1_050_000, + "dominant_surface": ( + candidate_topology["largest_edge_component_share"] >= 0.99 + ), + "boundary_near_reference": ( + candidate_topology["boundary_edges"] + <= max(1_000, 10 * ref_topology["boundary_edges"]) + ), + "surface_area_near_reference": ( + 0.9 + <= candidate_topology["surface_area"] + / ref_topology["surface_area"] + <= 1.1 + ), + } + return { + "checks": checks, + "passed": all(checks.values()), + "face_ratio": ( + candidate_topology["faces"] / ref_topology["faces"] + ), + "surface_area_ratio": ( + candidate_topology["surface_area"] + / ref_topology["surface_area"] + ), + "boundary_edge_delta": ( + candidate_topology["boundary_edges"] + - ref_topology["boundary_edges"] + ), + "nonmanifold_edge_delta": ( + candidate_topology["nonmanifold_edges"] + - ref_topology["nonmanifold_edges"] + ), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--reference", required=True) + parser.add_argument("--candidate", required=True) + parser.add_argument("--output", default=None) + args = parser.parse_args() + + reference = inspect(args.reference) + candidate = inspect(args.candidate) + report = { + "reference": reference, + "candidate": candidate, + "comparison": compare(reference, candidate), + } + rendered = json.dumps(report, indent=2, sort_keys=True) + print(rendered) + if args.output: + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(rendered + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/diagnose_checkpoint_remesh.py b/scripts/diagnose_checkpoint_remesh.py new file mode 100644 index 0000000..b9d996a --- /dev/null +++ b/scripts/diagnose_checkpoint_remesh.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Remesh a decoded Pixal3D checkpoint and persist the raw Metal candidate.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + +import numpy as np +import torch + +from backends.decoded_checkpoint import load_decoded_checkpoint +from backends.chunked_mtlbvh import ChunkedMtlBVH +from backends.mesh_postprocess import mesh_metrics + + +def surface_area(vertices: np.ndarray, faces: np.ndarray) -> float: + triangles = vertices[faces] + cross = np.cross( + triangles[:, 1] - triangles[:, 0], + triangles[:, 2] - triangles[:, 0], + ) + return float(np.linalg.norm(cross, axis=1).sum() * 0.5) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--resolution", type=int) + parser.add_argument("--band", type=float, default=1.0) + parser.add_argument("--project-back", type=float, default=0.0) + parser.add_argument("--refine-factor", type=float, default=0.87) + parser.add_argument("--source-face-chunk-size", type=int, default=0) + parser.add_argument("--query-chunk-size", type=int, default=262_144) + parser.add_argument("--output-prefix", type=Path, required=True) + parser.add_argument("--verbose", action="store_true") + args = parser.parse_args() + + if not torch.backends.mps.is_available(): + raise RuntimeError("Checkpoint remeshing requires Metal/MPS") + os.environ["CUMESH_REMESH_REFINE_FACTOR"] = str(args.refine_factor) + + mesh, source_resolution, _ = load_decoded_checkpoint(args.checkpoint) + resolution = args.resolution or source_resolution + base_scale = 1.0 + scale = ( + (resolution + 3 * args.band) + / resolution + * base_scale + ) + + from cumesh import remeshing + from mtlbvh import MtlBVH + + vertices_t = mesh.vertices.detach().cpu().float().contiguous() + faces_t = mesh.faces.detach().cpu().int().contiguous() + if args.source_face_chunk_size: + bvh = ChunkedMtlBVH( + MtlBVH, + vertices_t, + faces_t, + source_face_chunk_size=args.source_face_chunk_size, + query_chunk_size=args.query_chunk_size, + ) + else: + bvh = MtlBVH(vertices_t, faces_t) + started = time.perf_counter() + candidate_vertices_t, candidate_faces_t = remeshing.remesh_narrow_band_dc( + vertices_t, + faces_t, + center=torch.zeros(3, dtype=torch.float32), + scale=scale, + resolution=resolution, + band=args.band, + project_back=args.project_back, + verbose=args.verbose, + bvh=bvh, + ) + seconds = time.perf_counter() - started + + candidate_vertices = ( + candidate_vertices_t.detach().cpu().float().numpy() + ) + candidate_faces = candidate_faces_t.detach().cpu().int().numpy() + + args.output_prefix.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + args.output_prefix.with_suffix(".npz"), + vertices=candidate_vertices.astype(np.float32), + faces=candidate_faces.astype(np.int32), + ) + report = { + "checkpoint": str(args.checkpoint), + "source_resolution": source_resolution, + "remesh_resolution": resolution, + "band": args.band, + "project_back": args.project_back, + "refine_factor": args.refine_factor, + "source_face_chunk_size": args.source_face_chunk_size, + "query_chunk_size": args.query_chunk_size, + "seconds": seconds, + "source": { + "vertices": int(len(vertices_t)), + "faces": int(len(faces_t)), + }, + "candidate": { + **mesh_metrics(candidate_vertices, candidate_faces), + "surface_area": surface_area( + candidate_vertices, + candidate_faces, + ), + }, + } + args.output_prefix.with_suffix(".json").write_text( + json.dumps(report, indent=2, sort_keys=True), + encoding="utf-8", + ) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/diagnose_metal_cleanup.py b/scripts/diagnose_metal_cleanup.py new file mode 100644 index 0000000..d5fafd4 --- /dev/null +++ b/scripts/diagnose_metal_cleanup.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Probe CuMesh/Metal cleanup on a previously reconstructed mesh candidate. + +This keeps the expensive UDF reconstruction and the cleanup experiment +separate. It deliberately skips ``fill_holes`` (which can stall on large +Metal meshes) and ``unify_face_orientations`` (which has damaged disconnected +decoded surfaces in earlier tests). +""" + +from __future__ import annotations + +import argparse +import json +import resource +import time +from pathlib import Path + +import numpy as np +import torch + +from backends.mesh_postprocess import mesh_metrics + + +def _memory_snapshot() -> dict[str, float]: + snapshot = { + "process_peak_gib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + / 1024**3, + } + if torch.backends.mps.is_available(): + snapshot["mps_allocated_gib"] = torch.mps.current_allocated_memory() / 1024**3 + snapshot["mps_driver_gib"] = torch.mps.driver_allocated_memory() / 1024**3 + return snapshot + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("candidate", type=Path) + parser.add_argument("--target-faces", type=int, default=1_000_000) + parser.add_argument("--output-prefix", type=Path, required=True) + parser.add_argument("--fill-holes", action="store_true") + parser.add_argument("--fill-after-simplify", action="store_true") + parser.add_argument("--max-hole-perimeter", type=float, default=3e-2) + parser.add_argument("--verbose", action="store_true") + args = parser.parse_args() + + if not torch.backends.mps.is_available(): + raise RuntimeError("Metal cleanup requires MPS") + + archive = np.load(args.candidate) + vertices = np.ascontiguousarray(archive["vertices"], dtype=np.float32) + faces = np.ascontiguousarray(archive["faces"], dtype=np.int32) + + from cumesh import CuMesh + + mesh = CuMesh() + stages: list[dict[str, object]] = [] + + def run_stage(name: str, operation) -> None: + started = time.perf_counter() + operation() + torch.mps.synchronize() + stage = { + "name": name, + "seconds": time.perf_counter() - started, + "vertices": int(mesh.num_vertices), + "faces": int(mesh.num_faces), + "memory": _memory_snapshot(), + } + stages.append(stage) + if args.verbose: + print(json.dumps(stage), flush=True) + + run_stage( + "init", + lambda: mesh.init(torch.from_numpy(vertices), torch.from_numpy(faces)), + ) + run_stage("remove_duplicate_faces_1", mesh.remove_duplicate_faces) + run_stage("remove_degenerate_faces_1", mesh.remove_degenerate_faces) + run_stage("repair_non_manifold_edges_1", mesh.repair_non_manifold_edges) + run_stage( + "remove_small_components_1", + lambda: mesh.remove_small_connected_components(1e-5), + ) + if args.fill_holes: + run_stage( + "fill_holes_1", + lambda: mesh.fill_holes(args.max_hole_perimeter), + ) + run_stage( + "simplify", + lambda: mesh.simplify(args.target_faces, verbose=args.verbose), + ) + run_stage("remove_duplicate_faces_2", mesh.remove_duplicate_faces) + run_stage("remove_degenerate_faces_2", mesh.remove_degenerate_faces) + run_stage("repair_non_manifold_edges_2", mesh.repair_non_manifold_edges) + run_stage( + "remove_small_components_2", + lambda: mesh.remove_small_connected_components(1e-5), + ) + if args.fill_after_simplify: + run_stage( + "fill_holes_2", + lambda: mesh.fill_holes(args.max_hole_perimeter), + ) + + cleaned_vertices_t, cleaned_faces_t = mesh.read() + cleaned_vertices = cleaned_vertices_t.detach().cpu().float().numpy() + cleaned_faces = cleaned_faces_t.detach().cpu().int().numpy() + del cleaned_vertices_t, cleaned_faces_t, mesh + torch.mps.synchronize() + torch.mps.empty_cache() + + args.output_prefix.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + args.output_prefix.with_suffix(".npz"), + vertices=cleaned_vertices.astype(np.float32), + faces=cleaned_faces.astype(np.int32), + ) + report = { + "candidate": str(args.candidate), + "target_faces": args.target_faces, + "fill_holes": args.fill_holes, + "fill_after_simplify": args.fill_after_simplify, + "max_hole_perimeter": args.max_hole_perimeter, + "input": mesh_metrics(vertices, faces), + "output": mesh_metrics(cleaned_vertices, cleaned_faces), + "stages": stages, + "final_memory": _memory_snapshot(), + } + args.output_prefix.with_suffix(".json").write_text( + json.dumps(report, indent=2, sort_keys=True), + encoding="utf-8", + ) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/diagnose_metal_fill_holes.py b/scripts/diagnose_metal_fill_holes.py new file mode 100644 index 0000000..451b666 --- /dev/null +++ b/scripts/diagnose_metal_fill_holes.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +"""Measure CuMesh/Metal hole filling on a saved geometry candidate.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import torch + +from backends.mesh_postprocess import mesh_metrics + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("candidate", type=Path) + parser.add_argument("--max-perimeter", type=float, default=3e-2) + parser.add_argument("--output-prefix", type=Path, required=True) + args = parser.parse_args() + + if not torch.backends.mps.is_available(): + raise RuntimeError("Metal hole filling requires MPS") + + archive = np.load(args.candidate) + vertices = np.ascontiguousarray(archive["vertices"], dtype=np.float32) + faces = np.ascontiguousarray(archive["faces"], dtype=np.int32) + + from cumesh import CuMesh + + mesh = CuMesh() + mesh.init(torch.from_numpy(vertices), torch.from_numpy(faces)) + started = time.perf_counter() + mesh.fill_holes(max_hole_perimeter=args.max_perimeter) + torch.mps.synchronize() + seconds = time.perf_counter() - started + + output_vertices_t, output_faces_t = mesh.read() + output_vertices = output_vertices_t.detach().cpu().float().numpy() + output_faces = output_faces_t.detach().cpu().int().numpy() + + args.output_prefix.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + args.output_prefix.with_suffix(".npz"), + vertices=output_vertices.astype(np.float32), + faces=output_faces.astype(np.int32), + ) + report = { + "candidate": str(args.candidate), + "max_perimeter": args.max_perimeter, + "seconds": seconds, + "input": mesh_metrics(vertices, faces), + "output": mesh_metrics(output_vertices, output_faces), + } + args.output_prefix.with_suffix(".json").write_text( + json.dumps(report, indent=2, sort_keys=True), + encoding="utf-8", + ) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/diagnose_mtlbvh_coverage.py b/scripts/diagnose_mtlbvh_coverage.py new file mode 100644 index 0000000..49bb4c1 --- /dev/null +++ b/scripts/diagnose_mtlbvh_coverage.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""Check whether MtlBVH covers every region of a very large source mesh.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import torch + +from backends.decoded_checkpoint import load_decoded_checkpoint + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--bins", type=int, default=18) + parser.add_argument("--samples-per-bin", type=int, default=512) + parser.add_argument("--per-bin-bvh", action="store_true") + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + if args.bins <= 0 or args.samples_per_bin <= 0: + raise ValueError("bins and samples-per-bin must be positive") + + mesh, resolution, _ = load_decoded_checkpoint(args.checkpoint) + vertices = mesh.vertices.detach().cpu().float().contiguous() + faces = mesh.faces.detach().cpu().int().contiguous() + + from mtlbvh import MtlBVH + + bvh = None + build_seconds = 0.0 + if not args.per_bin_bvh: + built = time.perf_counter() + bvh = MtlBVH(vertices, faces) + build_seconds = time.perf_counter() - built + + bin_edges = np.linspace(0, len(faces), args.bins + 1, dtype=np.int64) + rows: list[dict[str, object]] = [] + all_distances: list[np.ndarray] = [] + queried = time.perf_counter() + for bin_index, (start, stop) in enumerate( + zip(bin_edges[:-1], bin_edges[1:]) + ): + if args.per_bin_bvh: + built = time.perf_counter() + bvh = MtlBVH(vertices, faces[int(start) : int(stop)]) + build_seconds += time.perf_counter() - built + assert bvh is not None + indices = torch.linspace( + int(start), + int(stop - 1), + steps=min(args.samples_per_bin, int(stop - start)), + dtype=torch.float64, + ).round().long() + triangles = vertices[faces[indices].long()] + centroids = triangles.mean(dim=1).contiguous() + distances, face_ids, _ = bvh.unsigned_distance( + centroids, + return_uvw=True, + ) + values = distances.detach().cpu().float().numpy() + all_distances.append(values) + rows.append( + { + "bin": bin_index, + "face_start": int(start), + "face_stop": int(stop), + "samples": int(len(values)), + "distance_voxels": { + "median": float(np.median(values) * resolution), + "p95": float(np.percentile(values, 95) * resolution), + "max": float(np.max(values) * resolution), + }, + "exact_source_face_share": float( + np.mean( + face_ids.detach().cpu().numpy() + == ( + indices.numpy() - int(start) + if args.per_bin_bvh + else indices.numpy() + ) + ) + ), + } + ) + if args.per_bin_bvh: + del bvh + bvh = None + query_seconds = time.perf_counter() - queried + combined = np.concatenate(all_distances) + report = { + "checkpoint": str(args.checkpoint), + "vertices": int(len(vertices)), + "faces": int(len(faces)), + "resolution": resolution, + "per_bin_bvh": args.per_bin_bvh, + "build_seconds": build_seconds, + "query_seconds": query_seconds, + "overall_distance_voxels": { + "median": float(np.median(combined) * resolution), + "p95": float(np.percentile(combined, 95) * resolution), + "max": float(np.max(combined) * resolution), + }, + "bins": rows, + } + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(report, indent=2, sort_keys=True), + encoding="utf-8", + ) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/diagnose_udf_remesh.py b/scripts/diagnose_udf_remesh.py new file mode 100644 index 0000000..b8cf8f4 --- /dev/null +++ b/scripts/diagnose_udf_remesh.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Run the Metal UDF remesher on an existing GLB and save diagnostics. + +This is intentionally a geometry-only probe. It lets us measure topology, +surface drift, runtime, and output density before spending another full +Pixal3D inference on a new export profile. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import trimesh + +from backends.mesh_postprocess import UDFConfig, run_udf_postprocess + + +def _load_welded_geometry(path: Path) -> tuple[np.ndarray, np.ndarray]: + scene = trimesh.load(path, force="scene") + if len(scene.geometry) != 1: + raise RuntimeError( + f"Expected one geometry in {path}, found {len(scene.geometry)}" + ) + mesh = next(iter(scene.geometry.values())).copy() + # GLB UV seams duplicate positions. Seven decimal digits preserve the + # float32 geometry while restoring the indexed surface used by remeshing. + mesh.merge_vertices(digits_vertex=7, merge_tex=True, merge_norm=True) + return ( + np.ascontiguousarray(mesh.vertices, dtype=np.float32), + np.ascontiguousarray(mesh.faces, dtype=np.int32), + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("input", type=Path) + parser.add_argument("--resolution", type=int, default=512) + parser.add_argument("--output-prefix", type=Path, required=True) + parser.add_argument("--verbose", action="store_true") + args = parser.parse_args() + + vertices, faces = _load_welded_geometry(args.input) + config = UDFConfig(resolution=args.resolution) + result = run_udf_postprocess( + vertices, + faces, + config, + verbose=args.verbose, + ) + + args.output_prefix.parent.mkdir(parents=True, exist_ok=True) + report_path = args.output_prefix.with_suffix(".json") + report_path.write_text( + json.dumps(result.report, indent=2, sort_keys=True), + encoding="utf-8", + ) + + if result.candidate_vertices is not None: + candidate_path = args.output_prefix.with_suffix(".npz") + np.savez_compressed( + candidate_path, + vertices=result.candidate_vertices.astype(np.float32), + faces=result.candidate_faces.astype(np.int32), + ) + + print(json.dumps(result.report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/doctor.py b/scripts/doctor.py new file mode 100644 index 0000000..0083705 --- /dev/null +++ b/scripts/doctor.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +"""Inspect whether this Mac is ready for quality-first TRELLIS.2 inference.""" + +from __future__ import annotations + +import argparse +import importlib +import json +import platform +import subprocess +import sys +from typing import Any + + +METAL_MODULES = ("flex_gemm", "cumesh", "mtldiffrast", "o_voxel") + + +def _command(*args: str) -> str | None: + try: + return subprocess.check_output(args, text=True, stderr=subprocess.STDOUT).strip() + except (OSError, subprocess.CalledProcessError): + return None + + +def _module_status(name: str) -> dict[str, str | bool]: + try: + module = importlib.import_module(name) + return {"available": True, "version": str(getattr(module, "__version__", "unknown"))} + except Exception as exc: # An incompatible metallib should be reported too. + return {"available": False, "error": f"{type(exc).__name__}: {exc}"} + + +def collect_report() -> dict[str, Any]: + report: dict[str, Any] = { + "platform": platform.platform(), + "machine": platform.machine(), + "python": sys.version.split()[0], + "macos": _command("sw_vers", "-productVersion"), + "metal_sdk": _command("xcrun", "--sdk", "macosx", "--show-sdk-version"), + "modules": {name: _module_status(name) for name in METAL_MODULES}, + } + try: + import torch + + report["torch"] = { + "version": torch.__version__, + "mps_built": torch.backends.mps.is_built(), + "mps_available": torch.backends.mps.is_available(), + } + except Exception as exc: + report["torch"] = {"available": False, "error": f"{type(exc).__name__}: {exc}"} + return report + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--json", action="store_true", help="Print compact JSON only.") + parser.add_argument("--require-metal", action="store_true", help="Exit non-zero unless all Metal modules work.") + args = parser.parse_args() + + report = collect_report() + if args.json: + print(json.dumps(report, sort_keys=True)) + else: + print(json.dumps(report, indent=2, sort_keys=True)) + + if args.require_metal: + torch_ok = bool(report.get("torch", {}).get("mps_available")) + modules_ok = all(bool(report["modules"][name].get("available")) for name in METAL_MODULES) + if not torch_ok or not modules_ok: + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/download_models.py b/scripts/download_models.py new file mode 100644 index 0000000..4a4bf93 --- /dev/null +++ b/scripts/download_models.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +"""Pre-cache the public Pixal3D inference models in the Hugging Face cache.""" + +from __future__ import annotations + +import argparse + +from huggingface_hub import snapshot_download + + +REPOS = ( + "TencentARC/Pixal3D", + "Ruicheng/moge-2-vitl", + "camenduru/dinov3-vitl16-pretrain-lvd1689m", + "ZhengPeng7/BiRefNet", +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--revision", default=None) + args = parser.parse_args() + for repo_id in REPOS: + print(f"\nDownloading {repo_id} ...") + snapshot_download(repo_id=repo_id, revision=args.revision) + print("\nPixal3D model files are cached.") + print("On MPS, the CUDA-only NAF upsampler is replaced by the built-in bilinear fallback.") + + +if __name__ == "__main__": + main() diff --git a/scripts/export_prepared_geometry.py b/scripts/export_prepared_geometry.py new file mode 100644 index 0000000..fd616aa --- /dev/null +++ b/scripts/export_prepared_geometry.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +"""Bake a decoded Pixal3D PBR volume onto an already prepared Metal mesh.""" + +from __future__ import annotations + +import argparse +import gc +import json +import time +from pathlib import Path + +import numpy as np +import torch + +from macos_compat import configure + +configure() + +from backends.cuda_parity_export import ( # noqa: E402 + force_cuda_material_semantics, + memory_bounded_o_voxel, +) +from backends.metal_preserve import ( # noqa: E402 + texture_projection_kwargs, + use_geometry_preserving_backend, +) + + +def _deserialize_layout( + layout: dict[str, list[int | None]], +) -> dict[str, slice]: + return {name: slice(*value) for name, value in layout.items()} + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "UV unwrap and bake a decoded Pixal3D attribute volume without " + "changing the supplied mesh geometry." + ) + ) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("geometry", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--texture-size", type=int, default=4096) + parser.add_argument("--grid-chunk-size", type=int, default=262_144) + parser.add_argument("--source-face-chunk-size", type=int, default=250_000) + parser.add_argument( + "--projection-query-chunk-size", + type=int, + default=262_144, + ) + parser.add_argument("--projection-cache", type=Path) + parser.add_argument("--skip-source-projection", action="store_true") + parser.add_argument("--report", type=Path) + parser.add_argument("--quiet", action="store_true") + args = parser.parse_args() + + if ( + args.texture_size <= 0 + or args.grid_chunk_size <= 0 + or args.source_face_chunk_size <= 8 + or args.projection_query_chunk_size <= 0 + ): + raise ValueError("texture and chunk sizes must be positive") + if not torch.backends.mps.is_available(): + raise RuntimeError("Prepared geometry export requires Metal/MPS") + + started = time.perf_counter() + payload = torch.load( + args.checkpoint, + map_location="cpu", + weights_only=True, + ) + resolution = int(payload["resolution"]) + coords = payload["coords"].detach().cpu().contiguous() + attrs = payload["attrs"].detach().cpu().contiguous() + attr_layout = _deserialize_layout(payload["layout"]) + checkpoint_metadata = dict(payload.get("metadata", {})) + + with np.load(args.geometry) as archive: + vertices_np = np.ascontiguousarray( + archive["vertices"], + dtype=np.float32, + ) + faces_np = np.ascontiguousarray( + archive["faces"], + dtype=np.int32, + ) + vertices = torch.from_numpy(vertices_np) + faces = torch.from_numpy(faces_np) + + projection_seconds = 0.0 + projection_distance_voxels: dict[str, float] | None = None + texture_sample_vertices: torch.Tensor | None = None + texture_fallback_projector = None + fallback_report: dict[str, float | int] = { + "queries": 0, + "seconds": 0.0, + } + if not args.skip_source_projection: + projection_started = time.perf_counter() + from backends.chunked_mtlbvh import ChunkedMtlBVH + from mtlbvh import MtlBVH + + source_vertices = ( + payload["vertices"].detach().cpu().float().contiguous() + ) + source_faces = ( + payload["faces"].detach().cpu().int().contiguous() + ) + source_bvh = ChunkedMtlBVH( + MtlBVH, + source_vertices, + source_faces, + source_face_chunk_size=args.source_face_chunk_size, + query_chunk_size=args.projection_query_chunk_size, + ) + + def project_to_source( + positions: torch.Tensor, + ) -> torch.Tensor: + query_started = time.perf_counter() + distances, source_face_ids, uvw = ( + source_bvh.unsigned_distance( + positions, + return_uvw=True, + ) + ) + assert uvw is not None + source_triangles = source_vertices[ + source_faces[source_face_ids.long()] + ] + projected = ( + source_triangles * uvw.unsqueeze(-1) + ).sum(dim=1).contiguous() + fallback_report["queries"] = ( + int(fallback_report["queries"]) + len(positions) + ) + fallback_report["seconds"] = ( + float(fallback_report["seconds"]) + + time.perf_counter() + - query_started + ) + del ( + distances, + source_face_ids, + uvw, + source_triangles, + ) + return projected + + texture_fallback_projector = project_to_source + if args.projection_cache is not None and args.projection_cache.exists(): + with np.load(args.projection_cache) as projection_archive: + projected_np = np.ascontiguousarray( + projection_archive["texture_vertices"], + dtype=np.float32, + ) + if projected_np.shape != vertices_np.shape: + raise ValueError( + "Cached texture vertices do not match prepared geometry" + ) + texture_sample_vertices = torch.from_numpy(projected_np) + else: + distances, source_face_ids, uvw = source_bvh.unsigned_distance( + vertices, + return_uvw=True, + ) + assert uvw is not None + source_triangles = source_vertices[ + source_faces[source_face_ids.long()] + ] + texture_sample_vertices = ( + source_triangles * uvw.unsqueeze(-1) + ).sum(dim=1).contiguous() + distance_voxels = distances.float() * resolution + projection_distance_voxels = { + "median": float(distance_voxels.median()), + "p95": float(torch.quantile(distance_voxels, 0.95)), + "p99": float(torch.quantile(distance_voxels, 0.99)), + "max": float(distance_voxels.max()), + } + if args.projection_cache is not None: + args.projection_cache.parent.mkdir( + parents=True, + exist_ok=True, + ) + np.savez_compressed( + args.projection_cache, + texture_vertices=( + texture_sample_vertices.detach().cpu().numpy() + ), + ) + del ( + distances, + source_face_ids, + uvw, + source_triangles, + distance_voxels, + ) + projection_seconds = time.perf_counter() - projection_started + + del payload + gc.collect() + torch.mps.empty_cache() + + import o_voxel.postprocess as postprocess + + if getattr(postprocess, "_BACKEND", None) != "metal": + raise RuntimeError("Prepared geometry export expected the Metal backend") + + projection_kwargs = texture_projection_kwargs(postprocess, "preserve") + bake_started = time.perf_counter() + with ( + memory_bounded_o_voxel( + postprocess, + bvh_chunk_size=262_144, + grid_chunk_size=args.grid_chunk_size, + source_resolution=resolution, + ), + use_geometry_preserving_backend(postprocess), + ): + glb = postprocess.to_glb( + vertices=vertices, + faces=faces, + attr_volume=attrs, + coords=coords, + attr_layout=attr_layout, + grid_size=resolution, + aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], + decimation_target=max(int(faces.shape[0]), 1), + texture_size=args.texture_size, + remesh=False, + remesh_band=1, + remesh_project=0, + verbose=not args.quiet, + use_tqdm=not args.quiet, + texture_sample_vertices=texture_sample_vertices, + texture_fallback_projector=texture_fallback_projector, + **projection_kwargs, + ) + bake_seconds = time.perf_counter() - bake_started + + force_cuda_material_semantics(glb) + rotation = np.array( + [ + [-1, 0, 0, 0], + [0, 0, -1, 0], + [0, -1, 0, 0], + [0, 0, 0, 1], + ], + dtype=np.float64, + ) + glb.apply_transform(rotation) + + args.output.parent.mkdir(parents=True, exist_ok=True) + glb.export(args.output, extension_webp=True) + total_seconds = time.perf_counter() - started + + report = { + "checkpoint": str(args.checkpoint), + "checkpoint_metadata": checkpoint_metadata, + "geometry": str(args.geometry), + "output": str(args.output), + "resolution": resolution, + "vertices": int(vertices.shape[0]), + "faces": int(faces.shape[0]), + "texture_size": args.texture_size, + "grid_chunk_size": args.grid_chunk_size, + "source_face_chunk_size": args.source_face_chunk_size, + "projection_query_chunk_size": args.projection_query_chunk_size, + "projection_cache": ( + str(args.projection_cache) + if args.projection_cache is not None + else None + ), + "projection_seconds": projection_seconds, + "projection_distance_voxels": projection_distance_voxels, + "texture_fallback": fallback_report, + "bake_seconds": bake_seconds, + "total_seconds": total_seconds, + } + report_path = args.report or args.output.with_suffix(".export.json") + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text( + json.dumps(report, indent=2, sort_keys=True), + encoding="utf-8", + ) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_iso_quality.sh b/scripts/run_iso_quality.sh new file mode 100755 index 0000000..ffc68ec --- /dev/null +++ b/scripts/run_iso_quality.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)" +cd "$ROOT_DIR" + +INPUT="${1:-output/inputs/0_img_2048.png}" +OUTPUT="${2:-output/pixal3d_mps_1536_cuda_parity.glb}" +LOG="${OUTPUT%.glb}.log" + +if [[ ! -x .venv/bin/python ]]; then + echo "Missing .venv; run bash setup_macos.sh first." >&2 + exit 2 +fi + +if [[ ! -f "$INPUT" && "$INPUT" == "output/inputs/0_img_2048.png" ]]; then + mkdir -p "$(dirname -- "$INPUT")" + sips --resampleHeightWidth 2048 2048 \ + assets/images/0_img.png \ + --out "$INPUT" >/dev/null +fi +if [[ ! -f "$INPUT" ]]; then + echo "Input image not found: $INPUT" >&2 + exit 2 +fi + +mkdir -p "$(dirname -- "$OUTPUT")" "$(dirname -- "$LOG")" +export PYTHONUNBUFFERED=1 +export PIXAL3D_NAF_CHUNK_ROWS="${PIXAL3D_NAF_CHUNK_ROWS:-1}" +export SPARSE_ATTN_BACKEND=sdpa + +echo "Input: $INPUT" +echo "Output: $OUTPUT" +echo "Log: $LOG" + +caffeinate -dimsu .venv/bin/python inference.py \ + --image "$INPUT" \ + --output "$OUTPUT" \ + --resolution 1536 \ + --low_vram \ + --export-profile cuda-parity \ + 2>&1 | tee "$LOG" diff --git a/scripts/smoke_cuda_parity_export.py b/scripts/smoke_cuda_parity_export.py new file mode 100644 index 0000000..3b7fa4f --- /dev/null +++ b/scripts/smoke_cuda_parity_export.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Small end-to-end check of native Metal remesh, raster and texture sampling.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import torch +import trimesh + +from macos_compat import configure + +configure() + +from backends.cuda_parity_export import to_glb_cuda_parity + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--output", + default="output/diagnostics/cuda_parity_smoke.glb", + ) + args = parser.parse_args() + + resolution = 32 + source = trimesh.creation.icosphere(subdivisions=3, radius=0.38) + vertices = torch.from_numpy( + np.ascontiguousarray(source.vertices, dtype=np.float32) + ) + faces = torch.from_numpy( + np.ascontiguousarray(source.faces, dtype=np.int32) + ) + + xyz = torch.stack( + torch.meshgrid( + *[torch.arange(resolution, dtype=torch.int32)] * 3, + indexing="ij", + ), + dim=-1, + ).reshape(-1, 3) + normalized = xyz.float() / (resolution - 1) + attrs = torch.cat( + [ + normalized, + torch.zeros(len(xyz), 1), + torch.full((len(xyz), 1), 0.6), + torch.ones(len(xyz), 1), + ], + dim=1, + ) + result = to_glb_cuda_parity( + vertices=vertices, + faces=faces, + attr_volume=attrs, + coords=xyz, + attr_layout={ + "base_color": slice(0, 3), + "metallic": slice(3, 4), + "roughness": slice(4, 5), + "alpha": slice(5, 6), + }, + resolution=resolution, + decimation_target=20_000, + texture_size=128, + bvh_chunk_size=4096, + grid_chunk_size=4096, + verbose=True, + use_tqdm=True, + ) + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + result.export(output, extension_webp=True) + material = result.visual.material + if material.alphaMode != "OPAQUE" or material.doubleSided: + raise RuntimeError("CUDA material semantics were not preserved") + print( + f"PASS: {output} ({len(result.vertices):,} vertices, " + f"{len(result.faces):,} faces)" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/smoke_naf_mps.py b/scripts/smoke_naf_mps.py new file mode 100644 index 0000000..5e02376 --- /dev/null +++ b/scripts/smoke_naf_mps.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Load the real pretrained NAF module and run its chunked MPS attention.""" + +from __future__ import annotations + +import argparse +import json +import time + +import torch +import torch.nn.functional as F + +from backends.naf_attention import install_chunked_naf_attention + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--target-size", type=int, default=128) + parser.add_argument("--feature-size", type=int, default=16) + args = parser.parse_args() + + if not torch.backends.mps.is_available(): + raise RuntimeError("This smoke test requires MPS") + if args.target_size % args.feature_size: + raise ValueError("target size must be divisible by feature size") + + model = torch.hub.load( + "valeoai/NAF", + "naf", + pretrained=True, + device="cpu", + trust_repo=True, + ) + install_chunked_naf_attention(model) + model.eval().requires_grad_(False).to("mps") + + generator = torch.Generator().manual_seed(17) + image = torch.rand( + 1, + 3, + args.target_size, + args.target_size, + generator=generator, + ).to("mps") + features = torch.randn( + 1, + 1024, + args.feature_size, + args.feature_size, + generator=generator, + ).to("mps") + started = time.perf_counter() + with torch.no_grad(): + output = model( + image, + features, + (args.target_size, args.target_size), + ) + torch.mps.synchronize() + elapsed = time.perf_counter() - started + bilinear = F.interpolate( + features, + size=(args.target_size, args.target_size), + mode="bilinear", + align_corners=False, + ) + report = { + "shape": list(output.shape), + "finite": bool(torch.isfinite(output).all().item()), + "seconds": elapsed, + "mps_allocated_gib": torch.mps.current_allocated_memory() / 1024**3, + "mps_driver_gib": torch.mps.driver_allocated_memory() / 1024**3, + "learned_vs_bilinear_mean_abs": float( + (output - bilinear).abs().mean().item() + ), + "learned_vs_bilinear_max_abs": float( + (output - bilinear).abs().max().item() + ), + } + print(json.dumps(report, indent=2, sort_keys=True)) + if report["shape"] != [ + 1, + 1024, + args.target_size, + args.target_size, + ] or not report["finite"]: + raise RuntimeError("Pretrained NAF MPS smoke test failed") + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_metal_backends.py b/scripts/validate_metal_backends.py new file mode 100644 index 0000000..6074a1f --- /dev/null +++ b/scripts/validate_metal_backends.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +"""Numerically validate the Metal hot paths before full Pixal3D inference.""" + +from __future__ import annotations + +import json +import math +import time + +import torch +import torch.nn.functional as F + +from backends.naf_attention import chunked_na2d + + +def _error_metrics(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + actual_f = actual.detach().cpu().float() + expected_f = expected.detach().cpu().float() + difference = (actual_f - expected_f).abs() + cosine = F.cosine_similarity( + actual_f.reshape(1, -1), + expected_f.reshape(1, -1), + ).item() + return { + "max_abs": float(difference.max().item()), + "mean_abs": float(difference.mean().item()), + "cosine": float(cosine), + } + + +def validate_sparse_convolution( + device: torch.device, + dtype: torch.dtype = torch.float16, +) -> dict[str, object]: + from flex_gemm.ops.spconv import ( + Algorithm, + set_algorithm, + sparse_submanifold_conv3d, + ) + + generator = torch.Generator().manual_seed(7) + resolution = 12 + grid = torch.stack( + torch.meshgrid( + torch.arange(resolution), + torch.arange(resolution), + torch.arange(resolution), + indexing="ij", + ), + dim=-1, + ) + active = ((grid.float() - 5.5) ** 2).sum(dim=-1).sqrt() < 5 + coords = torch.nonzero(active).int() + coords = torch.cat( + [torch.zeros(len(coords), 1, dtype=torch.int32), coords], + dim=1, + ).contiguous().to(device) + feats = torch.randn( + len(coords), 16, generator=generator, dtype=dtype + ).to(device) + weight = torch.randn( + 24, 3, 3, 3, 16, generator=generator, dtype=dtype + ).to(device) + bias = torch.randn(24, generator=generator, dtype=dtype).to(device) + shape = torch.Size([1, 16, resolution, resolution, resolution]) + + results = {} + coords_cpu = coords.cpu() + feats_cpu = feats.cpu().float() + weight_cpu = weight.cpu().float() + bias_cpu = bias.cpu().float() + coordinate_to_index = { + tuple(coordinate.tolist()): index + for index, coordinate in enumerate(coords_cpu) + } + reference = bias_cpu[None].expand(len(coords_cpu), -1).clone() + for target_index, coordinate in enumerate(coords_cpu.tolist()): + batch, x, y, z = coordinate + for kernel_x in range(3): + for kernel_y in range(3): + for kernel_z in range(3): + source_index = coordinate_to_index.get( + ( + batch, + x + kernel_x - 1, + y + kernel_y - 1, + z + kernel_z - 1, + ) + ) + if source_index is None: + continue + reference[target_index] += ( + feats_cpu[source_index] + @ weight_cpu[ + :, + kernel_x, + kernel_y, + kernel_z, + :, + ].T + ) + for name, algorithm in ( + ("implicit", Algorithm.IMPLICIT_GEMM), + ("masked", Algorithm.MASKED_IMPLICIT_GEMM), + ("masked_splitk", Algorithm.MASKED_IMPLICIT_GEMM_SPLITK), + ): + set_algorithm(algorithm) + started = time.perf_counter() + output, _ = sparse_submanifold_conv3d( + feats, coords, shape, weight, bias + ) + torch.mps.synchronize() + results[name] = { + **_error_metrics(output, reference), + "seconds": time.perf_counter() - started, + } + return results + + +def _sdpa_packed_reference(q, k, v, sequence_lengths): + outputs = [] + offset = 0 + for length in sequence_lengths: + qi = q[offset : offset + length].permute(1, 0, 2).unsqueeze(0) + ki = k[offset : offset + length].permute(1, 0, 2).unsqueeze(0) + vi = v[offset : offset + length].permute(1, 0, 2).unsqueeze(0) + outputs.append( + F.scaled_dot_product_attention(qi, ki, vi) + .squeeze(0) + .permute(1, 0, 2) + ) + offset += length + return torch.cat(outputs, dim=0) + + +def validate_sparse_attention( + device: torch.device, + dtype: torch.dtype = torch.float16, +) -> dict[str, object]: + import flex_gemm + + sequence_lengths = [73, 41, 19] + total = sum(sequence_lengths) + generator = torch.Generator().manual_seed(11) + q = torch.randn( + total, 4, 32, generator=generator, dtype=dtype + ).to(device) + k = torch.randn( + total, 4, 32, generator=generator, dtype=dtype + ).to(device) + v = torch.randn( + total, 4, 32, generator=generator, dtype=dtype + ).to(device) + prefix = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.int32, + device=device, + ) + + reference = _sdpa_packed_reference(q, k, v, sequence_lengths) + started = time.perf_counter() + output = flex_gemm.kernels.cuda.sparse_attention_fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + prefix, + prefix, + max(sequence_lengths), + max(sequence_lengths), + 1.0 / math.sqrt(q.shape[-1]), + ) + torch.mps.synchronize() + return { + **_error_metrics(output, reference), + "seconds": time.perf_counter() - started, + } + + +def validate_naf_chunking(device: torch.device) -> dict[str, object]: + generator = torch.Generator().manual_seed(13) + tensors = [ + torch.randn(1, 36, 36, 2, 8, generator=generator) + for _ in range(3) + ] + cpu_output = chunked_na2d( + *tensors, + kernel_size=3, + dilation=4, + chunk_rows=5, + ) + mps_tensors = [tensor.to(device) for tensor in tensors] + started = time.perf_counter() + mps_output = chunked_na2d( + *mps_tensors, + kernel_size=3, + dilation=4, + chunk_rows=5, + ) + torch.mps.synchronize() + return { + **_error_metrics(mps_output, cpu_output), + "seconds": time.perf_counter() - started, + } + + +def main() -> None: + if not torch.backends.mps.is_available(): + raise RuntimeError("Metal validation requires an available MPS device") + device = torch.device("mps") + report = { + "torch": torch.__version__, + "device": str(device), + "sparse_convolution": validate_sparse_convolution(device), + "sparse_attention": validate_sparse_attention(device), + "bf16": { + "sparse_convolution": validate_sparse_convolution( + device, + torch.bfloat16, + ), + "sparse_attention": validate_sparse_attention( + device, + torch.bfloat16, + ), + }, + "naf_chunking": validate_naf_chunking(device), + } + print(json.dumps(report, indent=2, sort_keys=True)) + + convolution_max = max( + result["max_abs"] + for result in report["sparse_convolution"].values() + ) + # The deliberately unscaled random convolution sums 27 fp16 products and + # reaches values around 100; 0.04 is below 0.05% relative error here. + if convolution_max > 0.04: + raise RuntimeError( + f"Metal sparse convolution parity failed: max_abs={convolution_max}" + ) + if report["sparse_attention"]["max_abs"] > 0.03: + raise RuntimeError("Metal sparse attention parity failed") + bf16_convolution_cosine = min( + result["cosine"] + for result in report["bf16"]["sparse_convolution"].values() + ) + if bf16_convolution_cosine < 0.999: + raise RuntimeError("Metal BF16 sparse convolution parity failed") + if report["bf16"]["sparse_attention"]["cosine"] < 0.999: + raise RuntimeError("Metal BF16 sparse attention parity failed") + if report["naf_chunking"]["max_abs"] > 2e-5: + raise RuntimeError("MPS NAF chunking parity failed") + + +if __name__ == "__main__": + main() diff --git a/setup_macos.sh b/setup_macos.sh new file mode 100644 index 0000000..12b040d --- /dev/null +++ b/setup_macos.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)" +cd "$ROOT_DIR" + +if [[ "$(uname -m)" != "arm64" ]]; then + echo "Pixal3D Metal setup requires Apple Silicon (arm64)." >&2 + exit 1 +fi + +DEPS_DIR="$ROOT_DIR/deps" +mkdir -p "$DEPS_DIR" + +clone_pinned() { + local url="$1" dir="$2" commit="$3" + if [[ ! -d "$DEPS_DIR/$dir/.git" ]]; then + git clone "$url" "$DEPS_DIR/$dir" + fi + git -C "$DEPS_DIR/$dir" fetch --depth 1 origin "$commit" + git -C "$DEPS_DIR/$dir" checkout --detach "$commit" +} + +clone_pinned https://github.com/pedronaugusto/mtlbvh.git mtlbvh 23f441c470ce1f537e1fd836f3ffb5b8245f7975 +clone_pinned https://github.com/pedronaugusto/mtldiffrast.git mtldiffrast 4668cd91cb6d27f5e264731f94a06841fbf7aab8 +clone_pinned https://github.com/pedronaugusto/mtlmesh.git mtlmesh 212079e55772cff3d648a21372392c37e0643f3b +clone_pinned https://github.com/pedronaugusto/mtlgemm.git mtlgemm 867aec8234299a7fe1ede7f802c8debe5a939a82 +clone_pinned https://github.com/pedronaugusto/trellis2-apple.git trellis2-apple 6055b868734af6e12769d229d90580e775fae9f0 + +if [[ ! -d .venv ]]; then + uv venv .venv --python python3.11 +fi +PYTHON="$ROOT_DIR/.venv/bin/python" +pip_install() { + uv pip install --python "$PYTHON" "$@" +} + +pip_install "torch>=2.13,<2.14" "torchvision>=0.28,<0.29" \ + setuptools wheel pybind11 +pip_install -r requirements.txt +pip_install -r requirements-macos.txt +# MoGe currently requires the newer utils3d.pt API. The older 0.0.2 wheel +# mentioned by the Pixal3D model card does not provide that module. +pip_install git+https://github.com/EasternJournalist/utils3d.git + +export MACOSX_DEPLOYMENT_TARGET="${MACOSX_DEPLOYMENT_TARGET:-12.0}" +pip_install --no-build-isolation "$DEPS_DIR/mtlbvh" +pip_install --no-build-isolation "$DEPS_DIR/mtldiffrast" +pip_install --no-build-isolation "$DEPS_DIR/mtlmesh" +pip_install --no-build-isolation "$DEPS_DIR/mtlgemm" +pip_install --no-build-isolation "$DEPS_DIR/trellis2-apple/o-voxel" + +"$PYTHON" patches/mps_compat.py +"$PYTHON" scripts/doctor.py + +echo +echo "Pixal3D macOS setup complete." +echo "Activate: source .venv/bin/activate" +echo "Run: python inference.py --image assets/images/0_img.png --output output.glb" +echo "Low RAM: python inference.py --image assets/images/0_img.png --output output.glb --low_vram" diff --git a/stubs/o_voxel/__init__.py b/stubs/o_voxel/__init__.py new file mode 100644 index 0000000..2ae2839 --- /dev/null +++ b/stubs/o_voxel/__init__.py @@ -0,0 +1 @@ +pass diff --git a/stubs/o_voxel/convert.py b/stubs/o_voxel/convert.py new file mode 100644 index 0000000..b548dd5 --- /dev/null +++ b/stubs/o_voxel/convert.py @@ -0,0 +1,146 @@ +""" +Pure-Python/PyTorch mesh extraction from sparse voxel dual-grid. + +Replaces the CUDA-only o_voxel._C hashmap operations with Python dicts. +Produces identical output to the CUDA version for inference. +""" + +import torch +import numpy as np +from typing import Union + + +def mesh_to_flexible_dual_grid(*args, **kwargs): + raise RuntimeError("mesh_to_flexible_dual_grid requires CUDA (o_voxel)") + + +# Static lookup tables (lazily initialized, cached per-device) +_edge_neighbor_voxel_offset = None +_quad_split_1 = None +_quad_split_2 = None + + +def flexible_dual_grid_to_mesh( + coords: torch.Tensor, + dual_vertices: torch.Tensor, + intersected_flag: torch.Tensor, + split_weight: Union[torch.Tensor, None], + aabb: Union[list, tuple, np.ndarray, torch.Tensor], + voxel_size: Union[float, list, tuple, np.ndarray, torch.Tensor] = None, + grid_size: Union[int, list, tuple, np.ndarray, torch.Tensor] = None, + train: bool = False, +): + """ + Extract a triangle mesh from sparse voxel dual-grid representation. + + Given a set of voxel coordinates with dual vertex positions and edge + intersection flags, builds quads connecting adjacent voxels at intersected + edges, then splits each quad into two triangles. + + Args: + coords: [N, 3] integer voxel coordinates. + dual_vertices: [N, 3] float vertex offsets within each voxel. + intersected_flag: [N, 3] bool flags indicating which edges are intersected. + split_weight: [N, 1] optional quad split weights (None = use normal alignment). + aabb: [[min_x, min_y, min_z], [max_x, max_y, max_z]] bounding box. + voxel_size: Size of each voxel (alternative to grid_size). + grid_size: Number of voxels per axis (alternative to voxel_size). + train: Must be False (training not supported in pure-Python version). + + Returns: + (vertices, triangles): mesh vertices [V, 3] and face indices [F, 3]. + """ + global _edge_neighbor_voxel_offset, _quad_split_1, _quad_split_2 + + device = coords.device + + if _edge_neighbor_voxel_offset is None or _edge_neighbor_voxel_offset.device != device: + _edge_neighbor_voxel_offset = torch.tensor([ + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], + [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], + [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]], + ], dtype=torch.int, device=device).unsqueeze(0) + _quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device) + _quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device) + + if isinstance(aabb, (list, tuple)): + aabb = np.array(aabb) + if isinstance(aabb, np.ndarray): + aabb = torch.tensor(aabb, dtype=torch.float32, device=device) + + if voxel_size is not None: + if isinstance(voxel_size, (int, float)): + voxel_size = [voxel_size] * 3 + if isinstance(voxel_size, (list, tuple, np.ndarray)): + voxel_size = torch.tensor(np.array(voxel_size), dtype=torch.float32, device=device) + grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int() + else: + if isinstance(grid_size, int): + grid_size = [grid_size] * 3 + if isinstance(grid_size, (list, tuple, np.ndarray)): + grid_size = torch.tensor(np.array(grid_size), dtype=torch.int32, device=device) + voxel_size = (aabb[1] - aabb[0]) / grid_size.float() + + N = dual_vertices.shape[0] + + # Build coordinate lookup on CPU + coords_cpu = coords.cpu() + coord_to_idx = {} + for i in range(N): + key = (coords_cpu[i, 0].item(), coords_cpu[i, 1].item(), coords_cpu[i, 2].item()) + coord_to_idx[key] = i + + # Find connected voxels for each intersected edge + edge_neighbor_voxel = coords.reshape(N, 1, 1, 3) + _edge_neighbor_voxel_offset + connected_voxel = edge_neighbor_voxel[intersected_flag] + M = connected_voxel.shape[0] + + if M == 0: + return torch.zeros(0, 3, device=device), torch.zeros(0, 3, dtype=torch.long, device=device) + + # Look up neighbor indices via dict + connected_cpu = connected_voxel.cpu().reshape(-1, 3) + indices = [] + for j in range(connected_cpu.shape[0]): + key = (connected_cpu[j, 0].item(), connected_cpu[j, 1].item(), connected_cpu[j, 2].item()) + indices.append(coord_to_idx.get(key, 0xFFFFFFFF)) + + connected_voxel_indices = torch.tensor(indices, dtype=torch.int64, device=device).reshape(M, 4) + connected_voxel_valid = (connected_voxel_indices != 0xFFFFFFFF).all(dim=1) + quad_indices = connected_voxel_indices[connected_voxel_valid].long() + L = quad_indices.shape[0] + + if L == 0: + return torch.zeros(0, 3, device=device), torch.zeros(0, 3, dtype=torch.long, device=device) + + # Compute world-space vertex positions + mesh_vertices = (coords.float() + dual_vertices) * voxel_size + aabb[0].reshape(1, 3) + + if train: + raise RuntimeError("Training mode not supported in pure-Python mesh extraction") + + # Triangulate quads: choose the diagonal split that produces better-aligned normals + if split_weight is None: + a1 = quad_indices[:, _quad_split_1] + n0 = torch.cross(mesh_vertices[a1[:, 1]] - mesh_vertices[a1[:, 0]], mesh_vertices[a1[:, 2]] - mesh_vertices[a1[:, 0]]) + n1 = torch.cross(mesh_vertices[a1[:, 2]] - mesh_vertices[a1[:, 1]], mesh_vertices[a1[:, 3]] - mesh_vertices[a1[:, 1]]) + align0 = (n0 * n1).sum(dim=1, keepdim=True).abs() + + a2 = quad_indices[:, _quad_split_2] + n0 = torch.cross(mesh_vertices[a2[:, 1]] - mesh_vertices[a2[:, 0]], mesh_vertices[a2[:, 2]] - mesh_vertices[a2[:, 0]]) + n1 = torch.cross(mesh_vertices[a2[:, 2]] - mesh_vertices[a2[:, 1]], mesh_vertices[a2[:, 3]] - mesh_vertices[a2[:, 1]]) + align1 = (n0 * n1).sum(dim=1, keepdim=True).abs() + + mesh_triangles = torch.where(align0 > align1, a1, a2).reshape(-1, 3) + else: + sw = split_weight[quad_indices] + sw_02 = (sw[:, 0] * sw[:, 2]).squeeze() + sw_13 = (sw[:, 1] * sw[:, 3]).squeeze() + cond = (sw_02 > sw_13).unsqueeze(1).expand(-1, 6) + mesh_triangles = torch.where( + cond, + quad_indices[:, _quad_split_1], + quad_indices[:, _quad_split_2], + ).reshape(-1, 3) + + return mesh_vertices, mesh_triangles diff --git a/stubs/o_voxel/io.py b/stubs/o_voxel/io.py new file mode 100644 index 0000000..9801f92 --- /dev/null +++ b/stubs/o_voxel/io.py @@ -0,0 +1,8 @@ +def read(*args, **kwargs): + raise RuntimeError("o_voxel.io requires CUDA") + +def write(*args, **kwargs): + raise RuntimeError("o_voxel.io requires CUDA") + +def read_vxz(*args, **kwargs): + raise RuntimeError("o_voxel.io requires CUDA") diff --git a/stubs/o_voxel/rasterize.py b/stubs/o_voxel/rasterize.py new file mode 100644 index 0000000..5e9fef5 --- /dev/null +++ b/stubs/o_voxel/rasterize.py @@ -0,0 +1,3 @@ +class VoxelRenderer: + def __init__(self, *args, **kwargs): + raise RuntimeError("o_voxel.rasterize requires CUDA") diff --git a/stubs/o_voxel_override_convert.py b/stubs/o_voxel_override_convert.py new file mode 100644 index 0000000..b548dd5 --- /dev/null +++ b/stubs/o_voxel_override_convert.py @@ -0,0 +1,146 @@ +""" +Pure-Python/PyTorch mesh extraction from sparse voxel dual-grid. + +Replaces the CUDA-only o_voxel._C hashmap operations with Python dicts. +Produces identical output to the CUDA version for inference. +""" + +import torch +import numpy as np +from typing import Union + + +def mesh_to_flexible_dual_grid(*args, **kwargs): + raise RuntimeError("mesh_to_flexible_dual_grid requires CUDA (o_voxel)") + + +# Static lookup tables (lazily initialized, cached per-device) +_edge_neighbor_voxel_offset = None +_quad_split_1 = None +_quad_split_2 = None + + +def flexible_dual_grid_to_mesh( + coords: torch.Tensor, + dual_vertices: torch.Tensor, + intersected_flag: torch.Tensor, + split_weight: Union[torch.Tensor, None], + aabb: Union[list, tuple, np.ndarray, torch.Tensor], + voxel_size: Union[float, list, tuple, np.ndarray, torch.Tensor] = None, + grid_size: Union[int, list, tuple, np.ndarray, torch.Tensor] = None, + train: bool = False, +): + """ + Extract a triangle mesh from sparse voxel dual-grid representation. + + Given a set of voxel coordinates with dual vertex positions and edge + intersection flags, builds quads connecting adjacent voxels at intersected + edges, then splits each quad into two triangles. + + Args: + coords: [N, 3] integer voxel coordinates. + dual_vertices: [N, 3] float vertex offsets within each voxel. + intersected_flag: [N, 3] bool flags indicating which edges are intersected. + split_weight: [N, 1] optional quad split weights (None = use normal alignment). + aabb: [[min_x, min_y, min_z], [max_x, max_y, max_z]] bounding box. + voxel_size: Size of each voxel (alternative to grid_size). + grid_size: Number of voxels per axis (alternative to voxel_size). + train: Must be False (training not supported in pure-Python version). + + Returns: + (vertices, triangles): mesh vertices [V, 3] and face indices [F, 3]. + """ + global _edge_neighbor_voxel_offset, _quad_split_1, _quad_split_2 + + device = coords.device + + if _edge_neighbor_voxel_offset is None or _edge_neighbor_voxel_offset.device != device: + _edge_neighbor_voxel_offset = torch.tensor([ + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], + [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], + [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]], + ], dtype=torch.int, device=device).unsqueeze(0) + _quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device) + _quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device) + + if isinstance(aabb, (list, tuple)): + aabb = np.array(aabb) + if isinstance(aabb, np.ndarray): + aabb = torch.tensor(aabb, dtype=torch.float32, device=device) + + if voxel_size is not None: + if isinstance(voxel_size, (int, float)): + voxel_size = [voxel_size] * 3 + if isinstance(voxel_size, (list, tuple, np.ndarray)): + voxel_size = torch.tensor(np.array(voxel_size), dtype=torch.float32, device=device) + grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int() + else: + if isinstance(grid_size, int): + grid_size = [grid_size] * 3 + if isinstance(grid_size, (list, tuple, np.ndarray)): + grid_size = torch.tensor(np.array(grid_size), dtype=torch.int32, device=device) + voxel_size = (aabb[1] - aabb[0]) / grid_size.float() + + N = dual_vertices.shape[0] + + # Build coordinate lookup on CPU + coords_cpu = coords.cpu() + coord_to_idx = {} + for i in range(N): + key = (coords_cpu[i, 0].item(), coords_cpu[i, 1].item(), coords_cpu[i, 2].item()) + coord_to_idx[key] = i + + # Find connected voxels for each intersected edge + edge_neighbor_voxel = coords.reshape(N, 1, 1, 3) + _edge_neighbor_voxel_offset + connected_voxel = edge_neighbor_voxel[intersected_flag] + M = connected_voxel.shape[0] + + if M == 0: + return torch.zeros(0, 3, device=device), torch.zeros(0, 3, dtype=torch.long, device=device) + + # Look up neighbor indices via dict + connected_cpu = connected_voxel.cpu().reshape(-1, 3) + indices = [] + for j in range(connected_cpu.shape[0]): + key = (connected_cpu[j, 0].item(), connected_cpu[j, 1].item(), connected_cpu[j, 2].item()) + indices.append(coord_to_idx.get(key, 0xFFFFFFFF)) + + connected_voxel_indices = torch.tensor(indices, dtype=torch.int64, device=device).reshape(M, 4) + connected_voxel_valid = (connected_voxel_indices != 0xFFFFFFFF).all(dim=1) + quad_indices = connected_voxel_indices[connected_voxel_valid].long() + L = quad_indices.shape[0] + + if L == 0: + return torch.zeros(0, 3, device=device), torch.zeros(0, 3, dtype=torch.long, device=device) + + # Compute world-space vertex positions + mesh_vertices = (coords.float() + dual_vertices) * voxel_size + aabb[0].reshape(1, 3) + + if train: + raise RuntimeError("Training mode not supported in pure-Python mesh extraction") + + # Triangulate quads: choose the diagonal split that produces better-aligned normals + if split_weight is None: + a1 = quad_indices[:, _quad_split_1] + n0 = torch.cross(mesh_vertices[a1[:, 1]] - mesh_vertices[a1[:, 0]], mesh_vertices[a1[:, 2]] - mesh_vertices[a1[:, 0]]) + n1 = torch.cross(mesh_vertices[a1[:, 2]] - mesh_vertices[a1[:, 1]], mesh_vertices[a1[:, 3]] - mesh_vertices[a1[:, 1]]) + align0 = (n0 * n1).sum(dim=1, keepdim=True).abs() + + a2 = quad_indices[:, _quad_split_2] + n0 = torch.cross(mesh_vertices[a2[:, 1]] - mesh_vertices[a2[:, 0]], mesh_vertices[a2[:, 2]] - mesh_vertices[a2[:, 0]]) + n1 = torch.cross(mesh_vertices[a2[:, 2]] - mesh_vertices[a2[:, 1]], mesh_vertices[a2[:, 3]] - mesh_vertices[a2[:, 1]]) + align1 = (n0 * n1).sum(dim=1, keepdim=True).abs() + + mesh_triangles = torch.where(align0 > align1, a1, a2).reshape(-1, 3) + else: + sw = split_weight[quad_indices] + sw_02 = (sw[:, 0] * sw[:, 2]).squeeze() + sw_13 = (sw[:, 1] * sw[:, 3]).squeeze() + cond = (sw_02 > sw_13).unsqueeze(1).expand(-1, 6) + mesh_triangles = torch.where( + cond, + quad_indices[:, _quad_split_1], + quad_indices[:, _quad_split_2], + ).reshape(-1, 3) + + return mesh_vertices, mesh_triangles diff --git a/tests/test_cuda_parity_export.py b/tests/test_cuda_parity_export.py new file mode 100644 index 0000000..10e940b --- /dev/null +++ b/tests/test_cuda_parity_export.py @@ -0,0 +1,128 @@ +from types import SimpleNamespace + +import torch + +from backends.cuda_parity_export import ( + _project_to_source, + force_cuda_material_semantics, + memory_bounded_o_voxel, +) + + +class _FakeBVH: + def __init__(self, vertices, faces): + self.vertices = vertices + self.faces = faces + + def unsigned_distance(self, positions, return_uvw=False): + distances = positions[:, 0] + face_ids = torch.arange(len(positions), dtype=torch.int64) + uvw = torch.ones(len(positions), 3) if return_uvw else None + return distances, face_ids, uvw + + +def test_memory_bounded_queries_preserve_order() -> None: + calls = [] + + def grid_sample(feats, coords, shape, grid, mode="trilinear"): + del feats, coords, shape, mode + calls.append(grid.shape[1]) + return grid.sum(dim=-1, keepdim=True) + + def remesh(*args, **kwargs): + del args + return kwargs + + module = SimpleNamespace( + _BVH=_FakeBVH, + _grid_sample_3d=grid_sample, + _remesh_narrow_band_dc=remesh, + ) + positions = torch.arange(30, dtype=torch.float32).reshape(10, 3) + grid = positions.reshape(1, 10, 3) + + with memory_bounded_o_voxel( + module, + bvh_chunk_size=4, + grid_chunk_size=3, + source_resolution=1536, + ): + bvh = module._BVH(torch.empty(0), torch.empty(0)) + distances, face_ids, uvw = bvh.unsigned_distance( + positions, + return_uvw=True, + ) + sampled = module._grid_sample_3d( + torch.empty(0), + torch.empty(0), + torch.Size(), + grid, + ) + remesh_kwargs = module._remesh_narrow_band_dc( + bvh=bvh, + scale=1.0, + resolution=1536, + band=1, + ) + + assert calls == [3, 3, 3, 1] + assert torch.equal(distances, positions[:, 0]) + assert torch.equal(face_ids, torch.tensor([0, 1, 2, 3, 0, 1, 2, 3, 0, 1])) + assert uvw.shape == (10, 3) + assert sampled.shape == (10, 1) + assert isinstance(remesh_kwargs["bvh"], _FakeBVH) + assert module._BVH is _FakeBVH + + +def test_force_cuda_material_semantics() -> None: + material = SimpleNamespace(alphaMode="BLEND", doubleSided=True) + mesh = SimpleNamespace(visual=SimpleNamespace(material=material)) + force_cuda_material_semantics(mesh) + assert material.alphaMode == "OPAQUE" + assert material.doubleSided is False + + +def test_project_to_source_uses_face_ids_and_barycentrics() -> None: + source_vertices = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + source_faces = torch.tensor( + [[0, 1, 2], [0, 2, 3]], + dtype=torch.int32, + ) + + class ProjectionBVH: + def unsigned_distance(self, positions, return_uvw=False): + assert return_uvw + assert len(positions) == 2 + return ( + torch.zeros(2), + torch.tensor([0, 1]), + torch.tensor( + [ + [0.25, 0.50, 0.25], + [0.50, 0.25, 0.25], + ] + ), + ) + + projected = _project_to_source( + ProjectionBVH(), + source_vertices, + source_faces, + torch.zeros(2, 3), + ) + assert torch.allclose( + projected, + torch.tensor( + [ + [0.50, 0.25, 0.0], + [0.0, 0.25, 0.25], + ] + ), + ) diff --git a/tests/test_decoded_checkpoint.py b/tests/test_decoded_checkpoint.py new file mode 100644 index 0000000..4891ff6 --- /dev/null +++ b/tests/test_decoded_checkpoint.py @@ -0,0 +1,46 @@ +from pathlib import Path + +import torch + +from backends.decoded_checkpoint import ( + load_decoded_checkpoint, + save_decoded_checkpoint, +) +from pixal3d.representations import MeshWithVoxel + + +def test_decoded_checkpoint_round_trip(tmp_path: Path) -> None: + mesh = MeshWithVoxel( + vertices=torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]] + ), + faces=torch.tensor([[0, 1, 2]], dtype=torch.int32), + origin=[-0.5, -0.5, -0.5], + voxel_size=1 / 16, + coords=torch.tensor([[1, 2, 3], [2, 3, 4]], dtype=torch.int32), + attrs=torch.arange(12, dtype=torch.float16).reshape(2, 6), + voxel_shape=torch.Size([1, 6, 16, 16, 16]), + layout={ + "base_color": slice(0, 3), + "metallic": slice(3, 4), + "roughness": slice(4, 5), + "alpha": slice(5, 6), + }, + ) + path = tmp_path / "decoded.pt" + save_decoded_checkpoint( + path, + mesh, + resolution=16, + metadata={"seed": 42}, + ) + + restored, resolution, metadata = load_decoded_checkpoint(path) + assert resolution == 16 + assert metadata == {"seed": 42} + assert restored.layout == mesh.layout + assert restored.voxel_shape == mesh.voxel_shape + assert torch.equal(restored.vertices, mesh.vertices) + assert torch.equal(restored.faces, mesh.faces) + assert torch.equal(restored.coords, mesh.coords) + assert torch.equal(restored.attrs, mesh.attrs) diff --git a/tests/test_naf_attention.py b/tests/test_naf_attention.py new file mode 100644 index 0000000..30ac3c9 --- /dev/null +++ b/tests/test_naf_attention.py @@ -0,0 +1,89 @@ +import torch + +from backends.naf_attention import chunked_na2d, neighborhood_indices + + +def _brute_na2d(query, key, value, kernel_size, dilation, scale): + batch, height, width, heads, head_dim = query.shape + y_indices = neighborhood_indices( + height, kernel_size, dilation, device=query.device + ) + x_indices = neighborhood_indices( + width, kernel_size, dilation, device=query.device + ) + output = torch.empty( + batch, + height, + width, + heads, + value.shape[-1], + dtype=value.dtype, + device=value.device, + ) + for y in range(height): + for x in range(width): + keys = [] + values = [] + for ny in y_indices[y]: + for nx in x_indices[x]: + keys.append(key[:, ny, nx]) + values.append(value[:, ny, nx]) + keys = torch.stack(keys, dim=-2) + values = torch.stack(values, dim=-2) + logits = torch.einsum("bhd,bhkd->bhk", query[:, y, x], keys) + weights = torch.softmax(logits * scale, dim=-1) + output[:, y, x] = torch.einsum( + "bhk,bhkd->bhd", weights, values + ) + return output + + +def test_neighborhood_indices_shift_windows_at_boundaries(): + indices = neighborhood_indices( + length=8, + kernel_size=3, + dilation=1, + device=torch.device("cpu"), + ) + assert indices.tolist() == [ + [0, 1, 2], + [0, 1, 2], + [1, 2, 3], + [2, 3, 4], + [3, 4, 5], + [4, 5, 6], + [5, 6, 7], + [5, 6, 7], + ] + + +def test_neighborhood_indices_preserve_dilation_groups(): + indices = neighborhood_indices( + length=12, + kernel_size=3, + dilation=2, + device=torch.device("cpu"), + ) + for query, neighbors in enumerate(indices.tolist()): + assert all(index % 2 == query % 2 for index in neighbors) + + +def test_chunked_na2d_matches_brute_reference(): + generator = torch.Generator().manual_seed(123) + query = torch.randn(1, 12, 10, 2, 4, generator=generator) + key = torch.randn(1, 12, 10, 2, 4, generator=generator) + value = torch.randn(1, 12, 10, 2, 7, generator=generator) + scale = 4**-0.5 + expected = _brute_na2d(query, key, value, 3, 2, scale) + + for chunk_rows in (1, 3, 12): + actual = chunked_na2d( + query, + key, + value, + kernel_size=3, + dilation=2, + scale=scale, + chunk_rows=chunk_rows, + ) + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6)