Add macOS Apple Silicon/MPS support with CUDA-parity export

.gitignore

@@ -24,3 +24,12 @@ env/
 Thumbs.db
 
 tmp
+__pycache__/
+.venv/
+deps/
+autotune_cache.json
+output/
+*.glb
+tmp/
+runpod/*.tar.gz
+RUNPOD_SESSION.md

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.

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)
+    app.launch(show_error=True, share=False)