inline env defaults in api_server, remove start script, drop intermediate mx.eval syncs
This commit is contained in:
parent
9cd27537d9
commit
d15e6de829
@ -8,6 +8,17 @@ Usage:
|
|||||||
python api_server.py --weights weights/TRELLIS.2-4B --port 8082
|
python api_server.py --weights weights/TRELLIS.2-4B --port 8082
|
||||||
"""
|
"""
|
||||||
import os
|
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 io
|
||||||
import base64
|
import base64
|
||||||
import time
|
import time
|
||||||
|
|||||||
@ -172,8 +172,6 @@ class MlxSparseStructureFlowModel(nn.Module):
|
|||||||
rope_cache = (self._rope_cos, self._rope_sin)
|
rope_cache = (self._rope_cos, self._rope_sin)
|
||||||
for i, block in enumerate(self.blocks):
|
for i, block in enumerate(self.blocks):
|
||||||
h = block(h, t_emb, cond, rope_cache=rope_cache)
|
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:
|
else:
|
||||||
if self._compiled_blocks is None:
|
if self._compiled_blocks is None:
|
||||||
try:
|
try:
|
||||||
@ -185,10 +183,8 @@ class MlxSparseStructureFlowModel(nn.Module):
|
|||||||
if self._compiled_blocks:
|
if self._compiled_blocks:
|
||||||
h = self._compiled_blocks(h, t_emb, cond)
|
h = self._compiled_blocks(h, t_emb, cond)
|
||||||
else:
|
else:
|
||||||
for i, block in enumerate(self.blocks):
|
for block in self.blocks:
|
||||||
h = block(h, t_emb, cond, rope_cache=None)
|
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())
|
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)
|
h = self._compiled_blocks(h, t_emb, cond, rope_cos, rope_sin)
|
||||||
else:
|
else:
|
||||||
rope_cache = (rope_cos, rope_sin)
|
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)
|
h = block(h, t_emb, cond, rope_cache=rope_cache)
|
||||||
if (i + 1) % 10 == 0:
|
|
||||||
mx.eval(h)
|
|
||||||
else:
|
else:
|
||||||
if self._compiled_blocks is None:
|
if self._compiled_blocks is None:
|
||||||
try:
|
try:
|
||||||
@ -335,10 +329,8 @@ class MlxSLatFlowModel(nn.Module):
|
|||||||
if self._compiled_blocks:
|
if self._compiled_blocks:
|
||||||
h = self._compiled_blocks(h, t_emb, cond)
|
h = self._compiled_blocks(h, t_emb, cond)
|
||||||
else:
|
else:
|
||||||
for i, block in enumerate(self.blocks):
|
for block in self.blocks:
|
||||||
h = block(h, t_emb, cond, rope_cache=None)
|
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())
|
logger.debug("[MLX] SLatFlow blocks done, N=%d, mem=%s", N, _metal_mem_mb())
|
||||||
|
|
||||||
|
|||||||
@ -396,7 +396,7 @@ def to_glb(
|
|||||||
metallicFactor=1.0,
|
metallicFactor=1.0,
|
||||||
roughnessFactor=1.0,
|
roughnessFactor=1.0,
|
||||||
alphaMode=alpha_mode,
|
alphaMode=alpha_mode,
|
||||||
doubleSided=True if not remesh else False,
|
doubleSided=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
# --- Coordinate System Conversion & Final Object ---
|
# --- Coordinate System Conversion & Final Object ---
|
||||||
@ -405,9 +405,9 @@ def to_glb(
|
|||||||
uvs_np = out_uvs.cpu().numpy()
|
uvs_np = out_uvs.cpu().numpy()
|
||||||
normals_np = out_normals.cpu().numpy()
|
normals_np = out_normals.cpu().numpy()
|
||||||
|
|
||||||
# Y-up to Z-up for GLB
|
# Y-up to Z-up for GLB (must copy to avoid in-place corruption)
|
||||||
vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2], -vertices_np[:, 1]
|
vertices_np[:, 1], vertices_np[:, 2] = vertices_np[:, 2].copy(), -vertices_np[:, 1].copy()
|
||||||
normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2], -normals_np[:, 1]
|
normals_np[:, 1], normals_np[:, 2] = normals_np[:, 2].copy(), -normals_np[:, 1].copy()
|
||||||
uvs_np[:, 1] = 1 - uvs_np[:, 1]
|
uvs_np[:, 1] = 1 - uvs_np[:, 1]
|
||||||
|
|
||||||
textured_mesh = trimesh.Trimesh(
|
textured_mesh = trimesh.Trimesh(
|
||||||
|
|||||||
@ -391,7 +391,7 @@ def to_glb(
|
|||||||
metallicFactor=1.0,
|
metallicFactor=1.0,
|
||||||
roughnessFactor=1.0,
|
roughnessFactor=1.0,
|
||||||
alphaMode=alpha_mode,
|
alphaMode=alpha_mode,
|
||||||
doubleSided=True if not remesh else False,
|
doubleSided=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Coordinate system conversion (Y-up to Z-up for GLB)
|
# Coordinate system conversion (Y-up to Z-up for GLB)
|
||||||
|
|||||||
@ -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"
|
|
||||||
Loading…
Reference in New Issue
Block a user