From d15e6de829e34bd933aea1c50d5816fa1c826e7c Mon Sep 17 00:00:00 2001
From: Pedro Augusto <16762743+pedronaugusto@users.noreply.github.com>
Date: Sun, 22 Mar 2026 14:11:49 +0000
Subject: [PATCH] inline env defaults in api_server, remove start script, drop
 intermediate mx.eval syncs

---
 api_server.py                      | 11 +++++++++++
 mlx_backend/flow_models.py         | 14 +++-----------
 o-voxel/o_voxel/postprocess.py     |  8 ++++----
 o-voxel/o_voxel/postprocess_cpu.py |  2 +-
 scripts/start_server.sh            | 23 -----------------------
 5 files changed, 19 insertions(+), 39 deletions(-)
 delete mode 100755 scripts/start_server.sh

diff --git a/api_server.py b/api_server.py
index 0849939..2bb923a 100644
--- a/api_server.py
+++ b/api_server.py
@@ -8,6 +8,17 @@ Usage:
     python api_server.py --weights weights/TRELLIS.2-4B --port 8082
 """
 import os
+import sys
+
+# Environment defaults — set before any torch/trellis imports
+os.environ.setdefault("SPARSE_CONV_BACKEND", "flex_gemm")
+os.environ.setdefault("ATTN_BACKEND", "sdpa")
+os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")  # deform_conv2d in RMBG
+
+# Add o-voxel to path if present
+_ovoxel = os.path.join(os.path.dirname(os.path.abspath(__file__)), "o-voxel")
+if os.path.isdir(_ovoxel) and _ovoxel not in sys.path:
+    sys.path.insert(0, _ovoxel)
 import io
 import base64
 import time
diff --git a/mlx_backend/flow_models.py b/mlx_backend/flow_models.py
index c9da0de..ba24024 100644
--- a/mlx_backend/flow_models.py
+++ b/mlx_backend/flow_models.py
@@ -172,8 +172,6 @@ class MlxSparseStructureFlowModel(nn.Module):
                 rope_cache = (self._rope_cos, self._rope_sin)
                 for i, block in enumerate(self.blocks):
                     h = block(h, t_emb, cond, rope_cache=rope_cache)
-                    if (i + 1) % 10 == 0:
-                        mx.eval(h)  # periodic eval to bound memory in fallback path
         else:
             if self._compiled_blocks is None:
                 try:
@@ -185,10 +183,8 @@ class MlxSparseStructureFlowModel(nn.Module):
             if self._compiled_blocks:
                 h = self._compiled_blocks(h, t_emb, cond)
             else:
-                for i, block in enumerate(self.blocks):
+                for block in self.blocks:
                     h = block(h, t_emb, cond, rope_cache=None)
-                    if (i + 1) % 10 == 0:
-                        mx.eval(h)
 
         logger.debug("[MLX] StructureFlow blocks done, mem=%s", _metal_mem_mb())
 
@@ -320,10 +316,8 @@ class MlxSLatFlowModel(nn.Module):
                 h = self._compiled_blocks(h, t_emb, cond, rope_cos, rope_sin)
             else:
                 rope_cache = (rope_cos, rope_sin)
-                for i, block in enumerate(self.blocks):
+                for block in self.blocks:
                     h = block(h, t_emb, cond, rope_cache=rope_cache)
-                    if (i + 1) % 10 == 0:
-                        mx.eval(h)
         else:
             if self._compiled_blocks is None:
                 try:
@@ -335,10 +329,8 @@ class MlxSLatFlowModel(nn.Module):
             if self._compiled_blocks:
                 h = self._compiled_blocks(h, t_emb, cond)
             else:
-                for i, block in enumerate(self.blocks):
+                for block in self.blocks:
                     h = block(h, t_emb, cond, rope_cache=None)
-                    if (i + 1) % 10 == 0:
-                        mx.eval(h)
 
         logger.debug("[MLX] SLatFlow blocks done, N=%d, mem=%s", N, _metal_mem_mb())
 
diff --git a/o-voxel/o_voxel/postprocess.py b/o-voxel/o_voxel/postprocess.py
index c9fa5b0..1f16265 100644
--- a/o-voxel/o_voxel/postprocess.py
+++ b/o-voxel/o_voxel/postprocess.py
@@ -396,7 +396,7 @@ def to_glb(
         metallicFactor=1.0,
         roughnessFactor=1.0,
         alphaMode=alpha_mode,
-        doubleSided=True if not remesh else False,
+        doubleSided=True,
     )
     
     # --- Coordinate System Conversion & Final Object ---
@@ -405,9 +405,9 @@ def to_glb(
     uvs_np = out_uvs.cpu().numpy()
     normals_np = out_normals.cpu().numpy()
     
-    # Y-up to Z-up for GLB
-    vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2], -vertices_np[:, 1]
-    normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2], -normals_np[:, 1]
+    # Y-up to Z-up for GLB (must copy to avoid in-place corruption)
+    vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2].copy(), -vertices_np[:, 1].copy()
+    normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2].copy(), -normals_np[:, 1].copy()
     uvs_np[:, 1] = 1 - uvs_np[:, 1]
     
     textured_mesh = trimesh.Trimesh(
diff --git a/o-voxel/o_voxel/postprocess_cpu.py b/o-voxel/o_voxel/postprocess_cpu.py
index e3c6e9c..de4f22d 100644
--- a/o-voxel/o_voxel/postprocess_cpu.py
+++ b/o-voxel/o_voxel/postprocess_cpu.py
@@ -391,7 +391,7 @@ def to_glb(
         metallicFactor=1.0,
         roughnessFactor=1.0,
         alphaMode=alpha_mode,
-        doubleSided=True if not remesh else False,
+        doubleSided=True,
     )
 
     # Coordinate system conversion (Y-up to Z-up for GLB)
diff --git a/scripts/start_server.sh b/scripts/start_server.sh
deleted file mode 100755
index f16f17d..0000000
--- a/scripts/start_server.sh
+++ /dev/null
@@ -1,23 +0,0 @@
-#!/bin/bash
-# Start the Trellis2 MLX API server on macOS
-set -e
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-ROOT_DIR="$(dirname "$SCRIPT_DIR")"
-cd "$ROOT_DIR"
-
-# Set macOS-specific environment
-export SPARSE_CONV_BACKEND=pytorch
-export ATTN_BACKEND=sdpa
-export PYTORCH_ENABLE_MPS_FALLBACK=1  # deform_conv2d in RMBG not yet on MPS
-export PYTHONPATH="$ROOT_DIR/o-voxel:${PYTHONPATH:-}"
-
-# Defaults
-WEIGHTS="${TRELLIS2_WEIGHTS:-weights/TRELLIS.2-4B}"
-PORT="${TRELLIS2_PORT:-8082}"
-
-echo "Starting Trellis2 MLX API server..."
-echo "  Weights: $WEIGHTS"
-echo "  Port: $PORT"
-
-python api_server.py --weights "$WEIGHTS" --port "$PORT"