Compare commits
5 Commits
pr175-head
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ba4e2ad42b | ||
|
|
0db816b55d | ||
|
|
7860148eb2 | ||
|
|
fa972aa345 | ||
|
|
746e727da2 |
159
mlx_backend/mlx_samplers.py
Normal file
159
mlx_backend/mlx_samplers.py
Normal file
@ -0,0 +1,159 @@
|
|||||||
|
"""Pure-MLX FlowEuler + CFG(+interval) sampler for the mlx-experimental backend.
|
||||||
|
|
||||||
|
Replaces the torch CPU sampler loop and its per-call adapter host bounces:
|
||||||
|
- tensors convert ONCE at stage entry/exit (was: every forward, 48+/stage)
|
||||||
|
- dense CFG pairs run as ONE batched forward (was: two synced forwards)
|
||||||
|
- `concat_cond` converts once per stage (was: full SparseTensor rebuild/step)
|
||||||
|
- one mx.eval per STEP (was: per forward), keeping the graph on-device
|
||||||
|
- TRELLIS2_MLX_COMPILE=1 additionally mx.compile's the step forward
|
||||||
|
(fixed shapes per stage; corridorkey/Hunyuan-T1 pattern)
|
||||||
|
|
||||||
|
Numerics mirror trellis2/pipelines/samplers/{flow_euler.py,
|
||||||
|
classifier_free_guidance_mixin.py, guidance_interval_mixin.py} exactly
|
||||||
|
(inclusive interval on t; torch-default unbiased std in CFG rescale).
|
||||||
|
|
||||||
|
Rollback: TRELLIS2_MLX_SAMPLER=0 keeps the upstream torch samplers.
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
from typing import Any, Optional
|
||||||
|
|
||||||
|
import mlx.core as mx
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from easydict import EasyDict as edict
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
from .adapters import _mlx_to_sparse, _mx_to_torch, _sparse_to_mlx, _torch_to_mx
|
||||||
|
from .sparse_tensor import MlxSparseTensor
|
||||||
|
|
||||||
|
|
||||||
|
class MlxFlowEulerGuidanceIntervalSampler:
|
||||||
|
"""MLX drop-in for FlowEulerGuidanceIntervalSampler (same .sample API)."""
|
||||||
|
|
||||||
|
def __init__(self, sigma_min: float):
|
||||||
|
self.sigma_min = sigma_min
|
||||||
|
|
||||||
|
# ---- math mirrors (flow_euler.py) ----
|
||||||
|
def _pred_to_xstart(self, x_t, t, pred):
|
||||||
|
return (1 - self.sigma_min) * x_t - \
|
||||||
|
(self.sigma_min + (1 - self.sigma_min) * t) * pred
|
||||||
|
|
||||||
|
def _xstart_to_pred(self, x_t, t, x_0):
|
||||||
|
return ((1 - self.sigma_min) * x_t - x_0) / \
|
||||||
|
(self.sigma_min + (1 - self.sigma_min) * t)
|
||||||
|
|
||||||
|
def _cfg_mix(self, x_t, t, v_pos, v_neg, g, g_rescale, feat_axes):
|
||||||
|
pred = g * v_pos + (1 - g) * v_neg
|
||||||
|
if g_rescale > 0:
|
||||||
|
x_0_pos = self._pred_to_xstart(x_t, t, v_pos)
|
||||||
|
x_0_cfg = self._pred_to_xstart(x_t, t, pred)
|
||||||
|
std_pos = mx.sqrt(mx.var(x_0_pos, axis=feat_axes,
|
||||||
|
keepdims=True, ddof=1))
|
||||||
|
std_cfg = mx.sqrt(mx.var(x_0_cfg, axis=feat_axes,
|
||||||
|
keepdims=True, ddof=1))
|
||||||
|
x_0_res = x_0_cfg * (std_pos / std_cfg)
|
||||||
|
x_0 = g_rescale * x_0_res + (1 - g_rescale) * x_0_cfg
|
||||||
|
pred = self._xstart_to_pred(x_t, t, x_0)
|
||||||
|
return pred
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _t_pairs(steps, rescale_t):
|
||||||
|
t_seq = np.linspace(1, 0, steps + 1)
|
||||||
|
t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
|
||||||
|
t_seq = t_seq.tolist()
|
||||||
|
return [(t_seq[i], t_seq[i + 1]) for i in range(steps)]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _maybe_compile(fn):
|
||||||
|
if os.environ.get("TRELLIS2_MLX_COMPILE", "0") == "1":
|
||||||
|
return mx.compile(fn)
|
||||||
|
return fn
|
||||||
|
|
||||||
|
# ---- public API (matches upstream sampler.sample) ----
|
||||||
|
def sample(self, model, noise, cond=None, neg_cond=None, steps: int = 50,
|
||||||
|
rescale_t: float = 1.0, guidance_strength: float = 7.5,
|
||||||
|
guidance_rescale: float = 0.0, guidance_interval=(0.0, 1.0),
|
||||||
|
verbose: bool = True, tqdm_desc: str = "Sampling", **kwargs):
|
||||||
|
raw = getattr(model, "_mlx", model) # unwrap MlxFlowModelAdapter
|
||||||
|
if isinstance(noise, torch.Tensor):
|
||||||
|
return self._sample_dense(
|
||||||
|
raw, noise, cond, neg_cond, steps, rescale_t,
|
||||||
|
guidance_strength, guidance_rescale, guidance_interval,
|
||||||
|
verbose, tqdm_desc)
|
||||||
|
return self._sample_sparse(
|
||||||
|
raw, noise, cond, neg_cond, steps, rescale_t,
|
||||||
|
guidance_strength, guidance_rescale, guidance_interval,
|
||||||
|
verbose, tqdm_desc, **kwargs)
|
||||||
|
|
||||||
|
# ---- dense (sparse-structure stage): batched CFG ----
|
||||||
|
def _sample_dense(self, raw, noise, cond, neg_cond, steps, rescale_t,
|
||||||
|
g, g_rescale, g_interval, verbose, desc):
|
||||||
|
x = _torch_to_mx(noise)
|
||||||
|
B = x.shape[0]
|
||||||
|
c = _torch_to_mx(cond)
|
||||||
|
n = _torch_to_mx(neg_cond) if neg_cond is not None else None
|
||||||
|
c2 = mx.concatenate([c, n], axis=0) if n is not None else None
|
||||||
|
feat_axes = tuple(range(1, x.ndim))
|
||||||
|
|
||||||
|
fwd = self._maybe_compile(lambda xx, tt, cc: raw(xx, tt, cc))
|
||||||
|
|
||||||
|
for t, t_prev in tqdm(self._t_pairs(steps, rescale_t), desc=desc,
|
||||||
|
disable=not verbose):
|
||||||
|
in_iv = g_interval[0] <= t <= g_interval[1]
|
||||||
|
eff_g = g if in_iv else 1.0
|
||||||
|
t_m = mx.full((B,), 1000.0 * t, dtype=mx.float32)
|
||||||
|
if eff_g == 1.0 or n is None:
|
||||||
|
pred = fwd(x, t_m, c)
|
||||||
|
elif eff_g == 0.0:
|
||||||
|
pred = fwd(x, t_m, n)
|
||||||
|
else:
|
||||||
|
x2 = mx.concatenate([x, x], axis=0)
|
||||||
|
t2 = mx.concatenate([t_m, t_m], axis=0)
|
||||||
|
v2 = fwd(x2, t2, c2)
|
||||||
|
pred = self._cfg_mix(x, t, v2[:B], v2[B:], eff_g,
|
||||||
|
g_rescale if in_iv else 0.0, feat_axes)
|
||||||
|
x = x - (t - t_prev) * pred
|
||||||
|
mx.eval(x)
|
||||||
|
return edict({"samples": _mx_to_torch(x)})
|
||||||
|
|
||||||
|
# ---- sparse (slat stages): once-per-stage conversion, cached tensor ----
|
||||||
|
def _sample_sparse(self, raw, noise, cond, neg_cond, steps, rescale_t,
|
||||||
|
g, g_rescale, g_interval, verbose, desc, **kwargs):
|
||||||
|
xs = _sparse_to_mlx(noise) # coords fixed for the whole stage
|
||||||
|
template = xs # .replace() keeps spatial caches
|
||||||
|
c = _torch_to_mx(cond)
|
||||||
|
n = _torch_to_mx(neg_cond) if neg_cond is not None else None
|
||||||
|
cc = None
|
||||||
|
if kwargs.get("concat_cond") is not None:
|
||||||
|
cc = _sparse_to_mlx(kwargs["concat_cond"]) # ONCE (was per-step)
|
||||||
|
|
||||||
|
def step_fwd(feats, t_m, cond_a):
|
||||||
|
xt = template.replace(feats)
|
||||||
|
out = raw(xt, t_m, cond_a, **({"concat_cond": cc} if cc is not None else {}))
|
||||||
|
return out.feats
|
||||||
|
|
||||||
|
fwd = self._maybe_compile(step_fwd)
|
||||||
|
|
||||||
|
feats = xs.feats
|
||||||
|
# batch=1 sparse: std over all feat elements (== upstream non-batch dims)
|
||||||
|
feat_axes = tuple(range(feats.ndim))
|
||||||
|
|
||||||
|
for t, t_prev in tqdm(self._t_pairs(steps, rescale_t), desc=desc,
|
||||||
|
disable=not verbose):
|
||||||
|
in_iv = g_interval[0] <= t <= g_interval[1]
|
||||||
|
eff_g = g if in_iv else 1.0
|
||||||
|
t_m = mx.array([1000.0 * t], dtype=mx.float32)
|
||||||
|
if eff_g == 1.0 or n is None:
|
||||||
|
pred = fwd(feats, t_m, c)
|
||||||
|
elif eff_g == 0.0:
|
||||||
|
pred = fwd(feats, t_m, n)
|
||||||
|
else:
|
||||||
|
v_pos = fwd(feats, t_m, c)
|
||||||
|
v_neg = fwd(feats, t_m, n)
|
||||||
|
pred = self._cfg_mix(feats, t, v_pos, v_neg, eff_g,
|
||||||
|
g_rescale if in_iv else 0.0, feat_axes)
|
||||||
|
feats = feats - (t - t_prev) * pred
|
||||||
|
mx.eval(feats)
|
||||||
|
|
||||||
|
return edict({"samples": _mlx_to_sparse(template.replace(feats),
|
||||||
|
ref=noise)})
|
||||||
@ -249,6 +249,18 @@ def create_mlx_pipeline(
|
|||||||
)(**args['tex_slat_sampler']['args'])
|
)(**args['tex_slat_sampler']['args'])
|
||||||
pipeline.tex_slat_sampler_params = args['tex_slat_sampler']['params']
|
pipeline.tex_slat_sampler_params = args['tex_slat_sampler']['params']
|
||||||
|
|
||||||
|
# Pure-MLX sampler loop (rung-3): kills per-forward host bounces + syncs.
|
||||||
|
# Rollback: TRELLIS2_MLX_SAMPLER=0 keeps the upstream torch samplers.
|
||||||
|
import os as _os
|
||||||
|
if _os.environ.get('TRELLIS2_MLX_SAMPLER', '1') == '1':
|
||||||
|
from .mlx_samplers import MlxFlowEulerGuidanceIntervalSampler
|
||||||
|
for _stage in ('sparse_structure_sampler', 'shape_slat_sampler',
|
||||||
|
'tex_slat_sampler'):
|
||||||
|
if args[_stage]['name'] == 'FlowEulerGuidanceIntervalSampler':
|
||||||
|
setattr(pipeline, _stage,
|
||||||
|
MlxFlowEulerGuidanceIntervalSampler(**args[_stage]['args']))
|
||||||
|
print('[MLX] Pure-MLX sampler loop active (TRELLIS2_MLX_SAMPLER=0 to revert)')
|
||||||
|
|
||||||
# Normalization
|
# Normalization
|
||||||
pipeline.shape_slat_normalization = args['shape_slat_normalization']
|
pipeline.shape_slat_normalization = args['shape_slat_normalization']
|
||||||
pipeline.tex_slat_normalization = args['tex_slat_normalization']
|
pipeline.tex_slat_normalization = args['tex_slat_normalization']
|
||||||
|
|||||||
@ -3,6 +3,7 @@ Submanifold sparse 3D convolution for MLX.
|
|||||||
Port of conv_pytorch.py algorithm: hash → neighbor map → gather → bmm → scatter.
|
Port of conv_pytorch.py algorithm: hash → neighbor map → gather → bmm → scatter.
|
||||||
"""
|
"""
|
||||||
import itertools
|
import itertools
|
||||||
|
import os
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
import mlx.core as mx
|
import mlx.core as mx
|
||||||
@ -129,6 +130,16 @@ class MlxSparseConv3d(nn.Module):
|
|||||||
w = self.weight.reshape(Co, K, Ci) # (Co, K, Ci)
|
w = self.weight.reshape(Co, K, Ci) # (Co, K, Ci)
|
||||||
w = w.transpose(1, 2, 0) # (K, Ci, Co)
|
w = w.transpose(1, 2, 0) # (K, Ci, Co)
|
||||||
|
|
||||||
|
# TRELLIS2_METAL_SPCONV=1: fused gather-GEMM Metal kernel — never
|
||||||
|
# materializes the (K,N,Ci) gather (stock path peaks ~77GB at
|
||||||
|
# decoder scale and eval-syncs between 1-2-offset chunks).
|
||||||
|
if os.environ.get("TRELLIS2_METAL_SPCONV", "0") == "1":
|
||||||
|
from .sparse_conv_metal import sparse_conv_metal
|
||||||
|
feats_padded = mx.concatenate([
|
||||||
|
x.feats, mx.zeros((1, Ci), dtype=x.feats.dtype)], axis=0)
|
||||||
|
result = sparse_conv_metal(feats_padded, neighbor_map, w, self.bias)
|
||||||
|
return x.replace(result.astype(x.feats.dtype))
|
||||||
|
|
||||||
# Pad feats with zero row for out-of-bounds indices
|
# Pad feats with zero row for out-of-bounds indices
|
||||||
# feats_padded[N] is zeros, so invalid neighbor lookups produce zero
|
# feats_padded[N] is zeros, so invalid neighbor lookups produce zero
|
||||||
# from matmul naturally — no valid_mask needed.
|
# from matmul naturally — no valid_mask needed.
|
||||||
|
|||||||
73
mlx_backend/sparse_conv_metal.py
Normal file
73
mlx_backend/sparse_conv_metal.py
Normal file
@ -0,0 +1,73 @@
|
|||||||
|
"""Fused gather-GEMM-accumulate Metal kernel for submanifold sparse conv.
|
||||||
|
|
||||||
|
The stock MlxSparseConv3d materializes a (K, N, Ci) gather — ~77GB at
|
||||||
|
decoder scale — then chunks it 1-2 offsets at a time with an mx.eval sync
|
||||||
|
between chunks. This kernel never materializes the gather: each thread
|
||||||
|
accumulates out[n, co] = Σ_k Σ_ci feats[nbr[k,n], ci] · w[k, ci, co]
|
||||||
|
straight from the source buffers. Correctness-first implementation
|
||||||
|
(scalar loops, fp32 accumulation); simdgroup-matrix tiling is a later
|
||||||
|
optimization if profiling ever says the decoders matter for *time*
|
||||||
|
(they are ~10s/gen — this kernel's prize is PEAK MEMORY, which gates the
|
||||||
|
M1 Ultra).
|
||||||
|
|
||||||
|
Env gate: TRELLIS2_METAL_SPCONV=1 routes MlxSparseConv3d through this.
|
||||||
|
"""
|
||||||
|
import mlx.core as mx
|
||||||
|
|
||||||
|
_SRC = """
|
||||||
|
uint n = thread_position_in_grid.x; // voxel index
|
||||||
|
uint co = thread_position_in_grid.y; // output channel
|
||||||
|
uint N = (uint)shape_info[0];
|
||||||
|
uint K = (uint)shape_info[1];
|
||||||
|
uint Ci = (uint)shape_info[2];
|
||||||
|
uint Co = (uint)shape_info[3];
|
||||||
|
if (n >= N || co >= Co) return;
|
||||||
|
|
||||||
|
float acc = 0.0f;
|
||||||
|
for (uint k = 0; k < K; ++k) {
|
||||||
|
int nbr = nmap[k * N + n]; // N (==pad row) when invalid
|
||||||
|
if (nbr >= (int)N) continue; // pad row is zeros anyway; skip
|
||||||
|
const device T* frow = feats + (size_t)nbr * Ci;
|
||||||
|
const device T* wrow = w + ((size_t)k * Ci) * Co + co;
|
||||||
|
for (uint ci = 0; ci < Ci; ++ci) {
|
||||||
|
acc += (float)frow[ci] * (float)wrow[(size_t)ci * Co];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out[(size_t)n * Co + co] = (T)acc;
|
||||||
|
"""
|
||||||
|
|
||||||
|
_kernel = None
|
||||||
|
|
||||||
|
|
||||||
|
def _get_kernel():
|
||||||
|
global _kernel
|
||||||
|
if _kernel is None:
|
||||||
|
_kernel = mx.fast.metal_kernel(
|
||||||
|
name="submconv3d_gather_gemm",
|
||||||
|
input_names=["feats", "nmap", "w", "shape_info"],
|
||||||
|
output_names=["out"],
|
||||||
|
source=_SRC,
|
||||||
|
)
|
||||||
|
return _kernel
|
||||||
|
|
||||||
|
|
||||||
|
def sparse_conv_metal(feats_padded: mx.array, neighbor_map: mx.array,
|
||||||
|
w_kcico: mx.array, bias: mx.array) -> mx.array:
|
||||||
|
"""feats_padded: (N+1, Ci) — row N is zeros (pad; skipped anyway).
|
||||||
|
neighbor_map: (K, N) int32 with N as the invalid sentinel.
|
||||||
|
w_kcico: (K, Ci, Co). bias: (Co,). Returns (N, Co) in w's dtype."""
|
||||||
|
K, N = neighbor_map.shape
|
||||||
|
Ci = feats_padded.shape[1]
|
||||||
|
Co = w_kcico.shape[2]
|
||||||
|
shape_info = mx.array([N, K, Ci, Co], dtype=mx.int32)
|
||||||
|
kern = _get_kernel()
|
||||||
|
(out,) = kern(
|
||||||
|
inputs=[feats_padded.astype(w_kcico.dtype), neighbor_map,
|
||||||
|
w_kcico, shape_info],
|
||||||
|
template=[("T", w_kcico.dtype)],
|
||||||
|
grid=(N, Co, 1),
|
||||||
|
threadgroup=(256, 1, 1),
|
||||||
|
output_shapes=[(N, Co)],
|
||||||
|
output_dtypes=[w_kcico.dtype],
|
||||||
|
)
|
||||||
|
return out + bias
|
||||||
@ -37,7 +37,23 @@ from trellis2.model_revisions import ( # noqa: E402
|
|||||||
TRELLIS_REVISION,
|
TRELLIS_REVISION,
|
||||||
)
|
)
|
||||||
|
|
||||||
SAFETY_FACE_TARGET = 200_000
|
def _default_safety_faces() -> int:
|
||||||
|
"""RAM-sized pre-bake face cap (MODELBEAST BENCHMARKS 2026-07-22): the
|
||||||
|
flat 200k laptop mtlbvh guard crushed 2.2M-tri decodes to ~95k-face GLBs
|
||||||
|
(fal ships ~495k from the same weights). 500k is verified crash-free on
|
||||||
|
the M3 Ultra Studio; small-memory boxes keep the safe cap. Env override:
|
||||||
|
TRELLIS2_MAX_BAKE_FACES (0 = uncapped)."""
|
||||||
|
env = os.environ.get("TRELLIS2_MAX_BAKE_FACES")
|
||||||
|
if env is not None:
|
||||||
|
return int(env) or 10_000_000
|
||||||
|
try:
|
||||||
|
mem_gb = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / 1e9
|
||||||
|
except (ValueError, OSError):
|
||||||
|
mem_gb = 0
|
||||||
|
return 500_000 if mem_gb >= 96 else 200_000
|
||||||
|
|
||||||
|
|
||||||
|
SAFETY_FACE_TARGET = _default_safety_faces()
|
||||||
WATCHDOG_SIGNATURES = (
|
WATCHDOG_SIGNATURES = (
|
||||||
"non-zero size",
|
"non-zero size",
|
||||||
"BVH needs at least 8 triangles",
|
"BVH needs at least 8 triangles",
|
||||||
@ -165,6 +181,43 @@ def _export_raw(mesh, path: Path) -> dict[str, Any]:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _force_opaque(glb_path: Path) -> None:
|
||||||
|
"""Rewrite every material's alphaMode to OPAQUE in a GLB, in place.
|
||||||
|
|
||||||
|
The o_voxel Metal baker samples the decoder's alpha attr noisily (same
|
||||||
|
family as the dark-patch base-color bug): solid hair bakes with scattered
|
||||||
|
near-zero alpha texels, to_glb's min<250 auto-detect flips the material
|
||||||
|
to BLEND, and the model renders as a transparent speckle veil.
|
||||||
|
TRELLIS2_ALPHA_MODE=auto keeps the baked behaviour for genuinely
|
||||||
|
transparent subjects (glass etc.); the default assumes solid objects.
|
||||||
|
"""
|
||||||
|
import struct
|
||||||
|
|
||||||
|
data = glb_path.read_bytes()
|
||||||
|
magic, version, _length = struct.unpack_from("<III", data, 0)
|
||||||
|
if magic != 0x46546C67: # b'glTF'
|
||||||
|
return
|
||||||
|
json_len, json_type = struct.unpack_from("<II", data, 12)
|
||||||
|
if json_type != 0x4E4F534A: # b'JSON'
|
||||||
|
return
|
||||||
|
gltf = json.loads(data[20:20 + json_len])
|
||||||
|
changed = False
|
||||||
|
for material in gltf.get("materials") or []:
|
||||||
|
if material.get("alphaMode", "OPAQUE") != "OPAQUE":
|
||||||
|
material["alphaMode"] = "OPAQUE"
|
||||||
|
material.pop("alphaCutoff", None)
|
||||||
|
changed = True
|
||||||
|
if not changed:
|
||||||
|
return
|
||||||
|
payload = json.dumps(gltf, separators=(",", ":")).encode()
|
||||||
|
payload += b" " * (-len(payload) % 4)
|
||||||
|
rest = data[20 + json_len:]
|
||||||
|
out = struct.pack("<III", magic, version, 20 + len(payload) + len(rest))
|
||||||
|
out += struct.pack("<II", len(payload), json_type) + payload + rest
|
||||||
|
glb_path.write_bytes(out)
|
||||||
|
print("Forced alphaMode=OPAQUE (TRELLIS2_ALPHA_MODE=auto keeps baked alpha)")
|
||||||
|
|
||||||
|
|
||||||
def _metal_baker_available() -> tuple[bool, str]:
|
def _metal_baker_available() -> tuple[bool, str]:
|
||||||
try:
|
try:
|
||||||
from trellis2.backends import probe_metal_backends
|
from trellis2.backends import probe_metal_backends
|
||||||
@ -212,6 +265,40 @@ def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_
|
|||||||
faces = torch.from_numpy(simplified_faces).to(dtype=mesh.faces.dtype)
|
faces = torch.from_numpy(simplified_faces).to(dtype=mesh.faces.dtype)
|
||||||
pre_simplified = True
|
pre_simplified = True
|
||||||
|
|
||||||
|
if baker == "vertex":
|
||||||
|
# Dark-patch fix (MODELBEAST BENCHMARKS 2026-07-19): the Metal texel
|
||||||
|
# sampler mangles clean decoder attrs; sample the voxel grid at mesh
|
||||||
|
# vertices (cKDTree IDW, ~0.3s) and ship linear vertex colors.
|
||||||
|
import numpy as np
|
||||||
|
import trimesh
|
||||||
|
from scipy.spatial import cKDTree
|
||||||
|
|
||||||
|
v_np = vertices.numpy()
|
||||||
|
f_np = faces.numpy()
|
||||||
|
coords_np = mesh.coords.cpu().numpy().astype(np.float32)
|
||||||
|
if coords_np.shape[1] == 4:
|
||||||
|
coords_np = coords_np[:, 1:4]
|
||||||
|
attrs_np = mesh.attrs.float().cpu().numpy()
|
||||||
|
vsz = float(mesh.voxel_size)
|
||||||
|
origin = np.array([-0.5, -0.5, -0.5], np.float32)
|
||||||
|
tree = cKDTree(coords_np * vsz + origin + vsz * 0.5)
|
||||||
|
d, idx = tree.query(v_np, k=4, workers=-1)
|
||||||
|
w = 1.0 / (d + vsz * 0.05) ** 2
|
||||||
|
w[d > vsz * 1.5] = 0.0
|
||||||
|
ws = w.sum(1, keepdims=True)
|
||||||
|
has = (ws > 0).squeeze()
|
||||||
|
ws[ws == 0] = 1.0
|
||||||
|
rgb = np.clip(((attrs_np[idx] * (w / ws)[..., None]).sum(1))[:, 0:3], 0, 1)
|
||||||
|
rgb[~has] = np.clip(attrs_np[idx[~has, 0], 0:3], 0, 1)
|
||||||
|
v_gltf = np.stack([v_np[:, 0], v_np[:, 2], -v_np[:, 1]], axis=1)
|
||||||
|
out = trimesh.Trimesh(vertices=v_gltf, faces=f_np, process=False)
|
||||||
|
rgba = np.concatenate([rgb, np.ones((len(v_gltf), 1))], 1)
|
||||||
|
out.visual = trimesh.visual.ColorVisuals(
|
||||||
|
out, vertex_colors=(rgba * 255).astype(np.uint8)) # LINEAR
|
||||||
|
_ = out.vertex_normals # materialize so glTF gets NORMAL (see _raw_trimesh)
|
||||||
|
out.export(str(path))
|
||||||
|
return out, pre_simplified
|
||||||
|
|
||||||
if baker == "metal":
|
if baker == "metal":
|
||||||
available, reason = _metal_baker_available()
|
available, reason = _metal_baker_available()
|
||||||
if not available:
|
if not available:
|
||||||
@ -239,11 +326,18 @@ def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_
|
|||||||
verbose=True,
|
verbose=True,
|
||||||
)
|
)
|
||||||
result.export(path)
|
result.export(path)
|
||||||
|
if os.environ.get("TRELLIS2_ALPHA_MODE", "opaque") != "auto":
|
||||||
|
_force_opaque(Path(path))
|
||||||
return result, pre_simplified
|
return result, pre_simplified
|
||||||
|
|
||||||
|
|
||||||
def _attempt_schedule(preferred: str, requested_target: Optional[int], raw_faces: int):
|
def _attempt_schedule(preferred: str, requested_target: Optional[int], raw_faces: int):
|
||||||
preferred_order = ["metal", "kdtree"] if preferred in {"auto", "metal"} else ["kdtree"]
|
if preferred == "vertex":
|
||||||
|
preferred_order = ["vertex"] # direct vertex-color bake; no fallback needed
|
||||||
|
elif preferred in {"auto", "metal"}:
|
||||||
|
preferred_order = ["metal", "kdtree"]
|
||||||
|
else:
|
||||||
|
preferred_order = ["kdtree"]
|
||||||
targets = [requested_target]
|
targets = [requested_target]
|
||||||
if requested_target is None and raw_faces > SAFETY_FACE_TARGET:
|
if requested_target is None and raw_faces > SAFETY_FACE_TARGET:
|
||||||
targets.append(SAFETY_FACE_TARGET)
|
targets.append(SAFETY_FACE_TARGET)
|
||||||
@ -394,11 +488,11 @@ def _parse_args():
|
|||||||
parser.add_argument("image", type=Path)
|
parser.add_argument("image", type=Path)
|
||||||
parser.add_argument("--output-dir", type=Path, required=True)
|
parser.add_argument("--output-dir", type=Path, required=True)
|
||||||
parser.add_argument("--backend", choices=("auto", "mps", "mlx-experimental"), default="auto")
|
parser.add_argument("--backend", choices=("auto", "mps", "mlx-experimental"), default="auto")
|
||||||
parser.add_argument("--baker", choices=("auto", "metal", "kdtree"), default="auto")
|
parser.add_argument("--baker", choices=("auto", "metal", "kdtree", "vertex"), default="auto")
|
||||||
parser.add_argument("--pipeline-type", choices=("512", "1024", "1024_cascade"), default="512")
|
parser.add_argument("--pipeline-type", choices=("512", "1024", "1024_cascade"), default="512")
|
||||||
parser.add_argument("--seed", type=int, default=42)
|
parser.add_argument("--seed", type=int, default=42)
|
||||||
parser.add_argument("--steps", type=int)
|
parser.add_argument("--steps", type=int)
|
||||||
parser.add_argument("--texture-size", type=int, choices=(512, 1024, 2048), default=1024)
|
parser.add_argument("--texture-size", type=int, choices=(512, 1024, 2048, 4096), default=1024)
|
||||||
parser.add_argument("--background", choices=("auto", "keep"), default="auto")
|
parser.add_argument("--background", choices=("auto", "keep"), default="auto")
|
||||||
parser.add_argument("--pbr-decimation-target", type=_parse_target, default=None, metavar="none|N")
|
parser.add_argument("--pbr-decimation-target", type=_parse_target, default=None, metavar="none|N")
|
||||||
parser.add_argument("--cache-dir", type=Path)
|
parser.add_argument("--cache-dir", type=Path)
|
||||||
|
|||||||
85
test_spconv_metal.py
Normal file
85
test_spconv_metal.py
Normal file
@ -0,0 +1,85 @@
|
|||||||
|
"""Parity + memory gate: fused Metal sparse conv vs the stock chunked path.
|
||||||
|
Small case first (exact check vs naive reference), then a decoder-scale
|
||||||
|
case for peak-memory + timing comparison.
|
||||||
|
"""
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import mlx.core as mx
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
ROOT = Path(__file__).resolve().parent
|
||||||
|
sys.path.insert(0, str(ROOT))
|
||||||
|
|
||||||
|
from mlx_backend.sparse_conv import MlxSparseConv3d, _build_neighbor_map
|
||||||
|
from mlx_backend.sparse_conv_metal import sparse_conv_metal
|
||||||
|
from mlx_backend.sparse_tensor import MlxSparseTensor
|
||||||
|
|
||||||
|
mx.random.seed(0)
|
||||||
|
|
||||||
|
|
||||||
|
def make_case(n_vox, ci, co, res):
|
||||||
|
# random unique voxel coords in a res^3 grid, batch 0
|
||||||
|
rng = np.random.default_rng(0)
|
||||||
|
flat = rng.choice(res ** 3, size=n_vox, replace=False)
|
||||||
|
x, rem = flat // (res * res), flat % (res * res)
|
||||||
|
y, z = rem // res, rem % res
|
||||||
|
coords = np.stack([np.zeros_like(x), x, y, z], 1).astype(np.int32)
|
||||||
|
st = MlxSparseTensor(
|
||||||
|
feats=mx.random.normal((n_vox, ci)).astype(mx.float16),
|
||||||
|
coords=mx.array(coords))
|
||||||
|
conv = MlxSparseConv3d(ci, co, 3)
|
||||||
|
conv.weight = mx.random.normal((co, 3, 3, 3, ci)).astype(mx.float16) * 0.05
|
||||||
|
conv.bias = mx.random.normal((co,)).astype(mx.float16) * 0.1
|
||||||
|
return st, conv
|
||||||
|
|
||||||
|
|
||||||
|
def run_metal(st, conv):
|
||||||
|
nmap = _build_neighbor_map(st.coords, st.shape[0], st.spatial_shape,
|
||||||
|
conv.kernel_size, conv.dilation)
|
||||||
|
mx.eval(nmap)
|
||||||
|
K = nmap.shape[0]
|
||||||
|
Co, Ci = conv.out_channels, conv.in_channels
|
||||||
|
w = conv.weight.reshape(Co, K, Ci).transpose(1, 2, 0) # (K, Ci, Co)
|
||||||
|
feats_padded = mx.concatenate(
|
||||||
|
[st.feats, mx.zeros((1, Ci), dtype=st.feats.dtype)], axis=0)
|
||||||
|
return sparse_conv_metal(feats_padded, nmap, w, conv.bias)
|
||||||
|
|
||||||
|
|
||||||
|
# ---- small exact-parity case ----
|
||||||
|
st, conv = make_case(5000, 64, 96, 64)
|
||||||
|
ref = conv(st).feats
|
||||||
|
got = run_metal(st, conv)
|
||||||
|
mx.eval(ref, got)
|
||||||
|
d = mx.max(mx.abs(ref.astype(mx.float32) - got.astype(mx.float32))).item()
|
||||||
|
rel = d / max(1e-6, mx.max(mx.abs(ref.astype(mx.float32))).item())
|
||||||
|
print(f"small: max|diff|={d:.3e} rel={rel:.3e}", "PASS" if rel < 2e-2 else "FAIL")
|
||||||
|
ok_small = rel < 2e-2
|
||||||
|
|
||||||
|
# ---- decoder-scale case: memory + time ----
|
||||||
|
st2, conv2 = make_case(400_000, 512, 512, 512)
|
||||||
|
mx.eval(st2.feats, conv2.weight)
|
||||||
|
|
||||||
|
mx.reset_peak_memory()
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
ref2 = conv2(st2).feats
|
||||||
|
mx.eval(ref2)
|
||||||
|
t_stock = time.perf_counter() - t0
|
||||||
|
m_stock = mx.get_peak_memory() / 2**30
|
||||||
|
|
||||||
|
st2._spatial_cache = {} # drop cached neighbor map so both build it
|
||||||
|
mx.reset_peak_memory()
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
got2 = run_metal(st2, conv2)
|
||||||
|
mx.eval(got2)
|
||||||
|
t_metal = time.perf_counter() - t0
|
||||||
|
m_metal = mx.get_peak_memory() / 2**30
|
||||||
|
|
||||||
|
d2 = mx.max(mx.abs(ref2.astype(mx.float32) - got2.astype(mx.float32))).item()
|
||||||
|
r2 = d2 / max(1e-6, mx.max(mx.abs(ref2.astype(mx.float32))).item())
|
||||||
|
print(f"large: stock {t_stock:.2f}s peak {m_stock:.2f}GB | "
|
||||||
|
f"metal {t_metal:.2f}s peak {m_metal:.2f}GB | rel diff {r2:.3e}")
|
||||||
|
ok_large = r2 < 2e-2
|
||||||
|
print("SPCONV_GATE:", "PASS" if (ok_small and ok_large) else "FAIL")
|
||||||
|
sys.exit(0 if (ok_small and ok_large) else 1)
|
||||||
@ -83,8 +83,13 @@ def inspect_glb(path: Path, *, require_pbr: bool) -> dict[str, Any]:
|
|||||||
if int(primitive.get("mode", 4)) != 4:
|
if int(primitive.get("mode", 4)) != 4:
|
||||||
raise ValueError("only TRIANGLES primitives are supported")
|
raise ValueError("only TRIANGLES primitives are supported")
|
||||||
attributes = primitive.get("attributes", {})
|
attributes = primitive.get("attributes", {})
|
||||||
required = ["POSITION", "NORMAL"] + (["TEXCOORD_0"] if require_pbr else [])
|
required = ["POSITION", "NORMAL"]
|
||||||
missing = [name for name in required if name not in attributes]
|
missing = [name for name in required if name not in attributes]
|
||||||
|
# a color-bearing static asset may carry UVs (texture bake) OR
|
||||||
|
# per-vertex colors (vertex bake) — either satisfies require_pbr
|
||||||
|
if require_pbr and "TEXCOORD_0" not in attributes \
|
||||||
|
and "COLOR_0" not in attributes:
|
||||||
|
missing.append("TEXCOORD_0|COLOR_0")
|
||||||
if missing:
|
if missing:
|
||||||
raise ValueError(f"primitive is missing attributes: {', '.join(missing)}")
|
raise ValueError(f"primitive is missing attributes: {', '.join(missing)}")
|
||||||
position = document["accessors"][int(attributes["POSITION"])]
|
position = document["accessors"][int(attributes["POSITION"])]
|
||||||
@ -103,7 +108,9 @@ def inspect_glb(path: Path, *, require_pbr: bool) -> dict[str, Any]:
|
|||||||
raise ValueError("invalid triangle index count")
|
raise ValueError("invalid triangle index count")
|
||||||
triangles += index_count // 3
|
triangles += index_count // 3
|
||||||
|
|
||||||
if require_pbr:
|
_vertex_colored = ("COLOR_0" in attributes
|
||||||
|
and "TEXCOORD_0" not in attributes)
|
||||||
|
if require_pbr and not _vertex_colored:
|
||||||
if "material" not in primitive:
|
if "material" not in primitive:
|
||||||
raise ValueError("PBR primitive has no material")
|
raise ValueError("PBR primitive has no material")
|
||||||
material = document["materials"][int(primitive["material"])]
|
material = document["materials"][int(primitive["material"])]
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user