monster/trellis-2-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

inline env defaults in api_server, remove start script, drop intermediate mx.eval syncs

Browse Source
This commit is contained in:
Pedro Augusto 2026-03-22 14:11:49 +00:00 committed by Jourloy
parent 9cd27537d9
commit d15e6de829
5 changed files with 19 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
api_server.py
View File

@ -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

14
mlx_backend/flow_models.py
View File

@ -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())

8
o-voxel/o_voxel/postprocess.py
View File

@ -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(

2
o-voxel/o_voxel/postprocess_cpu.py
View File

@ -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)

23
scripts/start_server.sh
View File

@ -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"