Compare commits
No commits in common. "main" and "pr175-head-754d403" have entirely different histories.
main
...
pr175-head
@ -1,159 +0,0 @@
|
||||
"""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,18 +249,6 @@ def create_mlx_pipeline(
|
||||
)(**args['tex_slat_sampler']['args'])
|
||||
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
|
||||
pipeline.shape_slat_normalization = args['shape_slat_normalization']
|
||||
pipeline.tex_slat_normalization = args['tex_slat_normalization']
|
||||
|
||||
@ -3,7 +3,6 @@ Submanifold sparse 3D convolution for MLX.
|
||||
Port of conv_pytorch.py algorithm: hash → neighbor map → gather → bmm → scatter.
|
||||
"""
|
||||
import itertools
|
||||
import os
|
||||
import logging
|
||||
import time
|
||||
import mlx.core as mx
|
||||
@ -130,16 +129,6 @@ class MlxSparseConv3d(nn.Module):
|
||||
w = self.weight.reshape(Co, K, Ci) # (Co, K, Ci)
|
||||
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
|
||||
# feats_padded[N] is zeros, so invalid neighbor lookups produce zero
|
||||
# from matmul naturally — no valid_mask needed.
|
||||
|
||||
@ -1,73 +0,0 @@
|
||||
"""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
|
||||
@ -212,40 +212,6 @@ def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_
|
||||
faces = torch.from_numpy(simplified_faces).to(dtype=mesh.faces.dtype)
|
||||
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":
|
||||
available, reason = _metal_baker_available()
|
||||
if not available:
|
||||
@ -277,12 +243,7 @@ def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_
|
||||
|
||||
|
||||
def _attempt_schedule(preferred: str, requested_target: Optional[int], raw_faces: int):
|
||||
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"]
|
||||
preferred_order = ["metal", "kdtree"] if preferred in {"auto", "metal"} else ["kdtree"]
|
||||
targets = [requested_target]
|
||||
if requested_target is None and raw_faces > SAFETY_FACE_TARGET:
|
||||
targets.append(SAFETY_FACE_TARGET)
|
||||
@ -433,7 +394,7 @@ def _parse_args():
|
||||
parser.add_argument("image", type=Path)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument("--backend", choices=("auto", "mps", "mlx-experimental"), default="auto")
|
||||
parser.add_argument("--baker", choices=("auto", "metal", "kdtree", "vertex"), default="auto")
|
||||
parser.add_argument("--baker", choices=("auto", "metal", "kdtree"), default="auto")
|
||||
parser.add_argument("--pipeline-type", choices=("512", "1024", "1024_cascade"), default="512")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--steps", type=int)
|
||||
|
||||
@ -1,85 +0,0 @@
|
||||
"""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,13 +83,8 @@ def inspect_glb(path: Path, *, require_pbr: bool) -> dict[str, Any]:
|
||||
if int(primitive.get("mode", 4)) != 4:
|
||||
raise ValueError("only TRIANGLES primitives are supported")
|
||||
attributes = primitive.get("attributes", {})
|
||||
required = ["POSITION", "NORMAL"]
|
||||
required = ["POSITION", "NORMAL"] + (["TEXCOORD_0"] if require_pbr else [])
|
||||
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:
|
||||
raise ValueError(f"primitive is missing attributes: {', '.join(missing)}")
|
||||
position = document["accessors"][int(attributes["POSITION"])]
|
||||
@ -108,9 +103,7 @@ def inspect_glb(path: Path, *, require_pbr: bool) -> dict[str, Any]:
|
||||
raise ValueError("invalid triangle index count")
|
||||
triangles += index_count // 3
|
||||
|
||||
_vertex_colored = ("COLOR_0" in attributes
|
||||
and "TEXCOORD_0" not in attributes)
|
||||
if require_pbr and not _vertex_colored:
|
||||
if require_pbr:
|
||||
if "material" not in primitive:
|
||||
raise ValueError("PBR primitive has no material")
|
||||
material = document["materials"][int(primitive["material"])]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user