diff --git a/docs/development.md b/docs/development.md index d24420c..5d11664 100644 --- a/docs/development.md +++ b/docs/development.md @@ -77,13 +77,20 @@ user-requested inference results, not test or audit scratch data. Preserve the temporary root only when its artifacts are needed for diagnosis; otherwise remove it after recording the relevant result. -## TRELLIS.2 Miniature Golden Fixture +## Model-Independent Golden Fixtures + +Golden coverage uses two complementary fixture types. Synthetic miniature +checkpoints exercise complete inference orchestration. Real-weight-derived +decoder patches preserve a reviewed real-model boundary without committing or +loading the source checkpoints. + +### TRELLIS.2 Synthetic Miniature The TRELLIS.2 golden test does not read `weights/` or download model assets. It generates a miniature source checkpoint, applies the production selective INT8 quantizer, and runs image conditioning, sparse sampling, shape and texture SLat -sampling, both decoders, artifact serialization, and GLB export without stage -mocking. +sampling, both decoders, artifact serialization, and the real SpatialKit export +path with miniature export settings. ```bash uv run pytest -m heavy tests/test_trellis2_golden_fixture.py -q @@ -94,6 +101,35 @@ Reviewed tensor and GLB expectations live in automatically during tests. Rebaseline it only after reviewing an intentional inference-contract change. +### Pixal3D Real-Weight-Derived Decoder Patch + +`tests/data/pixal3d_derived_golden/` contains a 16-cubed spatial patch captured +from an official-sample run of `TencentARC/Pixal3D` revision +`0b31f9160aa400719af409098bff7936a932f726`. The source used unquantized BF16 +flow and FP16 decoder checkpoints. The committed patch is under 64 KiB and +replays decoded O-Voxel validation, mesh extraction, remeshing, MLX QEM, UV, +texture baking, and GLB writing without the 22 GB source bundle. + +```bash +uv run pytest tests/test_pixal3d_derived_golden.py -q +uv run pytest -m heavy tests/test_pixal3d_derived_golden.py -q +``` + +The default test verifies provenance, checksums, and decoded contracts. The +Metal-backed replay is marked `heavy` and normally completes in under one +second. It does not replace the synthetic full-pipeline Pixal3D tests: the +derived fixture begins at the decoder-output boundary. + +Rebaseline only from a reviewed real inference result: + +```bash +uv run python scripts/pixal3d/write_derived_golden_fixture.py \ + outputs/pixal3d/real-smoke-moge-balanced-decoders1100k \ + outputs/pixal3d/real-smoke-moge-balanced-decoders1100k/trace.json \ + tests/data/pixal3d_derived_golden \ + --source-revision 0b31f9160aa400719af409098bff7936a932f726 +``` + ## Editing Constraints - Prefer existing module boundaries over new abstractions. diff --git a/scripts/pixal3d/write_derived_golden_fixture.py b/scripts/pixal3d/write_derived_golden_fixture.py new file mode 100644 index 0000000..0575933 --- /dev/null +++ b/scripts/pixal3d/write_derived_golden_fixture.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 +"""Write a compact real-weight-derived Pixal3D decoder/export fixture.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import tempfile +from pathlib import Path +from typing import Any + +import numpy as np + + +SOURCE_REPOSITORY = "TencentARC/Pixal3D" +SHAPE_FILENAME = "shape_decoder_fields.npz" +TEXTURE_FILENAME = "texture_decoder_pbr.npz" +MANIFEST_FILENAME = "golden.json" + + +def main() -> int: + args = _parser().parse_args() + # Keep --help and argument validation usable without a Metal device. + from mlx_spatial.spatialkit import export_decoded_ovoxel_glb, inspect_glb + + source_dir = args.source_decoded_dir + output_dir = args.output_dir + trace = json.loads(args.source_trace.read_text(encoding="utf-8")) + + with np.load(source_dir / SHAPE_FILENAME, allow_pickle=False) as payload: + source_coordinates = np.asarray(payload["coordinates"]) + source_fields = np.asarray(payload["fields"]) + with np.load(source_dir / TEXTURE_FILENAME, allow_pickle=False) as payload: + texture_coordinates = np.asarray(payload["coordinates"]) + source_attributes = np.asarray(payload["attributes"]) + source_grid_size = int(payload["decode_resolution"].item()) + + if not np.array_equal(source_coordinates, texture_coordinates): + raise ValueError("source shape and texture coordinates differ") + + origin = np.asarray(args.origin, dtype=np.int32) + local_grid_size = int(args.local_grid_size) + spatial = source_coordinates[:, 1:] + mask = np.all((spatial >= origin) & (spatial < origin + local_grid_size), axis=1) + if not np.any(mask): + raise ValueError("selected local patch is empty") + + coordinates = np.ascontiguousarray(source_coordinates[mask]) + coordinates[:, 1:] -= origin + fields = np.ascontiguousarray(source_fields[mask]) + attributes = np.ascontiguousarray(source_attributes[mask]) + metadata = { + "fixture_kind": "real-weight-derived-decoder-patch", + "local_grid_size": local_grid_size, + "local_origin": origin.tolist(), + "model_family": "pixal3d", + "source_coordinate_system": "gltf-y-up", + "source_grid_size": source_grid_size, + "source_repository": SOURCE_REPOSITORY, + "source_revision": args.source_revision, + } + + output_dir.mkdir(parents=True, exist_ok=True) + shape_path = output_dir / SHAPE_FILENAME + texture_path = output_dir / TEXTURE_FILENAME + metadata_json = np.asarray(json.dumps(metadata, sort_keys=True)) + np.savez_compressed( + shape_path, + coordinates=coordinates, + fields=fields, + metadata_json=metadata_json, + ) + np.savez_compressed( + texture_path, + coordinates=coordinates, + attributes=attributes, + spatial_shape=np.asarray([local_grid_size] * 3, dtype=np.int32), + batch_size=np.asarray(1, dtype=np.int32), + decode_resolution=np.asarray(local_grid_size, dtype=np.int32), + voxel_size=np.asarray(1.0 / local_grid_size, dtype=np.float32), + metadata_json=metadata_json, + ) + + export_settings = { + "grid_size": local_grid_size, + "quality_preset": "reference-target", + "remesh": True, + "remesh_resolution": local_grid_size, + "simplify_backend": "mlx-qem", + "target_faces": int(args.target_faces), + "texture_postprocess": "telea", + "texture_size": int(args.texture_size), + "uv_backend": "xatlas-equivalent-native", + } + with tempfile.TemporaryDirectory(prefix="pixal3d-derived-golden.") as directory: + result = export_decoded_ovoxel_glb( + output_dir, + Path(directory) / "model.glb", + **export_settings, + ) + glb = inspect_glb(result.glb.path) + + source_input = Path(str(trace.get("image_path", ""))) + manifest = { + "schema_version": 1, + "fixture_kind": "real-weight-derived-decoder-patch", + "scope": { + "covered": "decoded O-Voxel contract through native textured GLB export", + "not_covered": ( + "Pixal3D checkpoint loading, conditioning, flow sampling, " + "and decoder execution" + ), + }, + "source": { + "repository": SOURCE_REPOSITORY, + "revision": args.source_revision, + "precision": "unquantized mixed BF16 flow and FP16 decoder checkpoints", + "quantization": "none", + "pipeline_type": trace.get("pipeline_type"), + "seed": trace.get("seed"), + "grid_size": source_grid_size, + "input_sha256": _sha256(source_input) if source_input.is_file() else None, + "trace_sha256": _sha256(args.source_trace), + "completed_stages": trace.get("completed_stages", []), + }, + "selection": { + "origin": origin.tolist(), + "grid_size": local_grid_size, + "source_token_count": int(source_coordinates.shape[0]), + "selected_token_count": int(coordinates.shape[0]), + }, + "files": { + SHAPE_FILENAME: { + "bytes": shape_path.stat().st_size, + "sha256": _sha256(shape_path), + }, + TEXTURE_FILENAME: { + "bytes": texture_path.stat().st_size, + "sha256": _sha256(texture_path), + }, + }, + "arrays": { + "coordinates": _array_summary(coordinates), + "fields": _array_summary(fields), + "attributes": _array_summary(attributes), + }, + "expected_export": { + "settings": export_settings, + "source_vertices": int( + result.diagnostics["stages"]["extract_mesh"]["source_vertices"] + ), + "source_faces": int( + result.diagnostics["stages"]["extract_mesh"]["source_faces"] + ), + "final_faces": int( + result.diagnostics["stages"]["simplify_mesh"]["stats"]["final_faces"] + ), + "glb": { + "meshes": int(glb["mesh_count"]), + "primitives": int(glb["primitive_count"]), + "materials": int(glb["material_count"]), + "textures": int(glb["texture_count"]), + "images": int(glb["image_count"]), + "vertices": int(glb["total_vertices"]), + "faces": int(glb["total_faces"]), + }, + }, + } + (output_dir / MANIFEST_FILENAME).write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return 0 + + +def _array_summary(array: np.ndarray) -> dict[str, Any]: + summary: dict[str, Any] = {"shape": list(array.shape), "dtype": str(array.dtype)} + if np.issubdtype(array.dtype, np.integer): + summary["sha256"] = hashlib.sha256( + np.ascontiguousarray(array).tobytes() + ).hexdigest() + else: + values = np.asarray(array, dtype=np.float64) + summary["statistics"] = { + "min": float(values.min()), + "max": float(values.max()), + "mean": float(values.mean()), + "std": float(values.std()), + "l2": float(np.linalg.norm(values)), + } + return summary + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source_decoded_dir", type=Path) + parser.add_argument("source_trace", type=Path) + parser.add_argument("output_dir", type=Path) + parser.add_argument("--source-revision", required=True) + parser.add_argument( + "--origin", + type=int, + nargs=3, + default=(224, 496, 688), + metavar=("Z", "Y", "X"), + ) + parser.add_argument("--local-grid-size", type=int, default=16) + parser.add_argument("--target-faces", type=int, default=512) + parser.add_argument("--texture-size", type=int, default=16) + return parser + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/data/pixal3d_derived_golden/golden.json b/tests/data/pixal3d_derived_golden/golden.json new file mode 100644 index 0000000..9a33eb9 --- /dev/null +++ b/tests/data/pixal3d_derived_golden/golden.json @@ -0,0 +1,131 @@ +{ + "arrays": { + "attributes": { + "dtype": "float32", + "shape": [ + 1166, + 6 + ], + "statistics": { + "l2": 54.5424308341348, + "max": 1.0007288455963135, + "mean": 0.5237411886598465, + "min": 0.012211322784423828, + "std": 0.3884849388377122 + } + }, + "coordinates": { + "dtype": "int32", + "sha256": "1e352c7b7638293edf65e86c4afec0dac2078a0137d5249bdbf8accd37e4750e", + "shape": [ + 1166, + 4 + ] + }, + "fields": { + "dtype": "float32", + "shape": [ + 1166, + 7 + ], + "statistics": { + "l2": 882.6670673262066, + "max": 19.59222412109375, + "mean": -3.7514042070301383, + "min": -61.986366271972656, + "std": 9.021178029521382 + } + } + }, + "expected_export": { + "final_faces": 512, + "glb": { + "faces": 512, + "images": 2, + "materials": 1, + "meshes": 1, + "primitives": 1, + "textures": 2, + "vertices": 442 + }, + "settings": { + "grid_size": 16, + "quality_preset": "reference-target", + "remesh": true, + "remesh_resolution": 16, + "simplify_backend": "mlx-qem", + "target_faces": 512, + "texture_postprocess": "telea", + "texture_size": 16, + "uv_backend": "xatlas-equivalent-native" + }, + "source_faces": 2130, + "source_vertices": 1166 + }, + "files": { + "shape_decoder_fields.npz": { + "bytes": 34107, + "sha256": "d9e2f2aa86cbee5efe844d3bf303c6bb2424c06be3b929f8bc66664e9e2c8e99" + }, + "texture_decoder_pbr.npz": { + "bytes": 28004, + "sha256": "81fd1ce2b17d63c79fc318b098fdf44698cbee20bc1d9e63daeb9a3c935e3ee2" + } + }, + "fixture_kind": "real-weight-derived-decoder-patch", + "schema_version": 1, + "scope": { + "covered": "decoded O-Voxel contract through native textured GLB export", + "not_covered": "Pixal3D checkpoint loading, conditioning, flow sampling, and decoder execution" + }, + "selection": { + "grid_size": 16, + "origin": [ + 224, + 496, + 688 + ], + "selected_token_count": 1166, + "source_token_count": 4150336 + }, + "source": { + "completed_stages": [ + "input-image", + "asset-validation", + "pipeline-config", + "camera-setup", + "image-conditioning", + "projection-conditioning:ss", + "artifact:sparse_projection", + "sparse-structure-flow", + "sparse-structure-decoding", + "artifact:sparse_structure", + "projection-conditioning:shape_512", + "shape-slat-sampling:512", + "artifact:shape_slat_lr", + "shape-slat-cascade:upsample", + "artifact:shape_slat_hr_coordinates", + "projection-conditioning:shape_1024", + "shape-slat-sampling:1024", + "artifact:shape_slat_hr", + "projection-conditioning:tex_1024", + "texture-slat-sampling:1024", + "artifact:texture_slat", + "shape-decoder", + "artifact:shape_decoder_fields", + "texture-decoder", + "artifact:texture_decoder_pbr", + "mesh-export", + "artifact:textured_glb" + ], + "grid_size": 1024, + "input_sha256": "6959e517ee4bc6852791f69bd6ece696a435abcda8321727c07db8daf7f457cf", + "pipeline_type": "1024_cascade", + "precision": "unquantized mixed BF16 flow and FP16 decoder checkpoints", + "quantization": "none", + "repository": "TencentARC/Pixal3D", + "revision": "0b31f9160aa400719af409098bff7936a932f726", + "seed": 42, + "trace_sha256": "b37e1dd38acbdf7190010e4d0d1147ac5c2191c8df6233e7f3132c699282a533" + } +} diff --git a/tests/data/pixal3d_derived_golden/shape_decoder_fields.npz b/tests/data/pixal3d_derived_golden/shape_decoder_fields.npz new file mode 100644 index 0000000..e89855e Binary files /dev/null and b/tests/data/pixal3d_derived_golden/shape_decoder_fields.npz differ diff --git a/tests/data/pixal3d_derived_golden/texture_decoder_pbr.npz b/tests/data/pixal3d_derived_golden/texture_decoder_pbr.npz new file mode 100644 index 0000000..e5327a1 Binary files /dev/null and b/tests/data/pixal3d_derived_golden/texture_decoder_pbr.npz differ diff --git a/tests/data/trellis2_miniature_golden.json b/tests/data/trellis2_miniature_golden.json index 4df5a8b..efaf299 100644 --- a/tests/data/trellis2_miniature_golden.json +++ b/tests/data/trellis2_miniature_golden.json @@ -1,4 +1,23 @@ { + "export": { + "effective": { + "grid_size": 32, + "target_faces": 64, + "texture_size": 16 + }, + "requested": { + "diagnostics_path": null, + "grid_size": 512, + "quality_preset": "reference-target", + "remesh": true, + "remesh_resolution": 512, + "simplify_backend": "mlx-qem", + "target_faces": 256, + "texture_postprocess": "telea", + "texture_size": 32, + "uv_backend": "xatlas-equivalent-native" + } + }, "fixture": { "checkpoint_source": "generated synthetic tensors", "glb_target_faces": 256, @@ -18,14 +37,13 @@ "has_normal": true, "has_texcoord_0": true, "material": 0, - "positions": 4565, - "triangles": 3912 + "positions": 12, + "triangles": 4 } ], - "sha256": "44b55e7184c56449d87a9bed214a7ab24460f9ad3bd159e92867cec67efa9470", "textures": 2 }, - "schema_version": 1, + "schema_version": 2, "trace": { "completed_stages": [ "asset-config-validation", @@ -44,59 +62,95 @@ "tensor_outputs": { "cond_512": { "dtype": "float32", - "sha256": "aad5ce8209c55c84c2e64b92b8162e7e49ba89af0f19270a3762df796b90a5e5", "shape": [ 1, 65, 64 - ] + ], + "statistics": { + "l2": 46.05630751333882, + "max": 1.3148114681243896, + "mean": -1.0446456144563853e-09, + "min": -1.7143969535827637, + "std": 0.7140727348219936 + } }, "shape_flexidualgrid_fields": { "dtype": "float32", - "sha256": "d13460feb6ebeeb09b7fb153ee3c5c607bf1ec2205471396b5a8c5908e5f1f62", "shape": [ 512, 7 - ] + ], + "statistics": { + "l2": 156.75664142258063, + "max": 4.1127119064331055, + "mean": 1.7140144131821475, + "min": -0.14383302628993988, + "std": 1.9794850947718232 + } }, "shape_slat": { "dtype": "float32", - "sha256": "69670f899f1c5f681fb9b3fd71a229e3f4e22271ed6ee470f36af14f32cc31bd", "shape": [ 64, 32 - ] + ], + "statistics": { + "l2": 45.74745586307211, + "max": 3.7916488647460938, + "mean": -0.01963400755175826, + "min": -3.6229214668273926, + "std": 1.010694818658021 + } }, "sparse_latent": { "dtype": "float32", - "sha256": "2fa4d566772d77461fd6630c042e69b0002adf31ee4b724b6557c8a45666ef51", "shape": [ 1, 2, 2, 2, 2 - ] + ], + "statistics": { + "l2": 4.890631746635004, + "max": 2.205594778060913, + "mean": -0.06916660442948341, + "min": -2.760120153427124, + "std": 1.2206999676031494 + } }, "texture_slat": { "dtype": "float32", - "sha256": "99d23e5097506efb79c83865dc63c505ad4fad0db6c8d468d22845baa66b2141", "shape": [ 64, 32 - ] + ], + "statistics": { + "l2": 47.3129009807639, + "max": 3.7781307697296143, + "mean": 0.016508268732650322, + "min": -3.8875527381896973, + "std": 1.0453469426877748 + } }, "texture_voxel_attrs": { "dtype": "float32", - "sha256": "49b45f028a546f707bcdf37677c63eb279187de4561bbe06f1a91720320396b4", "shape": [ 512, 6 - ] + ], + "statistics": { + "l2": 27.73014181281817, + "max": 0.5602271556854248, + "mean": 0.4999949636403471, + "min": 0.43711984157562256, + "std": 0.01782653483942547 + } }, "texture_voxel_coordinates": { "dtype": "int32", - "sha256": "d572dc92213b089492d529519ba3356793c57f19955851405b63c16bc944b436", + "sha256": "893113016750894983cc6e51682d666b9697e3c1834c864977738ff8369a1586", "shape": [ 512, 4 diff --git a/tests/golden_assertions.py b/tests/golden_assertions.py new file mode 100644 index 0000000..a1a7253 --- /dev/null +++ b/tests/golden_assertions.py @@ -0,0 +1,143 @@ +"""Compact numerical summaries and tolerant golden comparisons for MLX fixtures.""" + +from __future__ import annotations + +import hashlib +import json +import struct +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +import mlx.core as mx +import numpy as np + + +def summarize_array(value: Any) -> dict[str, Any]: + """Summarize an array using stable structure and distribution statistics.""" + + if isinstance(value, mx.array): + mx.eval(value) + dtype = str(value.dtype).removeprefix("mlx.core.") + array = np.asarray(value.astype(mx.float32) if dtype == "bfloat16" else value) + else: + array = np.asarray(value) + dtype = str(array.dtype) + + summary: dict[str, Any] = {"shape": list(array.shape), "dtype": dtype} + flat = array.reshape(-1) + if np.issubdtype(array.dtype, np.floating): + finite = np.asarray(flat, dtype=np.float64) + if not np.all(np.isfinite(finite)): + raise ValueError("golden fixture arrays must contain only finite values") + summary["statistics"] = { + "min": float(np.min(finite)), + "max": float(np.max(finite)), + "mean": float(np.mean(finite)), + "std": float(np.std(finite)), + "l2": float(np.linalg.norm(finite)), + } + return summary + + canonical = np.ascontiguousarray(array.astype(array.dtype.newbyteorder("<"), copy=False)) + summary["sha256"] = hashlib.sha256(canonical.tobytes(order="C")).hexdigest() + return summary + + +def assert_golden_close( + actual: Any, + expected: Any, + *, + path: str = "golden", + rtol: float = 5e-4, + atol: float = 1e-4, +) -> None: + """Compare nested golden data exactly except for floating-point leaves.""" + + if isinstance(expected, Mapping): + if not isinstance(actual, Mapping): + raise AssertionError(f"{path}: expected a mapping, got {type(actual).__name__}") + if set(actual) != set(expected): + raise AssertionError( + f"{path}: keys differ; actual={sorted(actual)} expected={sorted(expected)}" + ) + for key in expected: + assert_golden_close( + actual[key], expected[key], path=f"{path}.{key}", rtol=rtol, atol=atol + ) + return + + if isinstance(expected, Sequence) and not isinstance( + expected, (str, bytes, bytearray) + ): + if not isinstance(actual, Sequence) or isinstance(actual, (str, bytes, bytearray)): + raise AssertionError( + f"{path}: expected a sequence, got {type(actual).__name__}" + ) + if len(actual) != len(expected): + raise AssertionError( + f"{path}: length differs; actual={len(actual)} expected={len(expected)}" + ) + for index, expected_value in enumerate(expected): + assert_golden_close( + actual[index], expected_value, path=f"{path}[{index}]", rtol=rtol, atol=atol + ) + return + + if isinstance(expected, float): + if not np.isclose(float(actual), expected, rtol=rtol, atol=atol): + raise AssertionError(f"{path}: actual={actual!r} expected={expected!r}") + return + + if actual != expected: + raise AssertionError(f"{path}: actual={actual!r} expected={expected!r}") + + +def summarize_glb(path: Path) -> dict[str, Any]: + """Read a GLB and summarize stable structural fields without third-party parsers.""" + + payload = path.read_bytes() + if len(payload) < 20: + raise ValueError("GLB payload is too short") + magic, version, declared_length = struct.unpack_from("<4sII", payload, 0) + if magic != b"glTF" or version != 2 or declared_length != len(payload): + raise ValueError("invalid GLB header") + json_length, json_type = struct.unpack_from(" dict: + return json.loads(MANIFEST_PATH.read_text(encoding="utf-8")) + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() diff --git a/tests/test_trellis2_golden_fixture.py b/tests/test_trellis2_golden_fixture.py index bc2b8a7..4a08255 100644 --- a/tests/test_trellis2_golden_fixture.py +++ b/tests/test_trellis2_golden_fixture.py @@ -7,11 +7,16 @@ import pytest +import mlx_spatial.trellis2_inference as trellis2_inference from mlx_spatial.trellis2_inference import Trellis2InferencePipeline from mlx_spatial.trellis2_quantization import read_trellis2_quantization_spec +from tests.golden_assertions import assert_golden_close, summarize_glb from tests.trellis2_golden_fixture import ( + TRELLIS2_MINIATURE_EXPORT_GRID_SIZE, + TRELLIS2_MINIATURE_EXPORT_TARGET_FACES, + TRELLIS2_MINIATURE_EXPORT_TEXTURE_SIZE, + Trellis2MiniatureSpatialKitExporter, build_trellis2_miniature_golden_fixture, - summarize_glb, summarize_trellis2_golden_trace, ) @@ -19,8 +24,14 @@ @pytest.mark.heavy -def test_miniature_int8_pipeline_emits_golden_trace_and_glb(tmp_path): +def test_miniature_int8_pipeline_emits_golden_trace_and_glb(tmp_path, monkeypatch): fixture = build_trellis2_miniature_golden_fixture(tmp_path) + exporter = Trellis2MiniatureSpatialKitExporter() + monkeypatch.setattr( + trellis2_inference, + "load_spatialkit_exporter", + lambda: (exporter, None), + ) quantization = read_trellis2_quantization_spec( fixture.quantized_root / "ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors" ) @@ -52,7 +63,7 @@ def test_miniature_int8_pipeline_emits_golden_trace_and_glb(tmp_path): assert result.ready, result.trace.blocker summary = { - "schema_version": 1, + "schema_version": 2, "fixture": { "checkpoint_source": "generated synthetic tensors", "pipeline_type": "512", @@ -65,6 +76,14 @@ def test_miniature_int8_pipeline_emits_golden_trace_and_glb(tmp_path): }, "trace": summarize_trellis2_golden_trace(result.trace), "glb": summarize_glb(fixture.output_path), + "export": { + "requested": exporter.requested_options, + "effective": { + "grid_size": TRELLIS2_MINIATURE_EXPORT_GRID_SIZE, + "target_faces": TRELLIS2_MINIATURE_EXPORT_TARGET_FACES, + "texture_size": TRELLIS2_MINIATURE_EXPORT_TEXTURE_SIZE, + }, + }, } expected = json.loads(GOLDEN_MANIFEST.read_text(encoding="utf-8")) - assert summary == expected + assert_golden_close(summary, expected) diff --git a/tests/trellis2_golden_fixture.py b/tests/trellis2_golden_fixture.py index 201d332..79c6712 100644 --- a/tests/trellis2_golden_fixture.py +++ b/tests/trellis2_golden_fixture.py @@ -2,10 +2,8 @@ from __future__ import annotations -import hashlib import json import math -import struct from dataclasses import dataclass from pathlib import Path from typing import Any, Mapping @@ -15,12 +13,19 @@ from PIL import Image from mlx_spatial.safetensors_io import save_safetensors +from mlx_spatial.spatialkit import export_decoded_ovoxel_glb from mlx_spatial.trellis2_decode import StructuredLatentDecoderConfig from mlx_spatial.trellis2_dinov3 import DinoV3ModelConfig from mlx_spatial.trellis2_forward import Trellis2ForwardTraceResult from mlx_spatial.trellis2_quantization import quantize_trellis2_weights from mlx_spatial.trellis2_slat import SLatFlowConfig from mlx_spatial.trellis2_sparse_structure import SparseStructureDecoderConfig, SparseStructureFlowConfig +from tests.golden_assertions import summarize_array + + +TRELLIS2_MINIATURE_EXPORT_GRID_SIZE = 32 +TRELLIS2_MINIATURE_EXPORT_TARGET_FACES = 64 +TRELLIS2_MINIATURE_EXPORT_TEXTURE_SIZE = 16 @dataclass(frozen=True) @@ -34,6 +39,30 @@ class Trellis2MiniatureGoldenFixture: output_path: Path +@dataclass +class Trellis2MiniatureSpatialKitExporter: + """Run the real SpatialKit exporter with a miniature remesh policy.""" + + requested_options: dict[str, Any] | None = None + + def __call__(self, decoded_dir: str | Path, output_path: str | Path, **options: Any): + self.requested_options = dict(options) + return export_decoded_ovoxel_glb( + decoded_dir, + output_path, + texture_size=TRELLIS2_MINIATURE_EXPORT_TEXTURE_SIZE, + target_faces=TRELLIS2_MINIATURE_EXPORT_TARGET_FACES, + quality_preset="reference-target", + grid_size=TRELLIS2_MINIATURE_EXPORT_GRID_SIZE, + uv_backend="xatlas-equivalent-native", + remesh=True, + remesh_resolution=TRELLIS2_MINIATURE_EXPORT_GRID_SIZE, + simplify_backend="mlx-qem", + texture_postprocess="telea", + diagnostics_path=options.get("diagnostics_path"), + ) + + def build_trellis2_miniature_golden_fixture(root: Path) -> Trellis2MiniatureGoldenFixture: """Create tiny source checkpoints, quantize them, and return runnable fixture paths.""" @@ -173,55 +202,20 @@ def summarize_trellis2_golden_trace(trace: Trellis2ForwardTraceResult) -> dict[s for output in trace.outputs: if output.payload is None: continue - tensor_outputs[output.name] = { - "shape": list(output.shape), - "dtype": output.dtype, - "sha256": _tensor_digest(output.payload), - } + summary = summarize_array(output.payload) + if summary["shape"] != list(output.shape) or summary["dtype"] != output.dtype: + raise AssertionError( + f"trace metadata disagrees with payload for {output.name}: " + f"declared shape={output.shape} dtype={output.dtype}; " + f"actual shape={summary['shape']} dtype={summary['dtype']}" + ) + tensor_outputs[output.name] = summary return { "completed_stages": list(trace.completed_stages), "tensor_outputs": tensor_outputs, } -def summarize_glb(path: Path) -> dict[str, Any]: - """Read a GLB without third-party parsers and summarize stable structural fields.""" - - payload = path.read_bytes() - if len(payload) < 20: - raise ValueError("GLB payload is too short") - magic, version, declared_length = struct.unpack_from("<4sII", payload, 0) - if magic != b"glTF" or version != 2 or declared_length != len(payload): - raise ValueError("invalid GLB header") - json_length, json_type = struct.unpack_from(" SLatFlowConfig: return SLatFlowConfig( name="SLatFlowModel", @@ -539,29 +533,12 @@ def _center_identity_conv(channels: int) -> mx.array: return mx.array(values) -def _tensor_digest(value: mx.array) -> str: - mx.eval(value) - dtype = str(value.dtype).removeprefix("mlx.core.") - if dtype.startswith("float") or dtype == "bfloat16": - array = np.asarray(value.astype(mx.float32), dtype="