diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 905fffa..197d129 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -4,6 +4,7 @@ on: push: branches: - "**" + pull_request: workflow_dispatch: permissions: @@ -11,7 +12,7 @@ permissions: jobs: test: - name: Fast pytest + name: Routine pytest runs-on: macos-15 steps: - name: Check out repository @@ -37,6 +38,6 @@ jobs: echo "PYTHONPYCACHEPREFIX=$scratch_root/cache/pycache" >> "$GITHUB_ENV" echo "UV_CACHE_DIR=$scratch_root/cache/uv" >> "$GITHUB_ENV" - - name: Run fast tests + - name: Run routine tests timeout-minutes: 10 run: uv run pytest --basetemp "$MLX_SPATIAL_TEST_SCRATCH/artifacts/pytest" diff --git a/docs/development.md b/docs/development.md index 5d11664..c462af2 100644 --- a/docs/development.md +++ b/docs/development.md @@ -77,12 +77,25 @@ 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. -## 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. +## Representative Pipeline Fixtures + +Routine integration coverage uses generated miniature checkpoints, composed +pipeline fixtures, and reviewed real-weight-derived boundary patches. Every +fixture states its provenance and limits. None requires the local production +weight bundles. + +- TRELLIS.2 generates and quantizes a miniature selective INT8 bundle, then + runs through the real SpatialKit GLB path. +- Pixal3D combines a synthetic full-orchestration fixture with a compact + real-weight-derived decoder patch. +- LiTo runs its complete source-contract generation path. +- SAM3D preserves CLI reconstruction, mesh extraction, and GLB writing around + deterministic conditioning and flow boundaries. +- HY-World 2 runs fixture reconstruction and writes its staged artifacts below + pytest temporary storage. +- MapAnything generates a miniature checkpoint and runs asset inspection, + safetensors loading, encoder, multi-view information sharing, prediction + heads, geometry postprocess, and NPZ writing. ### TRELLIS.2 Synthetic Miniature @@ -93,7 +106,7 @@ 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 +uv run pytest tests/test_trellis2_golden_fixture.py -q ``` Reviewed tensor and GLB expectations live in @@ -112,13 +125,15 @@ 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. +bounded Metal-backed replay normally completes in under one second. It does not +replace the synthetic full-pipeline Pixal3D test: the derived fixture begins at +the decoder-output boundary. Face counts and artifact structure remain exact; +the manifest permits a small explicit GLB vertex-count tolerance because native +UV seam splitting can duplicate a few vertices differently across macOS +hardware and driver versions. Rebaseline only from a reviewed real inference result: @@ -130,6 +145,25 @@ uv run python scripts/pixal3d/write_derived_golden_fixture.py \ --source-revision 0b31f9160aa400719af409098bff7936a932f726 ``` +### MapAnything Generated Miniature Scene + +The MapAnything fixture generates its checkpoint at runtime from deterministic +tiny tensors. The committed manifest records fixture provenance, covered and +excluded scope, stage order, output schemas, and tolerant numerical summaries. +It is a pipeline regression fixture, not official-weight parity. Stable outputs +retain numerical summaries; recovered intrinsics and world points use strict +shape, finite-value, homogeneous-matrix, sign, and boundedness invariants to +avoid amplifying cross-hardware noise from the deliberately tiny ray field. + +```bash +uv run pytest tests/test_mapanything_scene_pipeline.py \ + -m 'integration and not real_assets' -q +``` + +Reviewed expectations live in +`tests/data/mapanything_miniature_scene_golden.json`. Rebaseline them only after +reviewing an intentional production pipeline change. + ## Editing Constraints - Prefer existing module boundaries over new abstractions. diff --git a/pyproject.toml b/pyproject.toml index 90ed7b3..6e7f685 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -115,8 +115,12 @@ exclude = [ [tool.pytest.ini_options] testpaths = ["tests"] -addopts = "-m 'not heavy' -p no:cacheprovider" +addopts = "-m 'not (heavy or real_assets or torch_parity)' -p no:cacheprovider" markers = [ - "heavy: opt-in tests that load real weights, require Metal, or use large tensors", + "integration: representative tests that cross multiple production component boundaries", + "real_assets: tests that require local weights, inputs, or other uncommitted fixtures", + "metal: tests that require an Apple Metal device", + "heavy: opt-in tests with substantial runtime or memory cost", + "benchmark: resource and performance checks; always paired with heavy", "torch_parity: optional parity checks against the local PyTorch checkout", ] diff --git a/src/mlx_spatial/trellis2_texturing.py b/src/mlx_spatial/trellis2_texturing.py index 339e6a3..cb62e3c 100644 --- a/src/mlx_spatial/trellis2_texturing.py +++ b/src/mlx_spatial/trellis2_texturing.py @@ -202,14 +202,6 @@ def run( ), ) - preprocessed = preprocess_trellis2_image(image, rmbg_root=self.rmbg_root) - if not preprocessed.ready or preprocessed.image is None: - return Trellis2TexturingResult( - image_path=image, - mesh_path=mesh_file, - blocker=_preprocess_texturing_blocker(preprocessed.blocker), - ) - try: mesh_vertices, mesh_faces = _load_obj_mesh(mesh_file) except (OSError, ValueError) as error: @@ -225,30 +217,6 @@ def run( ), ) - fdg_coords, fdg_dual, fdg_intersected = mesh_to_flexible_dual_grid( - mesh_vertices, mesh_faces, grid_size=grid_size - ) - - if fdg_coords.shape[0] == 0: - return Trellis2TexturingResult( - image_path=image, - mesh_path=mesh_file, - blocker=Trellis2TexturingBlocker( - stage="mesh-preprocess", - operation="FlexiDualGrid voxelization", - reference=str(mesh_path), - reason="mesh_to_flexible_dual_grid produced no occupied voxels", - next_slice="increase grid_size or provide a mesh within the AABB", - ), - ) - - encoder_coords = np.column_stack( - [np.zeros(fdg_coords.shape[0], dtype=np.int32), fdg_coords] - ) - encoder_coords_mx = mx.array(encoder_coords, dtype=mx.int32) - dual_mx = mx.array(fdg_dual, dtype=mx.float32) - intersected_mx = mx.array(fdg_intersected.astype(np.float32), dtype=mx.float32) - discovery = discover_trellis2_conditioning_config(self.root) if not discovery.ready or discovery.config is None: return Trellis2TexturingResult( @@ -302,6 +270,38 @@ def run( texture_slat_sampler=replace(config.texture_slat_sampler, steps=slat_steps), ) + preprocessed = preprocess_trellis2_image(image, rmbg_root=self.rmbg_root) + if not preprocessed.ready or preprocessed.image is None: + return Trellis2TexturingResult( + image_path=image, + mesh_path=mesh_file, + blocker=_preprocess_texturing_blocker(preprocessed.blocker), + ) + + fdg_coords, fdg_dual, fdg_intersected = mesh_to_flexible_dual_grid( + mesh_vertices, mesh_faces, grid_size=grid_size + ) + + if fdg_coords.shape[0] == 0: + return Trellis2TexturingResult( + image_path=image, + mesh_path=mesh_file, + blocker=Trellis2TexturingBlocker( + stage="mesh-preprocess", + operation="FlexiDualGrid voxelization", + reference=str(mesh_path), + reason="mesh_to_flexible_dual_grid produced no occupied voxels", + next_slice="increase grid_size or provide a mesh within the AABB", + ), + ) + + encoder_coords = np.column_stack( + [np.zeros(fdg_coords.shape[0], dtype=np.int32), fdg_coords] + ) + encoder_coords_mx = mx.array(encoder_coords, dtype=mx.int32) + dual_mx = mx.array(fdg_dual, dtype=mx.float32) + intersected_mx = mx.array(fdg_intersected.astype(np.float32), dtype=mx.float32) + resolved_encoder_config_path = self.encoder_config_path or _SHAPE_ENCODER_CONFIG_CONVENTION try: encoder_config = read_structured_latent_encoder_config( diff --git a/tests/README.md b/tests/README.md index 5553b3a..91b42c6 100644 --- a/tests/README.md +++ b/tests/README.md @@ -1,10 +1,63 @@ # Testing Strategy -Default test runs are CPU-bound and fast. `tests/conftest.py` sets MLX to the CPU device at session start, and pytest uses `-m "not heavy"` by default. +The routine suite is bounded, self-contained, and representative. It includes +unit tests and miniature integration fixtures, while excluding tests marked +`heavy`, `real_assets`, or `torch_parity`. -Use tiny tensors in default tests. Prefer shapes like `shape=(1, 4, 32)` for behavioral checks, not full model-scale shapes like `shape=(1, 1024, 1024)`, unless the test is explicitly marked heavy. +```bash +uv run pytest +``` + +`tests/conftest.py` starts MLX on the CPU. A test may deliberately exercise +Metal when that boundary is important; such a test carries the `metal` marker. +Metal capability and test cost are separate concerns. + +## Markers + +| Marker | Meaning | Routine suite | +| --- | --- | --- | +| `integration` | Crosses multiple production component boundaries | Yes, unless paired with an exclusion marker | +| `metal` | Requires an Apple Metal device | Yes, when bounded | +| `real_assets` | Reads local weights, inputs, or uncommitted fixtures | No | +| `heavy` | Has substantial runtime or memory cost | No | +| `benchmark` | Measures resource or performance behavior; also marked `heavy` | No | +| `torch_parity` | Requires the opt-in PyTorch reference environment | No | + +Useful focused commands are: + +```bash +uv run pytest -m 'integration and not (heavy or real_assets or torch_parity)' +uv run pytest -m 'real_assets and not heavy' +uv run pytest -m heavy +uv run pytest -m benchmark +uv run pytest -m torch_parity +``` + +## Test Design + +- Prefer the smallest tensor shapes that preserve the production branch under + test. +- Keep one representative cross-component fixture per pipeline. Add narrower + tests only when they protect a distinct contract or failure mode. +- Keep each supported pipeline's primary guard in + `ROUTINE_PIPELINE_GUARDS`. A primary guard must remain an `integration` test + and must not carry a routine-exclusion marker. +- Assert stable stage, schema, shape, and numerical-summary contracts. Do not + accept a broad set of unrelated blockers as success. +- Record what a fixture covers and does not cover. Synthetic checkpoints must + not be presented as real-weight numerical parity. +- Write all generated files below `tmp_path` or the task scratch root. The + repository `outputs/` directory belongs to user-requested inference runs. + +The bounded pipeline fixtures cover TRELLIS.2 selective INT8 inference, +Pixal3D orchestration and derived decoder replay, LiTo source-contract +generation, SAM3D CLI reconstruction, HY-World 2 reconstruction, and the full +MapAnything scene path. They remain runnable after local model weights are +removed. + +## Isolated Runs -Mark tests with `@pytest.mark.heavy` when they load real files from `weights/`, require Metal-specific execution, or allocate model-scale tensors. Run them manually with: +Use a task-specific scratch root for any focused or expensive run: ```bash export MLX_SPATIAL_TEST_SCRATCH="$(mktemp -d /tmp/mlx-spatial-test.XXXXXX)" @@ -14,11 +67,10 @@ uv run pytest -m heavy \ --basetemp "$MLX_SPATIAL_TEST_SCRATCH/artifacts/pytest-heavy" ``` -All generated test inputs, outputs, parity bundles, caches, logs, and browser -artifacts must stay below that task root. Optional local reference checkouts use -explicit environment variables such as `MLX_SPATIAL_TORCH_ROOT`; committed -tests and anchor metadata must not contain developer-machine absolute paths. +Optional reference checkouts use explicit environment variables such as +`MLX_SPATIAL_TORCH_ROOT`. Committed fixtures and metadata must not contain +developer-machine absolute paths. -The GitHub Actions workflow runs the unified root suite, including -`tests/spatialkit`, on every branch push. The job creates the same isolated -scratch layout and has a 10-minute timeout so leaked heavy tests fail quickly. +GitHub Actions runs the routine root suite, including `tests/spatialkit`, on +pushes and pull requests. The job uses isolated scratch storage and a 10-minute +timeout. diff --git a/tests/data/mapanything_miniature_scene_golden.json b/tests/data/mapanything_miniature_scene_golden.json new file mode 100644 index 0000000..db1ceb4 --- /dev/null +++ b/tests/data/mapanything_miniature_scene_golden.json @@ -0,0 +1,193 @@ +{ + "artifact": { + "format": "npz", + "keys": [ + "__metadata_json__", + "camera_poses", + "confidence", + "depth", + "extrinsics", + "images", + "intrinsics", + "masks", + "world_points" + ] + }, + "fixture": { + "covered": [ + "asset inspection", + "safetensors loading", + "image preprocessing", + "full encoder", + "fusion norm", + "multi-view info sharing", + "dense pose and scale heads", + "scene geometry postprocess", + "NPZ artifact writing" + ], + "encoder_layers": 2, + "image_size": [ + 4, + 4 + ], + "info_sharing_layers": 2, + "kind": "generated-miniature-checkpoint", + "not_covered": [ + "numeric parity with official MapAnything weights", + "production image resolution", + "performance or memory benchmarking" + ], + "patch_size": 2, + "quantization": "none", + "random_seeds": { + "encoder": 42, + "heads": 123 + }, + "source": "deterministic synthetic tensors", + "views": 2 + }, + "geometry_invariants": { + "intrinsics": { + "dtype": "float32", + "finite": true, + "homogeneous_bottom_row": true, + "positive_focal_lengths": true, + "shape": [ + 2, + 3, + 3 + ] + }, + "world_points": { + "bounded": true, + "dtype": "float32", + "finite": true, + "nonzero": true, + "shape": [ + 2, + 4, + 4, + 3 + ] + } + }, + "predictions": { + "camera_poses": { + "dtype": "float32", + "shape": [ + 2, + 4, + 4 + ], + "statistics": { + "l2": 2.828591814609413, + "max": 1.0, + "mean": 0.00992451986735432, + "min": -0.8856602907180786, + "std": 0.4999306132682394 + } + }, + "confidence": { + "dtype": "float32", + "shape": [ + 2, + 4, + 4 + ], + "statistics": { + "l2": 11.313672927044621, + "max": 1.9999938011169434, + "mean": 1.9999937117099762, + "min": 1.9999936819076538, + "std": 5.1619136559035694e-08 + } + }, + "depth": { + "dtype": "float32", + "shape": [ + 2, + 4, + 4 + ], + "statistics": { + "l2": 5.630415995036871, + "max": 0.9953265190124512, + "mean": 0.9953263327479362, + "min": 0.9953261017799377, + "std": 1.199057470219663e-07 + } + }, + "extrinsics": { + "dtype": "float32", + "shape": [ + 2, + 4, + 4 + ], + "statistics": { + "l2": 2.8285918068254254, + "max": 1.0, + "mean": 0.012379770862025907, + "min": -0.8856602907180786, + "std": 0.4998758386586127 + } + }, + "images": { + "dtype": "float32", + "shape": [ + 2, + 4, + 4, + 3 + ], + "statistics": { + "l2": 6.344414571302413, + "max": 1.0, + "mean": 0.5000000035700699, + "min": 0.0, + "std": 0.4114455703905904 + } + }, + "masks": { + "dtype": "float32", + "shape": [ + 2, + 4, + 4 + ], + "statistics": { + "l2": 5.656854249492381, + "max": 1.0, + "mean": 1.0, + "min": 1.0, + "std": 0.0 + } + } + }, + "schema_version": 2, + "trace": { + "completed_stages": [ + "asset-config-validation", + "image-preprocessing", + "model-config", + "checkpoint-loading:encoder", + "full-encoder", + "checkpoint-loading:heads", + "fusion-norm", + "checkpoint-loading:info-sharing", + "info-sharing", + "prediction-heads", + "scene-postprocess" + ], + "frame_count": 2, + "implemented_boundary": "scene-generation", + "patch_grid": [ + 2, + 2 + ], + "target_size": [ + 4, + 4 + ] + } +} diff --git a/tests/data/pixal3d_derived_golden/golden.json b/tests/data/pixal3d_derived_golden/golden.json index 9a33eb9..7f64ef0 100644 --- a/tests/data/pixal3d_derived_golden/golden.json +++ b/tests/data/pixal3d_derived_golden/golden.json @@ -39,6 +39,7 @@ }, "expected_export": { "final_faces": 512, + "glb_vertex_tolerance": 8, "glb": { "faces": 512, "images": 2, @@ -73,7 +74,7 @@ } }, "fixture_kind": "real-weight-derived-decoder-patch", - "schema_version": 1, + "schema_version": 2, "scope": { "covered": "decoded O-Voxel contract through native textured GLB export", "not_covered": "Pixal3D checkpoint loading, conditioning, flow sampling, and decoder execution" diff --git a/tests/mapanything_scene_fixture.py b/tests/mapanything_scene_fixture.py new file mode 100644 index 0000000..9760f6b --- /dev/null +++ b/tests/mapanything_scene_fixture.py @@ -0,0 +1,451 @@ +"""Deterministic miniature assets for MapAnything scene-pipeline tests.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import mlx.core as mx +import numpy as np +from PIL import Image + +from mlx_spatial.mapanything_heads import MapAnythingHeadsConfig +from mlx_spatial.mapanything_model import ( + MapAnythingEncoderPrefixConfig, + MapAnythingInfoSharingConfig, +) +from tests.safetensors_test_utils import save_file + + +@dataclass(frozen=True) +class MapAnythingMiniatureSceneFixture: + """Paths and test-only config for a generated scene fixture.""" + + model_root: Path + image_root: Path + heads_config: MapAnythingHeadsConfig + + +def build_mapanything_miniature_scene_fixture( + root: Path, +) -> MapAnythingMiniatureSceneFixture: + """Generate a tiny, fully executable MapAnything checkpoint and two views.""" + + model_root = root / "model" + image_root = root / "images" + model_root.mkdir(parents=True) + image_root.mkdir(parents=True) + + encoder_config = tiny_encoder_config(layers=2) + info_config = tiny_info_config() + heads_config = tiny_heads_config(input_feature_dim=8) + weights = { + **tiny_encoder_weights(encoder_config), + **tiny_info_weights(info_config), + **tiny_heads_weights(heads_config), + } + (model_root / "config.json").write_text( + tiny_model_config_json(encoder_layers=2), + encoding="utf-8", + ) + save_file(weights, model_root / "model.safetensors") + + pixels = np.arange(4 * 4 * 3, dtype=np.uint8).reshape(4, 4, 3) + Image.fromarray(pixels, mode="RGB").save(image_root / "view-0.png") + Image.fromarray(255 - pixels, mode="RGB").save(image_root / "view-1.png") + return MapAnythingMiniatureSceneFixture( + model_root=model_root, + image_root=image_root, + heads_config=heads_config, + ) + + +def tiny_encoder_config(*, layers: int = 1) -> MapAnythingEncoderPrefixConfig: + return MapAnythingEncoderPrefixConfig( + embed_dim=8, + num_heads=2, + patch_size=2, + data_norm_type="dinov2", + encoder_size="giant", + keep_first_n_layers=layers, + ) + + +def tiny_info_config( + *, + depth: int = 2, + indices: tuple[int, ...] = (0, 1), + dim: int = 8, +) -> MapAnythingInfoSharingConfig: + return MapAnythingInfoSharingConfig( + input_embed_dim=dim, + dim=dim, + depth=depth, + num_heads=2, + indices=indices, + norm_intermediate=True, + ) + + +def tiny_heads_config(*, input_feature_dim: int = 4) -> MapAnythingHeadsConfig: + return MapAnythingHeadsConfig( + input_feature_dim=input_feature_dim, + patch_size=2, + layer_dims=(2, 2, 2, 2), + feature_dim=2, + dense_output_dim=6, + pose_resconv_blocks=2, + pose_rot_dim=4, + scale_hidden_dim=3, + scale_output_dim=1, + ) + + +def tiny_encoder_weights( + config: MapAnythingEncoderPrefixConfig | None = None, +) -> dict[str, mx.array]: + config = config or tiny_encoder_config() + hidden = config.swiglu_hidden_features + randn = _random_tensor_factory(seed=42) + weights: dict[str, mx.array] = { + "encoder.model.cls_token": randn((1, 1, config.embed_dim), 0.03), + "encoder.model.pos_embed": randn((1, 5, config.embed_dim), 0.02), + "encoder.model.patch_embed.proj.weight": randn( + (config.embed_dim, 3, config.patch_size, config.patch_size), 0.04 + ), + "encoder.model.patch_embed.proj.bias": randn((config.embed_dim,), 0.01), + } + for block_index in range(config.keep_first_n_layers): + prefix = f"encoder.model.blocks.{block_index}" + weights.update( + { + f"{prefix}.norm1.weight": randn((config.embed_dim,), 0.01, 1.0), + f"{prefix}.norm1.bias": randn((config.embed_dim,), 0.01), + f"{prefix}.attn.qkv.weight": randn( + (3 * config.embed_dim, config.embed_dim), 0.03 + ), + f"{prefix}.attn.qkv.bias": randn((3 * config.embed_dim,), 0.01), + f"{prefix}.attn.proj.weight": randn( + (config.embed_dim, config.embed_dim), 0.03 + ), + f"{prefix}.attn.proj.bias": randn((config.embed_dim,), 0.01), + f"{prefix}.ls1.gamma": randn((config.embed_dim,), 0.01, 0.1), + f"{prefix}.norm2.weight": randn((config.embed_dim,), 0.01, 1.0), + f"{prefix}.norm2.bias": randn((config.embed_dim,), 0.01), + f"{prefix}.mlp.w12.weight": randn( + (2 * hidden, config.embed_dim), 0.02 + ), + f"{prefix}.mlp.w12.bias": randn((2 * hidden,), 0.01), + f"{prefix}.mlp.w3.weight": randn( + (config.embed_dim, hidden), 0.02 + ), + f"{prefix}.mlp.w3.bias": randn((config.embed_dim,), 0.01), + f"{prefix}.ls2.gamma": randn((config.embed_dim,), 0.01, 0.1), + } + ) + return weights + + +def tiny_info_weights( + config: MapAnythingInfoSharingConfig, + *, + identity_odd_attention: bool = False, +) -> dict[str, mx.array]: + hidden = config.swiglu_hidden_features + weights: dict[str, mx.array] = { + "scale_token": mx.array( + np.linspace(-0.5, 0.6, config.dim, dtype=np.float32) + ), + "info_sharing.norm.weight": mx.ones((config.dim,), dtype=mx.float32), + "info_sharing.norm.bias": mx.zeros((config.dim,), dtype=mx.float32), + "info_sharing.view_pos_table": mx.zeros((1, config.dim), dtype=mx.float32), + } + for block_index in range(config.depth): + prefix = f"info_sharing.self_attention_blocks.{block_index}" + weights.update( + { + f"{prefix}.norm1.weight": mx.ones((config.dim,), dtype=mx.float32), + f"{prefix}.norm1.bias": mx.zeros((config.dim,), dtype=mx.float32), + f"{prefix}.attn.qkv.weight": mx.zeros( + (3 * config.dim, config.dim), dtype=mx.float32 + ), + f"{prefix}.attn.qkv.bias": mx.zeros( + (3 * config.dim,), dtype=mx.float32 + ), + f"{prefix}.attn.proj.weight": mx.zeros( + (config.dim, config.dim), dtype=mx.float32 + ), + f"{prefix}.attn.proj.bias": mx.zeros( + (config.dim,), dtype=mx.float32 + ), + f"{prefix}.ls1.gamma": mx.ones((config.dim,), dtype=mx.float32), + f"{prefix}.norm2.weight": mx.ones((config.dim,), dtype=mx.float32), + f"{prefix}.norm2.bias": mx.zeros((config.dim,), dtype=mx.float32), + f"{prefix}.mlp.w12.weight": mx.zeros( + (2 * hidden, config.dim), dtype=mx.float32 + ), + f"{prefix}.mlp.w12.bias": mx.zeros( + (2 * hidden,), dtype=mx.float32 + ), + f"{prefix}.mlp.w3.weight": mx.zeros( + (config.dim, hidden), dtype=mx.float32 + ), + f"{prefix}.mlp.w3.bias": mx.zeros((config.dim,), dtype=mx.float32), + f"{prefix}.ls2.gamma": mx.ones((config.dim,), dtype=mx.float32), + } + ) + if identity_odd_attention and config.depth > 1: + prefix = "info_sharing.self_attention_blocks.1" + identity = np.eye(config.dim, dtype=np.float32) + weights[f"{prefix}.attn.qkv.weight"] = mx.array( + np.concatenate((identity, identity, identity), axis=0) + ) + weights[f"{prefix}.attn.proj.weight"] = mx.array(identity) + return weights + + +def tiny_heads_weights(config: MapAnythingHeadsConfig) -> dict[str, mx.array]: + randn = _random_tensor_factory(seed=123) + weights: dict[str, mx.array] = { + "fusion_norm_layer.weight": mx.ones( + (config.input_feature_dim,), dtype=mx.float32 + ), + "fusion_norm_layer.bias": mx.zeros( + (config.input_feature_dim,), dtype=mx.float32 + ), + } + for index, channels in enumerate(config.layer_dims): + prefix = f"dense_head.0.input_process.{index}" + weights[f"{prefix}.0.0.weight"] = randn( + (channels, config.input_feature_dim, 1, 1) + ) + weights[f"{prefix}.0.0.bias"] = randn((channels,)) + weights[f"{prefix}.1.weight"] = randn( + (config.feature_dim, channels, 3, 3) + ) + for index, kernel_size in ((0, 4), (1, 2), (3, 3)): + channels = config.layer_dims[index] + prefix = f"dense_head.0.input_process.{index}.0.1" + weights[f"{prefix}.weight"] = randn( + (channels, channels, kernel_size, kernel_size) + ) + weights[f"{prefix}.bias"] = randn((channels,)) + + for block in ("refinenet1", "refinenet2", "refinenet3"): + for unit in ("resConfUnit1", "resConfUnit2"): + for conv in ("conv1", "conv2"): + prefix = f"dense_head.0.scratch.{block}.{unit}.{conv}" + weights[f"{prefix}.weight"] = randn( + (config.feature_dim, config.feature_dim, 3, 3) + ) + weights[f"{prefix}.bias"] = randn((config.feature_dim,)) + prefix = f"dense_head.0.scratch.{block}.out_conv" + weights[f"{prefix}.weight"] = randn( + (config.feature_dim, config.feature_dim, 1, 1) + ) + weights[f"{prefix}.bias"] = randn((config.feature_dim,)) + for conv in ("conv1", "conv2"): + prefix = f"dense_head.0.scratch.refinenet4.resConfUnit2.{conv}" + weights[f"{prefix}.weight"] = randn( + (config.feature_dim, config.feature_dim, 3, 3) + ) + weights[f"{prefix}.bias"] = randn((config.feature_dim,)) + weights["dense_head.0.scratch.refinenet4.out_conv.weight"] = randn( + (config.feature_dim, config.feature_dim, 1, 1) + ) + weights["dense_head.0.scratch.refinenet4.out_conv.bias"] = randn( + (config.feature_dim,) + ) + + final_channels = config.feature_dim // 2 + weights["dense_head.1.conv1.weight"] = randn( + (final_channels, config.feature_dim, 3, 3) + ) + weights["dense_head.1.conv1.bias"] = randn((final_channels,)) + weights["dense_head.1.conv2.0.weight"] = randn( + (final_channels, final_channels, 3, 3) + ) + weights["dense_head.1.conv2.0.bias"] = randn((final_channels,)) + weights["dense_head.1.conv2.2.weight"] = randn( + (config.dense_output_dim, final_channels, 1, 1) + ) + weights["dense_head.1.conv2.2.bias"] = mx.array( + [0.1, -0.1, 1.0, 0.0, 0.0, 2.0], + dtype=mx.float32, + ) + + pose_hidden = config.pose_hidden_dim + weights["pose_head.proj.weight"] = randn( + (pose_hidden, config.input_feature_dim, 1, 1) + ) + weights["pose_head.proj.bias"] = randn((pose_hidden,)) + for block_index in range(config.pose_resconv_blocks): + for conv in ("res_conv1", "res_conv2", "res_conv3"): + prefix = f"pose_head.res_conv.{block_index}.{conv}" + weights[f"{prefix}.weight"] = randn( + (pose_hidden, pose_hidden, 1, 1) + ) + weights[f"{prefix}.bias"] = randn((pose_hidden,)) + for layer in (0, 2): + weights[f"pose_head.more_mlps.{layer}.weight"] = randn( + (pose_hidden, pose_hidden) + ) + weights[f"pose_head.more_mlps.{layer}.bias"] = randn((pose_hidden,)) + weights["pose_head.fc_t.weight"] = randn((3, pose_hidden)) + weights["pose_head.fc_t.bias"] = randn((3,)) + weights["pose_head.fc_rot.weight"] = randn( + (config.pose_rot_dim, pose_hidden) + ) + weights["pose_head.fc_rot.bias"] = randn((config.pose_rot_dim,)) + + weights["scale_head.proj.weight"] = randn( + (config.scale_hidden_dim, config.input_feature_dim) + ) + weights["scale_head.proj.bias"] = randn((config.scale_hidden_dim,)) + for layer in (0, 1): + weights[f"scale_head.mlp.{layer}.0.weight"] = randn( + (config.scale_hidden_dim, config.scale_hidden_dim) + ) + weights[f"scale_head.mlp.{layer}.0.bias"] = randn( + (config.scale_hidden_dim,) + ) + weights["scale_head.output_proj.weight"] = randn( + (config.scale_output_dim, config.scale_hidden_dim) + ) + weights["scale_head.output_proj.bias"] = randn((config.scale_output_dim,)) + return weights + + +def _random_tensor_factory(seed: int): + rng = np.random.default_rng(seed) + + def randn( + shape: tuple[int, ...], + scale: float = 0.02, + offset: float = 0.0, + ) -> mx.array: + values = rng.normal(loc=offset, scale=scale, size=shape).astype(np.float32) + return mx.array(values) + + return randn + + +def tiny_features_and_registers( + config: MapAnythingInfoSharingConfig, +) -> tuple[tuple[mx.array, mx.array], tuple[mx.array, mx.array]]: + values = mx.arange(2 * config.dim * 2 * 2, dtype=mx.float32).reshape( + (2, config.dim, 2, 2) + ) / 100 + features = (values[0:1], values[1:2]) + registers = ( + mx.ones((1, config.dim, 1), dtype=mx.float32) * 0.1, + mx.ones((1, config.dim, 1), dtype=mx.float32) * -0.1, + ) + return features, registers + + +def tiny_prefix_asset_weights() -> dict[str, mx.array]: + config = tiny_encoder_config() + weights = tiny_encoder_weights(config) + weights.update( + { + "info_sharing.dummy": mx.zeros((1,), dtype=mx.float32), + "dense_head.dummy": mx.zeros((1,), dtype=mx.float32), + "pose_head.dummy": mx.zeros((1,), dtype=mx.float32), + "scale_head.dummy": mx.zeros((1,), dtype=mx.float32), + "fusion_norm_layer.weight": mx.ones( + (config.embed_dim,), dtype=mx.float32 + ), + "scale_token": mx.zeros((config.embed_dim,), dtype=mx.float32), + } + ) + return weights + + +def write_tiny_mapanything_prefix_fixture(root: Path) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "config.json").write_text( + tiny_model_config_json( + info_depth=1, + info_indices=(0,), + use_register_tokens=True, + ), + encoding="utf-8", + ) + save_file(tiny_prefix_asset_weights(), root / "model.safetensors") + return root + + +def tiny_model_config_json( + *, + encoder_layers: int = 1, + info_depth: int = 2, + info_dim: int = 8, + info_indices: tuple[int, ...] = (0, 1), + use_register_tokens: bool = True, +) -> str: + payload = _tiny_model_config( + encoder_layers=encoder_layers, + info_depth=info_depth, + info_dim=info_dim, + info_indices=info_indices, + use_register_tokens=use_register_tokens, + ) + return json.dumps(payload, indent=2) + "\n" + + +def _tiny_model_config( + *, + encoder_layers: int, + info_depth: int, + info_dim: int, + info_indices: tuple[int, ...], + use_register_tokens: bool, +) -> dict[str, object]: + return { + "encoder_config": { + "data_norm_type": "dinov2", + "name": "tiny-generated-fixture", + "size": "giant", + "keep_first_n_layers": encoder_layers, + "uses_torch_hub": False, + }, + "info_sharing_config": { + "model_type": "alternating_attention", + "model_return_type": "intermediate_features", + "module_args": { + "depth": info_depth, + "dim": info_dim, + "num_heads": 2, + "indices": list(info_indices), + }, + }, + "pred_head_config": { + "type": "dpt+pose", + "adaptor_type": "raydirs+depth+pose+confidence+mask", + "feature_head": {"patch_size": 2}, + "adaptor_config": { + "dense_pred_init_dict": { + "name": "raydirs+depth+pose+confidence+mask+scale" + } + }, + }, + "use_register_tokens_from_encoder": use_register_tokens, + } + + +__all__ = [ + "MapAnythingMiniatureSceneFixture", + "build_mapanything_miniature_scene_fixture", + "tiny_encoder_config", + "tiny_encoder_weights", + "tiny_features_and_registers", + "tiny_heads_config", + "tiny_heads_weights", + "tiny_info_config", + "tiny_info_weights", + "tiny_model_config_json", + "tiny_prefix_asset_weights", + "write_tiny_mapanything_prefix_fixture", +] diff --git a/tests/spatialkit/test_real_pixal3d_export.py b/tests/spatialkit/test_real_pixal3d_export.py index 32ed2a4..9872269 100644 --- a/tests/spatialkit/test_real_pixal3d_export.py +++ b/tests/spatialkit/test_real_pixal3d_export.py @@ -275,6 +275,8 @@ def test_pixal3d_run_manifest_supports_unmanifested_cached_output(tmp_path: Path @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_real_decoded_fixture_writes_glb_and_diagnostics(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -367,6 +369,8 @@ def test_export_pixal3d_glb_real_decoded_fixture_writes_glb_and_diagnostics(tmp_ @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_udf_remesh_closes_topology_and_matches_reference_mechanism(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -416,6 +420,8 @@ def test_export_pixal3d_glb_udf_remesh_closes_topology_and_matches_reference_mec @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_native_chart_backend_writes_real_fixture(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -533,6 +539,8 @@ def test_export_pixal3d_glb_native_chart_backend_writes_real_fixture(tmp_path: P @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_native_chart_violin_preprocessed_black_fixture(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -631,6 +639,8 @@ def test_export_pixal3d_glb_native_chart_violin_preprocessed_black_fixture(tmp_p @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_reference_target_preset_reports_thresholds(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -754,6 +764,8 @@ def test_export_pixal3d_glb_reference_target_preset_reports_thresholds(tmp_path: @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_reference_target_native_chart_backend_reports_readiness(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -964,6 +976,8 @@ def test_export_pixal3d_glb_reference_target_native_chart_backend_reports_readin @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_reference_target_4096_texture_reports_texture_resolution_gate(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -1013,6 +1027,8 @@ def test_export_pixal3d_glb_reference_target_4096_texture_reports_texture_resolu @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_upstream_settings_passes_readiness_gate(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -1084,6 +1100,8 @@ def test_export_pixal3d_glb_upstream_settings_passes_readiness_gate(tmp_path: Pa @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_native_chart_upstream_settings_passes_readiness_gate(tmp_path: Path) -> None: if not metal_device_available(): pytest.skip("Metal device unavailable for mlx-spatialkit real Pixal3D export") @@ -2144,6 +2162,8 @@ def test_glb_viewer_compatibility_summary_checks_normals_and_uint16_chunks() -> @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_qem_input_prep_real_fixture_manifold_and_bounded(tmp_path: Path) -> None: """S4 heavy: real fixture export with simplify_backend="qem" produces a fully manifold mesh. @@ -2508,6 +2528,8 @@ def _assert_qem_two_fixture_proof( @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_qem_two_fixture_main_manifold_and_beats_clustering(tmp_path: Path) -> None: """Slice-5 QEM proof: main fixture (pixal3d-1024-cascade-decoded-pbr) at res=256. @@ -2567,6 +2589,8 @@ def test_export_pixal3d_glb_qem_two_fixture_main_manifold_and_beats_clustering(t @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_pixal3d_glb_qem_two_fixture_violin_manifold_and_beats_clustering(tmp_path: Path) -> None: """Slice-5 QEM proof: violin-bow fixture at res=256. @@ -2709,6 +2733,8 @@ def _assert_chart_growth_parity_against_oracle(fixture: Path, anchor_name: str) @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_reference_uv_chart_growth_parity_main_fixture() -> None: """Slice-4 heavy: stage-B chart-count parity vs pip-xatlas oracle (main).""" if not metal_device_available(): @@ -2721,6 +2747,8 @@ def test_reference_uv_chart_growth_parity_main_fixture() -> None: @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_reference_uv_chart_growth_parity_violin_fixture() -> None: """Slice-4 heavy: stage-B chart-count parity vs pip-xatlas oracle (violin-bow).""" if not metal_device_available(): @@ -2857,6 +2885,8 @@ def _assert_parameterization_invariants_and_stretch_parity( @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_reference_uv_param_overlap_and_stretch_parity_main_fixture() -> None: """Slice-5 heavy: zero-overlap invariant + stretch parity (main).""" if not metal_device_available(): @@ -2869,6 +2899,8 @@ def test_reference_uv_param_overlap_and_stretch_parity_main_fixture() -> None: @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_reference_uv_param_overlap_and_stretch_parity_violin_fixture() -> None: """Slice-5 heavy: zero-overlap invariant + stretch parity (violin-bow).""" if not metal_device_available(): @@ -2907,6 +2939,8 @@ def test_uv_backend_validator_accepts_reference_backend() -> None: @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_reference_uv_packing_utilization_parity_violin_fixture() -> None: """Slice-6 heavy: full reference-backend atlas on the violin fixture. @@ -3159,6 +3193,8 @@ def _assert_reference_uv_e2e_proof( @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_reference_uv_e2e_proof_main_fixture(tmp_path: Path) -> None: """Slice-8 heavy: full reference-unwrap e2e proof (main fixture). @@ -3182,6 +3218,8 @@ def test_reference_uv_e2e_proof_main_fixture(tmp_path: Path) -> None: @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_reference_uv_e2e_proof_violin_fixture(tmp_path: Path) -> None: """Slice-8 heavy: full reference-unwrap e2e proof (violin-bow fixture). @@ -3208,6 +3246,8 @@ def test_reference_uv_e2e_proof_violin_fixture(tmp_path: Path) -> None: @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_inpaint_oracle_parity_against_cv2_telea() -> None: # TPP-02: our native Telea matches the cv2 INPAINT_TELEA oracle on the real # fixture inverse-coverage masks, within pinned tolerances (tight near the @@ -3275,6 +3315,8 @@ def test_inpaint_oracle_parity_against_cv2_telea() -> None: @pytest.mark.heavy +@pytest.mark.metal +@pytest.mark.real_assets def test_export_telea_postprocess_real_fixture_paints_gutter_without_black_seam(tmp_path: Path) -> None: # TPP-03: the reference Telea postprocess runs end-to-end on a real fixture, # painting the inverse-coverage gutter with no black-seam texels adjacent to @@ -3407,6 +3449,9 @@ def test_texture_postprocess_gate_passes_then_anti_gaming_flips_back() -> None: @pytest.mark.heavy +@pytest.mark.benchmark +@pytest.mark.metal +@pytest.mark.real_assets @pytest.mark.parametrize( "fixture_key,fixture_subpath", [ diff --git a/tests/test_gs_rasterize.py b/tests/test_gs_rasterize.py index d6e922e..89529a3 100644 --- a/tests/test_gs_rasterize.py +++ b/tests/test_gs_rasterize.py @@ -186,7 +186,7 @@ def test_anisotropic_scale_and_quaternion_rotate_screen_footprint(): assert np.asarray(vertical.rgba)[11, 8, 3] > np.asarray(vertical.rgba)[8, 11, 3] -@pytest.mark.heavy +@pytest.mark.metal def test_metal_matches_cpu_reference_for_tiny_image(): means = np.array( [ @@ -229,7 +229,7 @@ def test_metal_matches_cpu_reference_for_tiny_image(): np.testing.assert_allclose(np.asarray(metal.depth), np.asarray(cpu.depth), rtol=0.01, atol=0.01) -@pytest.mark.heavy +@pytest.mark.metal def test_metal_matches_cpu_reference_for_anisotropic_rotated_gaussian(): means = np.array([[0.0, 0.0, 2.0]], dtype=np.float32) angle = np.pi / 4.0 diff --git a/tests/test_hyworld2_inference.py b/tests/test_hyworld2_inference.py index 3602842..e6d622b 100644 --- a/tests/test_hyworld2_inference.py +++ b/tests/test_hyworld2_inference.py @@ -1,8 +1,8 @@ import json -import shutil from pathlib import Path import mlx.core as mx +import pytest from tests.safetensors_test_utils import save_file from mlx_spatial.hyworld2 import main @@ -72,9 +72,7 @@ def _write_fixture_images(path): def _output_dir(name): - path = Path("outputs") / "hyworld2" / name - shutil.rmtree(path, ignore_errors=True) - return path + return Path("outputs") / "hyworld2" / name def _assert_local_mlx_timing(metadata, completed_stages, *, blocker_stage, successful): @@ -304,7 +302,9 @@ def test_hyworld2_reconstruct_cli_rejects_output_outside_outputs(tmp_path, capsy assert "must stay under outputs" in output -def test_fixture_reconstruct_writes_staged_outputs_under_outputs(tmp_path): +@pytest.mark.integration +def test_fixture_reconstruct_writes_staged_outputs_under_outputs(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) _write_tiny_fixture_root(tmp_path) image_dir = tmp_path / "images" _write_fixture_images(image_dir) @@ -347,14 +347,13 @@ def test_fixture_reconstruct_writes_staged_outputs_under_outputs(tmp_path): ) -def test_fixture_reconstruct_writes_optional_mlx_parity_bundle(tmp_path): +def test_fixture_reconstruct_writes_optional_mlx_parity_bundle(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) _write_tiny_fixture_root(tmp_path) image_dir = tmp_path / "images" _write_fixture_images(image_dir) out = _output_dir("fixture-parity-bundle") parity_output = Path("outputs") / "hyworld2" / "fixture-parity-bundle.npz" - parity_output.unlink(missing_ok=True) - result = HyWorld2InferencePipeline(tmp_path).reconstruct( image_dir, output_path=out, @@ -372,7 +371,8 @@ def test_fixture_reconstruct_writes_optional_mlx_parity_bundle(tmp_path): assert "parity-mlx-bundle" in [output.name for output in result.trace.outputs] -def test_fixture_reconstruct_heads_depth_exports_only_depth(tmp_path): +def test_fixture_reconstruct_heads_depth_exports_only_depth(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) _write_tiny_fixture_root(tmp_path) image_dir = tmp_path / "images" _write_fixture_images(image_dir) @@ -399,7 +399,8 @@ def test_fixture_reconstruct_heads_depth_exports_only_depth(tmp_path): assert heads["camera"] == {"requested": False, "enabled": False, "export": False, "reason": "not requested"} -def test_fixture_reconstruct_cleans_stale_artifacts_when_heads_change(tmp_path): +def test_fixture_reconstruct_cleans_stale_artifacts_when_heads_change(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) _write_tiny_fixture_root(tmp_path) image_dir = tmp_path / "images" _write_fixture_images(image_dir) @@ -446,7 +447,8 @@ def test_fixture_reconstruct_cleans_stale_artifacts_when_heads_change(tmp_path): assert trace["metadata"]["heads"]["points"]["reason"] == "not requested" -def test_fixture_reconstruct_requested_gs_exports_gaussians_ply(tmp_path): +def test_fixture_reconstruct_requested_gs_exports_gaussians_ply(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) _write_tiny_fixture_root(tmp_path) image_dir = tmp_path / "images" _write_fixture_images(image_dir) @@ -491,7 +493,8 @@ def test_fixture_reconstruct_requested_gs_exports_gaussians_ply(tmp_path): assert (out / "trace.json").is_file() -def test_hyworld2_reconstruct_cli_fixture_tensors_depth_only(tmp_path, capsys): +def test_hyworld2_reconstruct_cli_fixture_tensors_depth_only(tmp_path, capsys, monkeypatch): + monkeypatch.chdir(tmp_path) _write_tiny_fixture_root(tmp_path) image_dir = tmp_path / "images" _write_fixture_images(image_dir) diff --git a/tests/test_lito_dit.py b/tests/test_lito_dit.py index b1f01bc..252ceaa 100644 --- a/tests/test_lito_dit.py +++ b/tests/test_lito_dit.py @@ -233,6 +233,8 @@ def test_dit_full_trajectory_matches_source_contract(caplog): @pytest.mark.heavy +@pytest.mark.benchmark +@pytest.mark.metal @pytest.mark.parametrize("profile", LITO_MEMORY_PROFILES) def test_dit_memory_profiles_stay_under_90gb(profile): _require_metal_memory_api() @@ -257,6 +259,8 @@ def test_dit_memory_profiles_stay_under_90gb(profile): @pytest.mark.heavy +@pytest.mark.benchmark +@pytest.mark.metal def test_dit_memory_safe_stays_well_under_threshold(): _require_metal_memory_api() tensors = _dit_input() diff --git a/tests/test_lito_inference.py b/tests/test_lito_inference.py index 5fcff13..d303671 100644 --- a/tests/test_lito_inference.py +++ b/tests/test_lito_inference.py @@ -25,6 +25,7 @@ ROOT = Path(__file__).resolve().parents[1] +@pytest.mark.integration def test_full_pipeline_runs_on_sample_input(tmp_path): image = _write_synthetic_image(tmp_path / "input.png") output = tmp_path / "result.ply" @@ -218,6 +219,7 @@ def _write_trellis_runtime_assets(root: Path) -> None: @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif(not (ROOT / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent") def test_generate_with_real_weight_headers_does_not_fall_back_to_smoke_on_backend_failure(tmp_path): image = _write_synthetic_image(tmp_path / "input.png") diff --git a/tests/test_lito_real_backend.py b/tests/test_lito_real_backend.py index 235726d..56ff0db 100644 --- a/tests/test_lito_real_backend.py +++ b/tests/test_lito_real_backend.py @@ -87,6 +87,7 @@ def test_inspect_lito_real_architecture_uses_safetensor_headers(tmp_path): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -545,6 +546,7 @@ def test_run_lito_voxel_decoder_lowres_latent_runs_fake_weights(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -583,6 +585,7 @@ def test_real_gaussian_output_heads_run_from_loaded_checkpoint_weights(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/image_to_3d/lito_dit_rgba.safetensors").is_file(), reason="LiTo weights absent", @@ -607,6 +610,7 @@ def test_real_dit_velocity_block0_runs_from_loaded_checkpoint_weights(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/image_to_3d/lito_dit_rgba.safetensors").is_file(), reason="LiTo weights absent", @@ -632,6 +636,7 @@ def test_real_dit_sampler_block0_runs_from_loaded_checkpoint_weights(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -653,6 +658,7 @@ def test_real_gaussian_query_point_stem_runs_from_loaded_checkpoint_weights(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -682,6 +688,7 @@ def test_real_gaussian_perceiver_block0_cross_only_runs_from_loaded_checkpoint_w @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -712,6 +719,7 @@ def test_real_gaussian_perceiver_block0_local_voxel_self_attention_runs_from_loa @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -742,6 +750,7 @@ def test_real_gaussian_perceiver_all_blocks_local_voxel_self_attention_runs_from @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -763,6 +772,7 @@ def test_real_voxel_decoder_lowres_latent_runs_from_loaded_checkpoint_weights(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not ( Path(__file__).resolve().parents[1] @@ -792,6 +802,7 @@ def test_real_trellis_sparse_structure_decoder_logits_run_from_local_mlx_weights @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not ( Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors" @@ -827,6 +838,7 @@ def test_real_init_coord_generation_from_latents_runs_with_local_mlx_weights(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -858,6 +870,7 @@ def test_load_lito_gaussian_decoder_weight_arrays_reads_real_safetensors_subset( @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -905,6 +918,7 @@ def test_load_lito_gaussian_decoder_weight_arrays_reads_real_cross_attention_sub @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/tokenizer/lito_new.safetensors").is_file(), reason="LiTo weights absent", @@ -936,6 +950,7 @@ def test_load_lito_voxel_decoder_weight_arrays_reads_real_safetensors_subset(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/image_to_3d/lito_dit_rgba.safetensors").is_file(), reason="LiTo weights absent", @@ -963,6 +978,7 @@ def test_load_lito_dit_weight_arrays_reads_real_safetensors_subset(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/image_to_3d/lito_dit_rgba.safetensors").is_file(), reason="LiTo weights absent", @@ -990,6 +1006,7 @@ def test_load_lito_patch_encoder_weight_arrays_reads_real_safetensors_subset(): @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif( not (Path(__file__).resolve().parents[1] / "weights/lito-research-mlx/image_to_3d/lito_dit_rgba.safetensors").is_file(), reason="LiTo weights absent", diff --git a/tests/test_mapanything_assets.py b/tests/test_mapanything_assets.py index bc93206..b8f74dc 100644 --- a/tests/test_mapanything_assets.py +++ b/tests/test_mapanything_assets.py @@ -184,6 +184,7 @@ def test_inspect_mapanything_model_assets_returns_blocker_for_corrupt_safetensor assert inspection.blocker.metadata["checkpoint"] == str(tmp_path / "model.safetensors") +@pytest.mark.real_assets def test_local_mapanything_checkpoint_layout_is_recognized_when_present(): root = Path(MAPANYTHING_DEFAULT_ROOT) if not root.is_dir(): diff --git a/tests/test_mapanything_full_encoder_parity.py b/tests/test_mapanything_full_encoder_parity.py index e13877c..b9cb4db 100644 --- a/tests/test_mapanything_full_encoder_parity.py +++ b/tests/test_mapanything_full_encoder_parity.py @@ -21,10 +21,13 @@ ROOT = Path(__file__).resolve().parents[1] -pytestmark = pytest.mark.skipif( - os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", - reason="opt-in MapAnything Torch reference parity", -) +pytestmark = [ + pytest.mark.torch_parity, + pytest.mark.skipif( + os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", + reason="opt-in MapAnything Torch reference parity", + ), +] def test_mapanything_full_encoder_matches_desk_scene_reference(): diff --git a/tests/test_mapanything_heads.py b/tests/test_mapanything_heads.py index 1793092..6433acd 100644 --- a/tests/test_mapanything_heads.py +++ b/tests/test_mapanything_heads.py @@ -23,6 +23,11 @@ run_mapanything_heads, validate_mapanything_heads_weights, ) +from tests.mapanything_scene_fixture import ( + tiny_heads_config as _tiny_heads_config, + tiny_heads_weights as _tiny_heads_weights, + tiny_model_config_json, +) def test_mapanything_heads_required_keys_are_explicit_and_unique(): @@ -129,89 +134,6 @@ def test_mapanything_heads_public_exports(): assert mlx_spatial.mapanything_heads_config_from_model_config is mapanything_heads_config_from_model_config -def _tiny_heads_config() -> MapAnythingHeadsConfig: - return MapAnythingHeadsConfig( - input_feature_dim=4, - patch_size=2, - layer_dims=(2, 2, 2, 2), - feature_dim=2, - dense_output_dim=6, - pose_resconv_blocks=2, - pose_rot_dim=4, - scale_hidden_dim=3, - scale_output_dim=1, - ) - - -def _tiny_heads_weights(config: MapAnythingHeadsConfig) -> dict[str, mx.array]: - rng = np.random.default_rng(123) - - def randn(shape: tuple[int, ...], scale: float = 0.02, offset: float = 0.0) -> mx.array: - return mx.array(rng.normal(loc=offset, scale=scale, size=shape).astype(np.float32)) - - weights: dict[str, mx.array] = { - "fusion_norm_layer.weight": mx.ones((config.input_feature_dim,), dtype=mx.float32), - "fusion_norm_layer.bias": mx.zeros((config.input_feature_dim,), dtype=mx.float32), - } - for index, channels in enumerate(config.layer_dims): - prefix = f"dense_head.0.input_process.{index}" - weights[f"{prefix}.0.0.weight"] = randn((channels, config.input_feature_dim, 1, 1)) - weights[f"{prefix}.0.0.bias"] = randn((channels,)) - weights[f"{prefix}.1.weight"] = randn((config.feature_dim, channels, 3, 3)) - weights["dense_head.0.input_process.0.0.1.weight"] = randn((2, 2, 4, 4)) - weights["dense_head.0.input_process.0.0.1.bias"] = randn((2,)) - weights["dense_head.0.input_process.1.0.1.weight"] = randn((2, 2, 2, 2)) - weights["dense_head.0.input_process.1.0.1.bias"] = randn((2,)) - weights["dense_head.0.input_process.3.0.1.weight"] = randn((2, 2, 3, 3)) - weights["dense_head.0.input_process.3.0.1.bias"] = randn((2,)) - - for block in ("refinenet1", "refinenet2", "refinenet3"): - for unit in ("resConfUnit1", "resConfUnit2"): - for conv in ("conv1", "conv2"): - weights[f"dense_head.0.scratch.{block}.{unit}.{conv}.weight"] = randn((2, 2, 3, 3)) - weights[f"dense_head.0.scratch.{block}.{unit}.{conv}.bias"] = randn((2,)) - weights[f"dense_head.0.scratch.{block}.out_conv.weight"] = randn((2, 2, 1, 1)) - weights[f"dense_head.0.scratch.{block}.out_conv.bias"] = randn((2,)) - for conv in ("conv1", "conv2"): - weights[f"dense_head.0.scratch.refinenet4.resConfUnit2.{conv}.weight"] = randn((2, 2, 3, 3)) - weights[f"dense_head.0.scratch.refinenet4.resConfUnit2.{conv}.bias"] = randn((2,)) - weights["dense_head.0.scratch.refinenet4.out_conv.weight"] = randn((2, 2, 1, 1)) - weights["dense_head.0.scratch.refinenet4.out_conv.bias"] = randn((2,)) - - weights["dense_head.1.conv1.weight"] = randn((1, 2, 3, 3)) - weights["dense_head.1.conv1.bias"] = randn((1,)) - weights["dense_head.1.conv2.0.weight"] = randn((1, 1, 3, 3)) - weights["dense_head.1.conv2.0.bias"] = randn((1,)) - weights["dense_head.1.conv2.2.weight"] = randn((config.dense_output_dim, 1, 1, 1)) - weights["dense_head.1.conv2.2.bias"] = randn((config.dense_output_dim,)) - - pose_hidden = config.pose_hidden_dim - weights["pose_head.proj.weight"] = randn((pose_hidden, config.input_feature_dim, 1, 1)) - weights["pose_head.proj.bias"] = randn((pose_hidden,)) - for block_index in range(config.pose_resconv_blocks): - for conv in ("res_conv1", "res_conv2", "res_conv3"): - weights[f"pose_head.res_conv.{block_index}.{conv}.weight"] = randn((pose_hidden, pose_hidden, 1, 1)) - weights[f"pose_head.res_conv.{block_index}.{conv}.bias"] = randn((pose_hidden,)) - weights["pose_head.more_mlps.0.weight"] = randn((pose_hidden, pose_hidden)) - weights["pose_head.more_mlps.0.bias"] = randn((pose_hidden,)) - weights["pose_head.more_mlps.2.weight"] = randn((pose_hidden, pose_hidden)) - weights["pose_head.more_mlps.2.bias"] = randn((pose_hidden,)) - weights["pose_head.fc_t.weight"] = randn((3, pose_hidden)) - weights["pose_head.fc_t.bias"] = randn((3,)) - weights["pose_head.fc_rot.weight"] = randn((config.pose_rot_dim, pose_hidden)) - weights["pose_head.fc_rot.bias"] = randn((config.pose_rot_dim,)) - - weights["scale_head.proj.weight"] = randn((config.scale_hidden_dim, config.input_feature_dim)) - weights["scale_head.proj.bias"] = randn((config.scale_hidden_dim,)) - weights["scale_head.mlp.0.0.weight"] = randn((config.scale_hidden_dim, config.scale_hidden_dim)) - weights["scale_head.mlp.0.0.bias"] = randn((config.scale_hidden_dim,)) - weights["scale_head.mlp.1.0.weight"] = randn((config.scale_hidden_dim, config.scale_hidden_dim)) - weights["scale_head.mlp.1.0.bias"] = randn((config.scale_hidden_dim,)) - weights["scale_head.output_proj.weight"] = randn((config.scale_output_dim, config.scale_hidden_dim)) - weights["scale_head.output_proj.bias"] = randn((config.scale_output_dim,)) - return weights - - def _numpy_nchw_layer_norm(values: np.ndarray, eps: float = 1e-6) -> np.ndarray: nhwc = np.transpose(values, (0, 2, 3, 1)) mean = nhwc.mean(axis=-1, keepdims=True) @@ -221,31 +143,4 @@ def _numpy_nchw_layer_norm(values: np.ndarray, eps: float = 1e-6) -> np.ndarray: def _tiny_config_json() -> str: - return """{ - "encoder_config": { - "data_norm_type": "dinov2", - "name": "tiny-test", - "size": "giant", - "keep_first_n_layers": 1, - "uses_torch_hub": false - }, - "info_sharing_config": { - "model_type": "alternating_attention", - "model_return_type": "intermediate_features", - "module_args": { - "depth": 2, - "dim": 4, - "num_heads": 2, - "indices": [0, 1] - } - }, - "pred_head_config": { - "type": "dpt+pose", - "adaptor_type": "raydirs+depth+pose+confidence+mask", - "feature_head": {"patch_size": 2}, - "adaptor_config": { - "dense_pred_init_dict": {"name": "raydirs+depth+pose+confidence+mask+scale"} - } - }, - "use_register_tokens_from_encoder": true -}""" + return tiny_model_config_json(info_dim=4) diff --git a/tests/test_mapanything_heads_parity.py b/tests/test_mapanything_heads_parity.py index 9ee5ad9..ada582e 100644 --- a/tests/test_mapanything_heads_parity.py +++ b/tests/test_mapanything_heads_parity.py @@ -21,10 +21,13 @@ ROOT = Path(__file__).resolve().parents[1] -pytestmark = pytest.mark.skipif( - os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", - reason="opt-in MapAnything Torch reference parity", -) +pytestmark = [ + pytest.mark.torch_parity, + pytest.mark.skipif( + os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", + reason="opt-in MapAnything Torch reference parity", + ), +] def test_mapanything_fusion_and_heads_match_desk_scene_reference(): diff --git a/tests/test_mapanything_inference.py b/tests/test_mapanything_inference.py index a3bb769..0a32d20 100644 --- a/tests/test_mapanything_inference.py +++ b/tests/test_mapanything_inference.py @@ -4,11 +4,9 @@ import tomllib from pathlib import Path -import mlx.core as mx import numpy as np import pytest from PIL import Image -from tests.safetensors_test_utils import save_file import mlx_spatial from mlx_spatial.mapanything_inference import ( @@ -22,6 +20,9 @@ MapAnythingScenePredictions, write_mapanything_scene_npz, ) +from tests.mapanything_scene_fixture import ( + write_tiny_mapanything_prefix_fixture as _write_tiny_model_root, +) ROOT = Path(__file__).resolve().parents[1] @@ -135,6 +136,7 @@ def test_mapanything_scene_prediction_bundle_schema(tmp_path): assert "case" in str(data["__metadata_json__"]) +@pytest.mark.real_assets def test_mapanything_prefix_pipeline_runs_local_desk_when_assets_present(): model_root = ROOT / "weights/map-anything" image_root = ROOT / "inputs/map-anything/desk" @@ -241,75 +243,3 @@ def test_mapanything_prefix_pipeline_public_exports(): assert mlx_spatial.MAPANYTHING_PREFIX_PARITY_ATOL == MAPANYTHING_PREFIX_PARITY_ATOL assert mlx_spatial.MapAnythingScenePipeline is MapAnythingScenePipeline assert mlx_spatial.MAPANYTHING_SCENE_OUTPUT_KEYS == MAPANYTHING_SCENE_OUTPUT_KEYS - - -def _write_tiny_model_root(root: Path) -> None: - root.mkdir(parents=True, exist_ok=True) - (root / "config.json").write_text(_tiny_config_json(), encoding="utf-8") - save_file(_tiny_encoder_prefix_weights(), root / "model.safetensors") - - -def _tiny_encoder_prefix_weights() -> dict[str, mx.array]: - embed_dim = 8 - patch_size = 2 - hidden = 24 - zeros = lambda shape: mx.zeros(shape, dtype=mx.float32) - ones = lambda shape: mx.ones(shape, dtype=mx.float32) - return { - "encoder.model.cls_token": zeros((1, 1, embed_dim)), - "encoder.model.pos_embed": zeros((1, 5, embed_dim)), - "encoder.model.patch_embed.proj.weight": zeros((embed_dim, 3, patch_size, patch_size)), - "encoder.model.patch_embed.proj.bias": zeros((embed_dim,)), - "encoder.model.blocks.0.norm1.weight": ones((embed_dim,)), - "encoder.model.blocks.0.norm1.bias": zeros((embed_dim,)), - "encoder.model.blocks.0.attn.qkv.weight": zeros((3 * embed_dim, embed_dim)), - "encoder.model.blocks.0.attn.qkv.bias": zeros((3 * embed_dim,)), - "encoder.model.blocks.0.attn.proj.weight": zeros((embed_dim, embed_dim)), - "encoder.model.blocks.0.attn.proj.bias": zeros((embed_dim,)), - "encoder.model.blocks.0.ls1.gamma": ones((embed_dim,)), - "encoder.model.blocks.0.norm2.weight": ones((embed_dim,)), - "encoder.model.blocks.0.norm2.bias": zeros((embed_dim,)), - "encoder.model.blocks.0.mlp.w12.weight": zeros((2 * hidden, embed_dim)), - "encoder.model.blocks.0.mlp.w12.bias": zeros((2 * hidden,)), - "encoder.model.blocks.0.mlp.w3.weight": zeros((embed_dim, hidden)), - "encoder.model.blocks.0.mlp.w3.bias": zeros((embed_dim,)), - "encoder.model.blocks.0.ls2.gamma": ones((embed_dim,)), - "info_sharing.dummy": zeros((1,)), - "dense_head.dummy": zeros((1,)), - "pose_head.dummy": zeros((1,)), - "scale_head.dummy": zeros((1,)), - "fusion_norm_layer.weight": ones((embed_dim,)), - "scale_token": zeros((embed_dim,)), - } - - -def _tiny_config_json() -> str: - return """{ - "encoder_config": { - "data_norm_type": "dinov2", - "name": "tiny-test", - "size": "giant", - "keep_first_n_layers": 1, - "uses_torch_hub": false, - "with_registers": false - }, - "info_sharing_config": { - "model_type": "alternating_attention", - "model_return_type": "intermediate_features", - "module_args": { - "depth": 1, - "dim": 8, - "num_heads": 2, - "indices": [0] - } - }, - "pred_head_config": { - "type": "dpt+pose", - "adaptor_type": "raydirs+depth+pose+confidence+mask", - "feature_head": {"patch_size": 2}, - "adaptor_config": { - "dense_pred_init_dict": {"name": "raydirs+depth+pose+confidence+mask+scale"} - } - }, - "use_register_tokens_from_encoder": true -}""" diff --git a/tests/test_mapanything_info_sharing.py b/tests/test_mapanything_info_sharing.py index f927f37..87d0528 100644 --- a/tests/test_mapanything_info_sharing.py +++ b/tests/test_mapanything_info_sharing.py @@ -18,6 +18,12 @@ run_mapanything_info_sharing, validate_mapanything_info_sharing_weights, ) +from tests.mapanything_scene_fixture import ( + tiny_features_and_registers as _tiny_features_and_registers, + tiny_info_config as _tiny_info_config, + tiny_info_weights as _tiny_info_weights, + tiny_model_config_json, +) def test_mapanything_info_sharing_required_keys_cover_configured_layers(): @@ -136,66 +142,6 @@ def test_mapanything_info_sharing_public_exports(): ) -def _tiny_info_config(depth: int = 2, indices: tuple[int, ...] = (0, 1)) -> MapAnythingInfoSharingConfig: - return MapAnythingInfoSharingConfig( - input_embed_dim=8, - dim=8, - depth=depth, - num_heads=2, - indices=indices, - norm_intermediate=True, - ) - - -def _tiny_info_weights( - config: MapAnythingInfoSharingConfig, - *, - identity_odd_attention: bool = False, -) -> dict[str, mx.array]: - hidden = config.swiglu_hidden_features - weights: dict[str, mx.array] = { - "scale_token": mx.array(np.linspace(-0.5, 0.6, config.dim, dtype=np.float32)), - "info_sharing.norm.weight": mx.ones((config.dim,), dtype=mx.float32), - "info_sharing.norm.bias": mx.zeros((config.dim,), dtype=mx.float32), - "info_sharing.view_pos_table": mx.zeros((1, config.dim), dtype=mx.float32), - } - for block_index in range(config.depth): - prefix = f"info_sharing.self_attention_blocks.{block_index}" - weights[f"{prefix}.norm1.weight"] = mx.ones((config.dim,), dtype=mx.float32) - weights[f"{prefix}.norm1.bias"] = mx.zeros((config.dim,), dtype=mx.float32) - weights[f"{prefix}.attn.qkv.weight"] = mx.zeros((3 * config.dim, config.dim), dtype=mx.float32) - weights[f"{prefix}.attn.qkv.bias"] = mx.zeros((3 * config.dim,), dtype=mx.float32) - weights[f"{prefix}.attn.proj.weight"] = mx.zeros((config.dim, config.dim), dtype=mx.float32) - weights[f"{prefix}.attn.proj.bias"] = mx.zeros((config.dim,), dtype=mx.float32) - weights[f"{prefix}.ls1.gamma"] = mx.ones((config.dim,), dtype=mx.float32) - weights[f"{prefix}.norm2.weight"] = mx.ones((config.dim,), dtype=mx.float32) - weights[f"{prefix}.norm2.bias"] = mx.zeros((config.dim,), dtype=mx.float32) - weights[f"{prefix}.mlp.w12.weight"] = mx.zeros((2 * hidden, config.dim), dtype=mx.float32) - weights[f"{prefix}.mlp.w12.bias"] = mx.zeros((2 * hidden,), dtype=mx.float32) - weights[f"{prefix}.mlp.w3.weight"] = mx.zeros((config.dim, hidden), dtype=mx.float32) - weights[f"{prefix}.mlp.w3.bias"] = mx.zeros((config.dim,), dtype=mx.float32) - weights[f"{prefix}.ls2.gamma"] = mx.ones((config.dim,), dtype=mx.float32) - - if identity_odd_attention and config.depth > 1: - prefix = "info_sharing.self_attention_blocks.1" - identity = np.eye(config.dim, dtype=np.float32) - weights[f"{prefix}.attn.qkv.weight"] = mx.array(np.concatenate((identity, identity, identity), axis=0)) - weights[f"{prefix}.attn.proj.weight"] = mx.array(identity) - return weights - - -def _tiny_features_and_registers( - config: MapAnythingInfoSharingConfig, -) -> tuple[tuple[mx.array, mx.array], tuple[mx.array, mx.array]]: - values = mx.arange(2 * config.dim * 2 * 2, dtype=mx.float32).reshape((2, config.dim, 2, 2)) / 100 - features = (values[0:1], values[1:2]) - registers = ( - mx.ones((1, config.dim, 1), dtype=mx.float32) * 0.1, - mx.ones((1, config.dim, 1), dtype=mx.float32) * -0.1, - ) - return features, registers - - def _numpy_channel_layer_norm(values: np.ndarray, eps: float = 1e-6) -> np.ndarray: transposed = np.transpose(values, (0, 2, 1)) mean = transposed.mean(axis=-1, keepdims=True) @@ -205,31 +151,4 @@ def _numpy_channel_layer_norm(values: np.ndarray, eps: float = 1e-6) -> np.ndarr def _tiny_config_json() -> str: - return """{ - "encoder_config": { - "data_norm_type": "dinov2", - "name": "tiny-test", - "size": "giant", - "keep_first_n_layers": 1, - "uses_torch_hub": false - }, - "info_sharing_config": { - "model_type": "alternating_attention", - "model_return_type": "intermediate_features", - "module_args": { - "depth": 2, - "dim": 8, - "num_heads": 2, - "indices": [0, 1] - } - }, - "pred_head_config": { - "type": "dpt+pose", - "adaptor_type": "raydirs+depth+pose+confidence+mask", - "feature_head": {"patch_size": 2}, - "adaptor_config": { - "dense_pred_init_dict": {"name": "raydirs+depth+pose+confidence+mask+scale"} - } - }, - "use_register_tokens_from_encoder": true -}""" + return tiny_model_config_json() diff --git a/tests/test_mapanything_info_sharing_parity.py b/tests/test_mapanything_info_sharing_parity.py index e224689..800d2f9 100644 --- a/tests/test_mapanything_info_sharing_parity.py +++ b/tests/test_mapanything_info_sharing_parity.py @@ -20,10 +20,13 @@ ROOT = Path(__file__).resolve().parents[1] -pytestmark = pytest.mark.skipif( - os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", - reason="opt-in MapAnything Torch reference parity", -) +pytestmark = [ + pytest.mark.torch_parity, + pytest.mark.skipif( + os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", + reason="opt-in MapAnything Torch reference parity", + ), +] def test_mapanything_info_sharing_matches_desk_scene_reference(): diff --git a/tests/test_mapanything_model.py b/tests/test_mapanything_model.py index ddabffd..4a79c76 100644 --- a/tests/test_mapanything_model.py +++ b/tests/test_mapanything_model.py @@ -29,6 +29,11 @@ load_mapanything_parity_bundle, mapanything_parity_report_to_dict, ) +from tests.mapanything_scene_fixture import ( + tiny_encoder_config, + tiny_encoder_weights, + tiny_model_config_json, +) ROOT = Path(__file__).resolve().parents[1] @@ -223,98 +228,21 @@ def test_mapanything_encoder_prefix_public_exports(): ) -def _tiny_config() -> MapAnythingEncoderPrefixConfig: - return MapAnythingEncoderPrefixConfig( - embed_dim=8, - num_heads=2, - patch_size=2, - data_norm_type="dinov2", - encoder_size="giant", - keep_first_n_layers=1, - ) +_tiny_config = tiny_encoder_config +_tiny_encoder_prefix_weights = tiny_encoder_weights def _tiny_full_config() -> MapAnythingEncoderPrefixConfig: - return MapAnythingEncoderPrefixConfig( - embed_dim=8, - num_heads=2, - patch_size=2, - data_norm_type="dinov2", - encoder_size="giant", - keep_first_n_layers=2, - ) - - -def _tiny_encoder_prefix_weights() -> dict[str, mx.array]: - config = _tiny_config() - hidden = config.swiglu_hidden_features - rng = np.random.default_rng(42) - - def randn(shape: tuple[int, ...], scale: float = 0.02, offset: float = 0.0) -> mx.array: - return mx.array(rng.normal(loc=offset, scale=scale, size=shape).astype(np.float32)) - - weights = { - "encoder.model.cls_token": randn((1, 1, config.embed_dim), 0.03), - "encoder.model.pos_embed": randn((1, 5, config.embed_dim), 0.02), - "encoder.model.patch_embed.proj.weight": randn( - (config.embed_dim, 3, config.patch_size, config.patch_size), - 0.04, - ), - "encoder.model.patch_embed.proj.bias": randn((config.embed_dim,), 0.01), - "encoder.model.blocks.0.norm1.weight": randn((config.embed_dim,), 0.01, 1.0), - "encoder.model.blocks.0.norm1.bias": randn((config.embed_dim,), 0.01), - "encoder.model.blocks.0.attn.qkv.weight": randn((3 * config.embed_dim, config.embed_dim), 0.03), - "encoder.model.blocks.0.attn.qkv.bias": randn((3 * config.embed_dim,), 0.01), - "encoder.model.blocks.0.attn.proj.weight": randn((config.embed_dim, config.embed_dim), 0.03), - "encoder.model.blocks.0.attn.proj.bias": randn((config.embed_dim,), 0.01), - "encoder.model.blocks.0.ls1.gamma": randn((config.embed_dim,), 0.01, 0.1), - "encoder.model.blocks.0.norm2.weight": randn((config.embed_dim,), 0.01, 1.0), - "encoder.model.blocks.0.norm2.bias": randn((config.embed_dim,), 0.01), - "encoder.model.blocks.0.mlp.w12.weight": randn((2 * hidden, config.embed_dim), 0.02), - "encoder.model.blocks.0.mlp.w12.bias": randn((2 * hidden,), 0.01), - "encoder.model.blocks.0.mlp.w3.weight": randn((config.embed_dim, hidden), 0.02), - "encoder.model.blocks.0.mlp.w3.bias": randn((config.embed_dim,), 0.01), - "encoder.model.blocks.0.ls2.gamma": randn((config.embed_dim,), 0.01, 0.1), - } - return weights + return tiny_encoder_config(layers=2) def _tiny_full_encoder_weights() -> dict[str, mx.array]: - prefix = _tiny_encoder_prefix_weights() - full = dict(prefix) - for key, value in prefix.items(): - block1_key = key.replace("encoder.model.blocks.0.", "encoder.model.blocks.1.") - if block1_key != key: - full[block1_key] = value - return full + return tiny_encoder_weights(_tiny_full_config()) def _tiny_config_json() -> str: - return """{ - "encoder_config": { - "data_norm_type": "dinov2", - "name": "tiny-test", - "size": "giant", - "keep_first_n_layers": 1, - "uses_torch_hub": false - }, - "info_sharing_config": { - "model_type": "alternating_attention", - "model_return_type": "intermediate_features", - "module_args": { - "depth": 1, - "dim": 8, - "num_heads": 2, - "indices": [0] - } - }, - "pred_head_config": { - "type": "dpt+pose", - "adaptor_type": "raydirs+depth+pose+confidence+mask", - "feature_head": {"patch_size": 2}, - "adaptor_config": { - "dense_pred_init_dict": {"name": "raydirs+depth+pose+confidence+mask+scale"} - } - }, - "use_register_tokens_from_encoder": false -}""" + return tiny_model_config_json( + info_depth=1, + info_indices=(0,), + use_register_tokens=False, + ) diff --git a/tests/test_mapanything_scene_pipeline.py b/tests/test_mapanything_scene_pipeline.py index 8163736..d50f984 100644 --- a/tests/test_mapanything_scene_pipeline.py +++ b/tests/test_mapanything_scene_pipeline.py @@ -1,18 +1,152 @@ +import json from pathlib import Path import numpy as np import pytest +import mlx_spatial.mapanything_scene as mapanything_scene from mlx_spatial.mapanything_scene import ( MAPANYTHING_SCENE_OUTPUT_KEYS, MapAnythingScenePipeline, + MapAnythingSceneResult, write_mapanything_scene_npz, ) +from tests.golden_assertions import assert_golden_close, summarize_array +from tests.mapanything_scene_fixture import ( + build_mapanything_miniature_scene_fixture, +) ROOT = Path(__file__).resolve().parents[1] +GOLDEN_MANIFEST = ROOT / "tests/data/mapanything_miniature_scene_golden.json" + + +@pytest.mark.integration +def test_mapanything_miniature_scene_pipeline_matches_golden( + tmp_path, + monkeypatch, +): + fixture = build_mapanything_miniature_scene_fixture(tmp_path) + monkeypatch.setattr( + mapanything_scene, + "mapanything_heads_config_from_model_config", + lambda _: fixture.heads_config, + ) + + result = MapAnythingScenePipeline(fixture.model_root).generate( + fixture.image_root, + resize_mode="fixed_size", + size=(4, 4), + ) + + assert result.ready, result.trace.blocker + assert result.predictions is not None + output_path = write_mapanything_scene_npz( + tmp_path / "outputs/mapanything/miniature-scene.npz", + result.predictions, + metadata={"completed_stages": list(result.trace.completed_stages)}, + ) + with np.load(output_path, allow_pickle=False) as payload: + assert set(MAPANYTHING_SCENE_OUTPUT_KEYS).issubset(payload.files) + assert "scene-generation" in str(payload["__metadata_json__"]) + + expected = json.loads(GOLDEN_MANIFEST.read_text(encoding="utf-8")) + assert_golden_close(_summarize_miniature_scene(result), expected) + + +def _summarize_miniature_scene( + result: MapAnythingSceneResult, +) -> dict[str, object]: + assert result.predictions is not None + stable_prediction_keys = tuple( + key + for key in MAPANYTHING_SCENE_OUTPUT_KEYS + if key not in {"intrinsics", "world_points"} + ) + return { + "schema_version": 2, + "fixture": { + "kind": "generated-miniature-checkpoint", + "source": "deterministic synthetic tensors", + "quantization": "none", + "random_seeds": {"encoder": 42, "heads": 123}, + "views": 2, + "image_size": [4, 4], + "patch_size": 2, + "encoder_layers": 2, + "info_sharing_layers": 2, + "covered": [ + "asset inspection", + "safetensors loading", + "image preprocessing", + "full encoder", + "fusion norm", + "multi-view info sharing", + "dense pose and scale heads", + "scene geometry postprocess", + "NPZ artifact writing", + ], + "not_covered": [ + "numeric parity with official MapAnything weights", + "production image resolution", + "performance or memory benchmarking", + ], + }, + "trace": { + "completed_stages": list(result.trace.completed_stages), + "frame_count": result.trace.frame_count, + "target_size": list(result.trace.target_size or ()), + "patch_grid": result.trace.metadata["patch_grid"], + "implemented_boundary": result.trace.metadata["implemented_boundary"], + }, + "predictions": { + key: summarize_array(getattr(result.predictions, key)) + for key in stable_prediction_keys + }, + "geometry_invariants": _summarize_geometry_invariants(result), + "artifact": { + "keys": sorted( + (*MAPANYTHING_SCENE_OUTPUT_KEYS, "__metadata_json__") + ), + "format": "npz", + }, + } + + +def _summarize_geometry_invariants( + result: MapAnythingSceneResult, +) -> dict[str, object]: + assert result.predictions is not None + intrinsics = np.asarray(result.predictions.intrinsics) + world_points = np.asarray(result.predictions.world_points) + focal_lengths = np.stack((intrinsics[:, 0, 0], intrinsics[:, 1, 1]), axis=1) + return { + "intrinsics": { + "shape": list(intrinsics.shape), + "dtype": str(intrinsics.dtype), + "finite": bool(np.isfinite(intrinsics).all()), + "positive_focal_lengths": bool((focal_lengths > 0.0).all()), + "homogeneous_bottom_row": bool( + np.allclose( + intrinsics[:, 2, :], + np.array([0.0, 0.0, 1.0], dtype=np.float32), + atol=1e-6, + rtol=0.0, + ) + ), + }, + "world_points": { + "shape": list(world_points.shape), + "dtype": str(world_points.dtype), + "finite": bool(np.isfinite(world_points).all()), + "nonzero": bool(np.linalg.norm(world_points) > 1e-3), + "bounded": bool(np.max(np.abs(world_points)) < 10.0), + }, + } +@pytest.mark.integration +@pytest.mark.real_assets def test_mapanything_scene_pipeline_generates_local_desk_npz(tmp_path): model_root = ROOT / "weights/map-anything" image_root = ROOT / "inputs/map-anything/desk" diff --git a/tests/test_mapanything_scene_postprocess_parity.py b/tests/test_mapanything_scene_postprocess_parity.py index bbd5461..e4ee53a 100644 --- a/tests/test_mapanything_scene_postprocess_parity.py +++ b/tests/test_mapanything_scene_postprocess_parity.py @@ -20,10 +20,13 @@ from mlx_spatial.mapanything_preprocess import MapAnythingPreprocessedInput, MapAnythingPreprocessedView -pytestmark = pytest.mark.skipif( - os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", - reason="opt-in MapAnything Torch reference parity", -) +pytestmark = [ + pytest.mark.torch_parity, + pytest.mark.skipif( + os.environ.get(MAPANYTHING_TORCH_PARITY_ENV) != "1", + reason="opt-in MapAnything Torch reference parity", + ), +] def test_mapanything_scene_postprocess_matches_desk_reference(): diff --git a/tests/test_pixal3d_derived_golden.py b/tests/test_pixal3d_derived_golden.py index aa4a7bc..c700e4e 100644 --- a/tests/test_pixal3d_derived_golden.py +++ b/tests/test_pixal3d_derived_golden.py @@ -20,9 +20,11 @@ MANIFEST_PATH = FIXTURE_ROOT / "golden.json" +@pytest.mark.integration def test_pixal3d_derived_golden_manifest_matches_committed_decoder_patch(): manifest = _manifest() + assert manifest["schema_version"] == 2 assert manifest["fixture_kind"] == "real-weight-derived-decoder-patch" assert manifest["source"]["repository"] == "TencentARC/Pixal3D" assert manifest["source"]["quantization"] == "none" @@ -54,9 +56,12 @@ def test_pixal3d_derived_golden_manifest_matches_committed_decoder_patch(): ) -@pytest.mark.heavy +@pytest.mark.integration +@pytest.mark.metal def test_pixal3d_derived_golden_replays_native_textured_export(tmp_path): expected = _manifest()["expected_export"] + vertex_tolerance = int(expected.pop("glb_vertex_tolerance")) + expected_vertices = int(expected["glb"].pop("vertices")) result = export_decoded_ovoxel_glb( FIXTURE_ROOT, tmp_path / "model.glb", @@ -85,7 +90,9 @@ def test_pixal3d_derived_golden_replays_native_textured_export(tmp_path): "faces": int(glb["total_faces"]), }, } + actual_vertices = int(actual["glb"].pop("vertices")) assert actual == expected + assert abs(actual_vertices - expected_vertices) <= vertex_tolerance def _manifest() -> dict: diff --git a/tests/test_pixal3d_pipeline.py b/tests/test_pixal3d_pipeline.py index 33000df..74689a5 100644 --- a/tests/test_pixal3d_pipeline.py +++ b/tests/test_pixal3d_pipeline.py @@ -4,6 +4,7 @@ import mlx.core as mx import numpy as np +import pytest from PIL import Image import mlx_spatial.pixal3d_inference as pixal3d_inference @@ -829,6 +830,7 @@ def test_pixal3d_pipeline_writes_texture_decoder_pbr_artifact_with_fake_decode_a assert result.trace.blocker.metadata["texture_decoder_attributes_shape"] == (4096, 6) +@pytest.mark.integration def test_pixal3d_pipeline_writes_textured_glb_with_fake_export_route(tmp_path, monkeypatch): root = write_fake_pixal3d_decode_root(tmp_path / "weights", proj_in_channels=3, sparse_steps=1, shape_steps=1, texture_steps=1) image = tmp_path / "image.png" diff --git a/tests/test_pytest_config.py b/tests/test_pytest_config.py index 8bf8c6f..467b3b6 100644 --- a/tests/test_pytest_config.py +++ b/tests/test_pytest_config.py @@ -1,15 +1,102 @@ +import ast +from pathlib import Path + import mlx.core as mx +ROOT = Path(__file__).resolve().parents[1] +ROUTINE_PIPELINE_GUARDS = { + "hyworld2": ( + "tests/test_hyworld2_inference.py", + "test_fixture_reconstruct_writes_staged_outputs_under_outputs", + ), + "lito": ( + "tests/test_lito_inference.py", + "test_full_pipeline_runs_on_sample_input", + ), + "mapanything": ( + "tests/test_mapanything_scene_pipeline.py", + "test_mapanything_miniature_scene_pipeline_matches_golden", + ), + "pixal3d": ( + "tests/test_pixal3d_pipeline.py", + "test_pixal3d_pipeline_writes_textured_glb_with_fake_export_route", + ), + "sam3d": ( + "tests/test_sam3d_tools.py", + "test_sam3d_cli_reconstruct_writes_gaussian_ply_and_textured_glb_with_fixture_pipeline", + ), + "trellis2": ( + "tests/test_trellis2_golden_fixture.py", + "test_miniature_int8_pipeline_emits_golden_trace_and_glb", + ), +} +ROUTINE_EXCLUSION_MARKERS = {"heavy", "real_assets", "torch_parity"} + + def test_pytest_sets_mlx_cpu_default(): assert mx.default_device() == mx.cpu -def test_pytest_defaults_skip_heavy_tests(pytestconfig): +def test_pytest_defaults_select_bounded_self_contained_tests(pytestconfig): addopts = pytestconfig.getini("addopts") joined = " ".join(addopts if isinstance(addopts, list) else [addopts]) markers = pytestconfig.getini("markers") assert "-m" in addopts - assert "not heavy" in joined - assert any(marker.startswith("heavy:") for marker in markers) + assert "not (heavy or real_assets or torch_parity)" in joined + for name in ( + "integration", + "real_assets", + "metal", + "heavy", + "benchmark", + "torch_parity", + ): + assert any(marker.startswith(f"{name}:") for marker in markers) + + +def test_supported_pipelines_keep_a_routine_integration_guard(): + errors = [] + for model, (relative_path, function_name) in ROUTINE_PIPELINE_GUARDS.items(): + path = ROOT / relative_path + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + function = next( + ( + node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == function_name + ), + None, + ) + if function is None: + errors.append(f"{model}: missing {relative_path}::{function_name}") + continue + + markers = { + marker + for decorator in function.decorator_list + if (marker := _pytest_marker_name(decorator)) is not None + } + if "integration" not in markers: + errors.append(f"{model}: representative guard is not marked integration") + excluded = markers & ROUTINE_EXCLUSION_MARKERS + if excluded: + errors.append( + f"{model}: representative guard is excluded from routine CI by {sorted(excluded)}" + ) + + assert not errors, "\n".join(errors) + + +def _pytest_marker_name(decorator: ast.expr) -> str | None: + target = decorator.func if isinstance(decorator, ast.Call) else decorator + if not isinstance(target, ast.Attribute): + return None + mark = target.value + if not isinstance(mark, ast.Attribute) or mark.attr != "mark": + return None + if not isinstance(mark.value, ast.Name) or mark.value.id != "pytest": + return None + return target.attr diff --git a/tests/test_sam3d_contract.py b/tests/test_sam3d_contract.py index 4c73f27..cf688a3 100644 --- a/tests/test_sam3d_contract.py +++ b/tests/test_sam3d_contract.py @@ -120,6 +120,7 @@ def _write_checkpoint(path: Path, prefixes: tuple[str, ...], *, omit_prefix: str @pytest.mark.heavy +@pytest.mark.real_assets @pytest.mark.skipif(not SAM3D_MLX_PIPELINE.is_file(), reason="SAM3D MLX weights absent") def test_real_sam3d_mlx_contract_maps_source_targets_and_weight_prefixes(): audit = audit_sam3d_source_weight_contract(SAM3D_MLX_ROOT) diff --git a/tests/test_sam3d_tools.py b/tests/test_sam3d_tools.py index 106354c..a338019 100644 --- a/tests/test_sam3d_tools.py +++ b/tests/test_sam3d_tools.py @@ -437,6 +437,7 @@ def test_sam3d_external_pointmap_cli_reports_validation_blocker(tmp_path, capsys assert not (tmp_path / output).exists() +@pytest.mark.integration def test_sam3d_cli_reconstruct_writes_gaussian_ply_and_textured_glb_with_fixture_pipeline(tmp_path, capsys, monkeypatch): monkeypatch.chdir(tmp_path) weights = tmp_path / "weights" diff --git a/tests/test_trellis2_forward.py b/tests/test_trellis2_forward.py index eedbfff..b9cf10f 100644 --- a/tests/test_trellis2_forward.py +++ b/tests/test_trellis2_forward.py @@ -747,6 +747,7 @@ def test_assess_dinov3_conditioning_reports_present_asset_config_blocker(tmp_pat @pytest.mark.heavy +@pytest.mark.metal def test_assess_dinov3_conditioning_reports_precise_transformer_blocker(tmp_path): _write_trellis2_root(tmp_path / "trellis") config = discover_trellis2_conditioning_config(tmp_path / "trellis").config @@ -1045,6 +1046,7 @@ def test_attempt_forward_trace_with_fake_dinov3_assets_reaches_sparse_boundary(t @pytest.mark.heavy +@pytest.mark.metal def test_attempt_forward_trace_with_executable_dinov3_assets_reaches_sparse_boundary(tmp_path): _write_trellis2_root(tmp_path / "trellis", conditioning_resolution=2) dino_root = tmp_path / "dinov3" diff --git a/tests/test_trellis2_golden_fixture.py b/tests/test_trellis2_golden_fixture.py index 4a08255..e41ae85 100644 --- a/tests/test_trellis2_golden_fixture.py +++ b/tests/test_trellis2_golden_fixture.py @@ -23,7 +23,8 @@ GOLDEN_MANIFEST = Path(__file__).parent / "data/trellis2_miniature_golden.json" -@pytest.mark.heavy +@pytest.mark.integration +@pytest.mark.metal def test_miniature_int8_pipeline_emits_golden_trace_and_glb(tmp_path, monkeypatch): fixture = build_trellis2_miniature_golden_fixture(tmp_path) exporter = Trellis2MiniatureSpatialKitExporter() diff --git a/tests/test_trellis2_texturing.py b/tests/test_trellis2_texturing.py index 5903e87..23df4aa 100644 --- a/tests/test_trellis2_texturing.py +++ b/tests/test_trellis2_texturing.py @@ -1,5 +1,4 @@ import json -import struct from pathlib import Path import mlx.core as mx @@ -12,7 +11,6 @@ from mlx_spatial.trellis2_texturing import ( Trellis2TexturingBlocker, Trellis2TexturingPipeline, - Trellis2TexturingResult, _load_obj_mesh, TRELLIS2_TEXTURING_DEFAULT_SEED, TRELLIS2_TEXTURING_DEFAULT_TEXTURE_SIZE, @@ -406,7 +404,7 @@ def _write_fixture_outputs_root(tmp_path: Path): rmbg_root = tmp_path / "fixture_weights/rmbg2" img_path = tmp_path / "fixture_inputs/demo.png" mesh_path = tmp_path / "fixture_inputs/cube.obj" - outputs_dir = Path("outputs/fixture_textured") + outputs_dir = tmp_path / "outputs/fixture_textured" outputs_dir.mkdir(parents=True, exist_ok=True) output_path = outputs_dir / "fixture_textured.glb" @@ -455,7 +453,8 @@ def test_run_rejects_non_glb_output(self, tmp_path): assert "only writes .glb" in result.blocker.reason def test_run_allows_export_path_outside_repository_outputs(self, tmp_path): - pipeline = Trellis2TexturingPipeline(root=tmp_path / "weights/trellis2") + root = tmp_path / "weights/trellis2" + pipeline = Trellis2TexturingPipeline(root=root) _write_rgb_image(tmp_path / "img.png") _write_obj_mesh(tmp_path / "mesh.obj") result = pipeline.run( @@ -465,7 +464,9 @@ def test_run_allows_export_path_outside_repository_outputs(self, tmp_path): ) assert not result.ready assert result.blocker is not None - assert result.blocker.operation != "export path validation" + assert result.blocker.stage == "asset-config" + assert result.blocker.operation == "TRELLIS.2 conditioning config discovery" + assert result.blocker.reason == f"pipeline config file not found: {root / 'pipeline.json'}" def test_run_rejects_missing_image(self, tmp_path): pipeline = Trellis2TexturingPipeline(root=tmp_path / "weights/trellis2") @@ -510,7 +511,8 @@ def test_run_rejects_bad_grid_size(self, tmp_path): assert result.blocker is not None assert result.blocker.stage == "mesh-preprocess" - def test_run_with_fixture_assets_produces_textured_glb(self, tmp_path): + @pytest.mark.integration + def test_run_with_fixture_assets_reaches_spatialkit_export_boundary(self, tmp_path): root, dinov3_root, rmbg_root, img_path, mesh_path, output_path = _write_fixture_outputs_root(tmp_path) pipeline = Trellis2TexturingPipeline( @@ -529,52 +531,12 @@ def test_run_with_fixture_assets_produces_textured_glb(self, tmp_path): glb_target_faces=100, ) - assert isinstance(result, Trellis2TexturingResult) - if result.ready: - assert result.artifact is not None - assert result.artifact.format == "glb" - assert output_path.is_file() - payload = output_path.read_bytes() - assert payload[:4] == b"glTF" - _verify_textured_glb_channels(payload) - else: - assert result.blocker is not None - assert result.blocker.stage in { - "mesh-export", "image-conditioning", "fdg-encoder", - "shape-decoder", "texture-slat", "texture-decoder", "decoded-artifact-write", - } - - def test_run_with_fixture_assets_512_pipeline_type(self, tmp_path): - root, dinov3_root, rmbg_root, img_path, mesh_path, output_path = _write_fixture_outputs_root(tmp_path) - - pipeline = Trellis2TexturingPipeline( - root=root, - dino_root=dinov3_root, - rmbg_root=rmbg_root, - ) - result = pipeline.run( - img_path, - mesh_path, - output_path=output_path, - pipeline_type="512", - seed=42, - grid_size=16, - slat_steps=1, - glb_target_faces=100, - ) - - assert isinstance(result, Trellis2TexturingResult) - if result.ready: - assert result.artifact is not None - assert result.artifact.format == "glb" - assert output_path.is_file() - _verify_textured_glb_channels(output_path.read_bytes()) - else: - assert result.blocker is not None - assert result.blocker.stage in { - "mesh-export", "image-conditioning", "fdg-encoder", - "shape-decoder", "texture-slat", "texture-decoder", "decoded-artifact-write", - } + assert not result.ready + assert result.blocker is not None + assert result.blocker.stage == "mesh-export" + assert result.blocker.operation == "export decoded TRELLIS.2 O-Voxel artifacts through SpatialKit" + assert result.blocker.reason == "mesh vertices must contain at least one vertex" + assert not output_path.exists() def test_run_missing_encoder_config_is_blocked(self, tmp_path): root = tmp_path / "weights/trellis2" @@ -599,6 +561,10 @@ def test_run_missing_encoder_config_is_blocked(self, tmp_path): ) assert not result.ready assert result.blocker is not None + assert result.blocker.stage == "fdg-encoder" + assert result.blocker.operation == "FDG encoder config validation" + assert result.blocker.reference == str(root / "shape_encoder.json") + assert "No such file or directory" in result.blocker.reason def test_run_missing_pipeline_config_is_blocked(self, tmp_path): img = tmp_path / "test.png" @@ -612,6 +578,11 @@ def test_run_missing_pipeline_config_is_blocked(self, tmp_path): result = pipeline.run(img, mesh, output_path=output) assert not result.ready assert result.blocker is not None + assert result.blocker.stage == "asset-config" + assert result.blocker.operation == "TRELLIS.2 conditioning config discovery" + assert result.blocker.reason == ( + f"pipeline config file not found: {tmp_path / 'nonexistent_weights/pipeline.json'}" + ) def test_run_bad_slat_steps_is_blocked(self, tmp_path): root = tmp_path / "weights/trellis2" @@ -634,6 +605,9 @@ def test_run_bad_slat_steps_is_blocked(self, tmp_path): ) assert not result.ready assert result.blocker is not None + assert result.blocker.stage == "texture-slat" + assert result.blocker.operation == "SLat step override validation" + assert result.blocker.reason == "slat_steps must be positive, got 0" def test_run_bad_pipeline_type_is_blocked(self, tmp_path): root = tmp_path / "weights/trellis2" @@ -656,6 +630,9 @@ def test_run_bad_pipeline_type_is_blocked(self, tmp_path): ) assert not result.ready assert result.blocker is not None + assert result.blocker.stage == "texture-slat" + assert result.blocker.operation == "texture SLat route selection" + assert result.blocker.reason == "unsupported texture SLat pipeline type: invalid" class TestLoadObjMesh: @@ -691,26 +668,3 @@ def test_handles_negative_indices(self, tmp_path): path.write_text("v 0 0 0\nv 1 0 0\nv 0 1 0\nf -3 -2 -1\n") verts, faces = _load_obj_mesh(path) assert faces.shape == (1, 3) - - -def _glb_json(payload: bytes) -> dict: - magic, version, total_length = struct.unpack_from("