From b30ba06b9a26b912d68c4410a746bd2b6ee2fba8 Mon Sep 17 00:00:00 2001 From: BP <11394934+benjipeng@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:49:26 -0400 Subject: [PATCH 1/2] test(trellis2): add self-contained INT8 golden fixture --- docs/development.md | 17 + tests/data/trellis2_miniature_golden.json | 107 ++++ tests/test_trellis2_golden_fixture.py | 70 +++ tests/trellis2_golden_fixture.py | 567 ++++++++++++++++++++++ 4 files changed, 761 insertions(+) create mode 100644 tests/data/trellis2_miniature_golden.json create mode 100644 tests/test_trellis2_golden_fixture.py create mode 100644 tests/trellis2_golden_fixture.py diff --git a/docs/development.md b/docs/development.md index 2d05e68..d24420c 100644 --- a/docs/development.md +++ b/docs/development.md @@ -77,6 +77,23 @@ 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 + +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. + +```bash +uv run pytest -m heavy tests/test_trellis2_golden_fixture.py -q +``` + +Reviewed tensor and GLB expectations live in +`tests/data/trellis2_miniature_golden.json`. Do not regenerate that manifest +automatically during tests. Rebaseline it only after reviewing an intentional +inference-contract change. + ## Editing Constraints - Prefer existing module boundaries over new abstractions. diff --git a/tests/data/trellis2_miniature_golden.json b/tests/data/trellis2_miniature_golden.json new file mode 100644 index 0000000..4df5a8b --- /dev/null +++ b/tests/data/trellis2_miniature_golden.json @@ -0,0 +1,107 @@ +{ + "fixture": { + "checkpoint_source": "generated synthetic tensors", + "glb_target_faces": 256, + "pipeline_type": "512", + "quantization_bits": 8, + "quantization_group_size": 64, + "sampler_steps": 1, + "seed": 7, + "texture_size": 32 + }, + "glb": { + "images": 2, + "materials": 1, + "meshes": 1, + "primitives": [ + { + "has_normal": true, + "has_texcoord_0": true, + "material": 0, + "positions": 4565, + "triangles": 3912 + } + ], + "sha256": "44b55e7184c56449d87a9bed214a7ab24460f9ad3bd159e92867cec67efa9470", + "textures": 2 + }, + "schema_version": 1, + "trace": { + "completed_stages": [ + "asset-config-validation", + "checkpoint-probe-readiness", + "input-image", + "image-preprocessing-background", + "image-conditioning", + "sparse-structure-sampling", + "shape-slat-sampling", + "texture-slat-sampling", + "shape-decoder", + "texture-decoder", + "decoded-artifact-write", + "mesh-export" + ], + "tensor_outputs": { + "cond_512": { + "dtype": "float32", + "sha256": "aad5ce8209c55c84c2e64b92b8162e7e49ba89af0f19270a3762df796b90a5e5", + "shape": [ + 1, + 65, + 64 + ] + }, + "shape_flexidualgrid_fields": { + "dtype": "float32", + "sha256": "d13460feb6ebeeb09b7fb153ee3c5c607bf1ec2205471396b5a8c5908e5f1f62", + "shape": [ + 512, + 7 + ] + }, + "shape_slat": { + "dtype": "float32", + "sha256": "69670f899f1c5f681fb9b3fd71a229e3f4e22271ed6ee470f36af14f32cc31bd", + "shape": [ + 64, + 32 + ] + }, + "sparse_latent": { + "dtype": "float32", + "sha256": "2fa4d566772d77461fd6630c042e69b0002adf31ee4b724b6557c8a45666ef51", + "shape": [ + 1, + 2, + 2, + 2, + 2 + ] + }, + "texture_slat": { + "dtype": "float32", + "sha256": "99d23e5097506efb79c83865dc63c505ad4fad0db6c8d468d22845baa66b2141", + "shape": [ + 64, + 32 + ] + }, + "texture_voxel_attrs": { + "dtype": "float32", + "sha256": "49b45f028a546f707bcdf37677c63eb279187de4561bbe06f1a91720320396b4", + "shape": [ + 512, + 6 + ] + }, + "texture_voxel_coordinates": { + "dtype": "int32", + "sha256": "d572dc92213b089492d529519ba3356793c57f19955851405b63c16bc944b436", + "shape": [ + 512, + 4 + ] + } + } + } +} diff --git a/tests/test_trellis2_golden_fixture.py b/tests/test_trellis2_golden_fixture.py new file mode 100644 index 0000000..bc2b8a7 --- /dev/null +++ b/tests/test_trellis2_golden_fixture.py @@ -0,0 +1,70 @@ +"""End-to-end regression coverage for the miniature TRELLIS.2 golden fixture.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from mlx_spatial.trellis2_inference import Trellis2InferencePipeline +from mlx_spatial.trellis2_quantization import read_trellis2_quantization_spec +from tests.trellis2_golden_fixture import ( + build_trellis2_miniature_golden_fixture, + summarize_glb, + summarize_trellis2_golden_trace, +) + +GOLDEN_MANIFEST = Path(__file__).parent / "data/trellis2_miniature_golden.json" + + +@pytest.mark.heavy +def test_miniature_int8_pipeline_emits_golden_trace_and_glb(tmp_path): + fixture = build_trellis2_miniature_golden_fixture(tmp_path) + quantization = read_trellis2_quantization_spec( + fixture.quantized_root / "ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors" + ) + assert quantization is not None + assert quantization.bits == 8 + assert quantization.group_size == 64 + assert quantization.tensors + + dino_quantization = read_trellis2_quantization_spec( + fixture.dino_root / "model.safetensors" + ) + assert dino_quantization is not None + assert dino_quantization.bits == 8 + assert dino_quantization.group_size == 64 + assert dino_quantization.tensors + + result = Trellis2InferencePipeline(fixture.quantized_root).generate_textured_glb( + fixture.image_path, + output_path=fixture.output_path, + dino_root=fixture.dino_root, + pipeline_type="512", + seed=7, + max_num_tokens=4_096, + decoder_token_limit=10_000, + texture_size=32, + glb_target_faces=256, + retain_trace_payloads=True, + ) + + assert result.ready, result.trace.blocker + summary = { + "schema_version": 1, + "fixture": { + "checkpoint_source": "generated synthetic tensors", + "pipeline_type": "512", + "seed": 7, + "sampler_steps": 1, + "quantization_bits": 8, + "quantization_group_size": 64, + "texture_size": 32, + "glb_target_faces": 256, + }, + "trace": summarize_trellis2_golden_trace(result.trace), + "glb": summarize_glb(fixture.output_path), + } + expected = json.loads(GOLDEN_MANIFEST.read_text(encoding="utf-8")) + assert summary == expected diff --git a/tests/trellis2_golden_fixture.py b/tests/trellis2_golden_fixture.py new file mode 100644 index 0000000..201d332 --- /dev/null +++ b/tests/trellis2_golden_fixture.py @@ -0,0 +1,567 @@ +"""Self-contained miniature TRELLIS.2 golden fixture construction.""" + +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 + +import mlx.core as mx +import numpy as np +from PIL import Image + +from mlx_spatial.safetensors_io import save_safetensors +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 + + +@dataclass(frozen=True) +class Trellis2MiniatureGoldenFixture: + """Paths for one generated miniature affine-INT8 TRELLIS.2 fixture.""" + + source_root: Path + quantized_root: Path + dino_root: Path + image_path: Path + output_path: Path + + +def build_trellis2_miniature_golden_fixture(root: Path) -> Trellis2MiniatureGoldenFixture: + """Create tiny source checkpoints, quantize them, and return runnable fixture paths.""" + + source_root = root / "trellis2-source" + quantized_root = root / "trellis2-mlx-8bit" + dino_root = root / "dinov3" + image_path = root / "input.png" + output_path = root / "output" / "model.glb" + + sparse_flow = SparseStructureFlowConfig( + name="SparseStructureFlowModel", + resolution=2, + in_channels=2, + out_channels=2, + model_channels=64, + cond_channels=64, + num_blocks=1, + num_heads=4, + mlp_ratio=2.0, + pe_mode="rope", + share_mod=True, + initialization="scaled", + qk_rms_norm=True, + qk_rms_norm_cross=True, + dtype="float32", + ) + sparse_decoder = SparseStructureDecoderConfig( + name="SparseStructureDecoder", + out_channels=1, + latent_channels=2, + num_res_blocks=0, + channels=(4, 4), + num_res_blocks_middle=0, + norm_type="layer", + use_fp16=False, + ) + shape_slat = _slat_config(in_channels=32) + texture_slat = _slat_config(in_channels=64) + shape_decoder = _decoder_config(name="FlexiDualGridVaeDecoder", out_channels=7, pred_subdiv=True) + texture_decoder = _decoder_config(name="SparseUnetVaeDecoder", out_channels=6, pred_subdiv=False) + dino_config = DinoV3ModelConfig( + model_type="dinov3_vit", + image_size=512, + patch_size=64, + hidden_size=64, + num_hidden_layers=1, + num_attention_heads=4, + intermediate_size=128, + layer_norm_eps=1e-5, + use_swiglu_ffn=False, + num_register_tokens=0, + expected_feature_width=64, + rope_theta=100.0, + pos_embed_rescale=2.0, + ) + + _write_pipeline(source_root) + _write_json( + source_root / "ckpts/ss_flow_img_dit_1_3B_64_bf16.json", + _sparse_flow_config_payload(sparse_flow), + ) + _write_flow_checkpoint( + source_root / "ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors", + in_channels=sparse_flow.in_channels, + out_channels=sparse_flow.out_channels, + model_channels=sparse_flow.model_channels, + cond_channels=sparse_flow.cond_channels, + num_heads=sparse_flow.num_heads, + mlp_ratio=sparse_flow.mlp_ratio, + ) + + sparse_decoder_base = source_root / "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16" + _write_json(sparse_decoder_base.with_suffix(".json"), _sparse_decoder_config_payload(sparse_decoder)) + _write_sparse_decoder_checkpoint(sparse_decoder_base.with_suffix(".safetensors"), sparse_decoder) + + for resolution in (512, 1024): + shape_base = source_root / f"ckpts/slat_flow_img2shape_dit_1_3B_{resolution}_bf16" + _write_json(shape_base.with_suffix(".json"), _slat_config_payload(shape_slat, resolution=32)) + _write_flow_checkpoint( + shape_base.with_suffix(".safetensors"), + in_channels=shape_slat.in_channels, + out_channels=shape_slat.out_channels, + model_channels=shape_slat.model_channels, + cond_channels=shape_slat.cond_channels, + num_heads=shape_slat.num_heads, + mlp_ratio=shape_slat.mlp_ratio, + ) + texture_base = source_root / f"ckpts/slat_flow_imgshape2tex_dit_1_3B_{resolution}_bf16" + _write_json(texture_base.with_suffix(".json"), _slat_config_payload(texture_slat, resolution=32)) + _write_flow_checkpoint( + texture_base.with_suffix(".safetensors"), + in_channels=texture_slat.in_channels, + out_channels=texture_slat.out_channels, + model_channels=texture_slat.model_channels, + cond_channels=texture_slat.cond_channels, + num_heads=texture_slat.num_heads, + mlp_ratio=texture_slat.mlp_ratio, + ) + + shape_decoder_base = source_root / "ckpts/shape_dec_next_dc_f16c32_fp16" + _write_json(shape_decoder_base.with_suffix(".json"), _decoder_config_payload(shape_decoder)) + _write_decoder_checkpoint(shape_decoder_base.with_suffix(".safetensors"), shape_decoder) + texture_decoder_base = source_root / "ckpts/tex_dec_next_dc_f16c32_fp16" + _write_json(texture_decoder_base.with_suffix(".json"), _decoder_config_payload(texture_decoder)) + _write_decoder_checkpoint(texture_decoder_base.with_suffix(".safetensors"), texture_decoder) + + for name in ("shape_enc_next_dc_f16c32_fp16", "tex_enc_next_dc_f16c32_fp16"): + base = source_root / "ckpts" / name + _write_json(base.with_suffix(".json"), {"name": "MiniatureUnusedEncoder", "args": {}}) + save_safetensors( + base.with_suffix(".safetensors"), + {"blocks.0.0.mlp.0.weight": _pattern((64, 64), scale=0.01)}, + ) + + _write_dinov3(dino_root, dino_config) + quantize_trellis2_weights( + source_root, + quantized_root, + dinov3_root=dino_root, + bits=8, + group_size=64, + ) + _write_rgba_fixture(image_path) + return Trellis2MiniatureGoldenFixture( + source_root=source_root, + quantized_root=quantized_root, + dino_root=quantized_root / "dinov3", + image_path=image_path, + output_path=output_path, + ) + + +def summarize_trellis2_golden_trace(trace: Trellis2ForwardTraceResult) -> dict[str, Any]: + """Return a stable, JSON-compatible summary of tensor-bearing stage outputs.""" + + tensor_outputs = {} + 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), + } + 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", + resolution=32, + in_channels=in_channels, + out_channels=32, + model_channels=64, + cond_channels=64, + num_blocks=1, + num_heads=4, + mlp_ratio=2.0, + pe_mode="rope", + share_mod=True, + initialization="scaled", + qk_rms_norm=True, + qk_rms_norm_cross=True, + dtype="float32", + ) + + +def _decoder_config(*, name: str, out_channels: int, pred_subdiv: bool) -> StructuredLatentDecoderConfig: + return StructuredLatentDecoderConfig( + name=name, + latent_channels=32, + model_channels=(64, 32), + num_blocks=(1, 0), + block_type=("SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d"), + up_block_type=("SparseResBlockC2S3d",), + use_fp16=False, + out_channels=out_channels, + resolution=256 if name == "FlexiDualGridVaeDecoder" else None, + pred_subdiv=pred_subdiv, + ) + + +def _write_pipeline(root: Path) -> None: + models = { + "sparse_structure_decoder": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16", + "sparse_structure_flow_model": "ckpts/ss_flow_img_dit_1_3B_64_bf16", + "shape_slat_decoder": "ckpts/shape_dec_next_dc_f16c32_fp16", + "shape_slat_flow_model_512": "ckpts/slat_flow_img2shape_dit_1_3B_512_bf16", + "shape_slat_flow_model_1024": "ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16", + "tex_slat_decoder": "ckpts/tex_dec_next_dc_f16c32_fp16", + "tex_slat_flow_model_512": "ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16", + "tex_slat_flow_model_1024": "ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16", + } + sampler = { + "name": "FlowEulerGuidanceIntervalSampler", + "args": {"sigma_min": 1e-5}, + "params": { + "steps": 1, + "guidance_strength": 1.0, + "guidance_rescale": 0.0, + "guidance_interval": [0.0, 1.0], + "rescale_t": 1.0, + }, + } + normalization = {"mean": [0.0] * 32, "std": [1.0] * 32} + _write_json( + root / "pipeline.json", + { + "name": "Trellis2ImageTo3DPipeline", + "args": { + "models": models, + "sparse_structure_sampler": sampler, + "shape_slat_sampler": sampler, + "shape_slat_normalization": normalization, + "tex_slat_sampler": sampler, + "tex_slat_normalization": normalization, + "image_cond_model": { + "name": "DinoV3FeatureExtractor", + "args": {"model_name": "miniature/dinov3"}, + }, + "rembg_model": {"name": "UnusedForRgbaFixture", "args": {}}, + "default_pipeline_type": "512", + }, + }, + ) + _write_json(root / "texturing_pipeline.json", {"name": "MiniatureTexturingPipeline", "args": {}}) + + +def _write_flow_checkpoint( + path: Path, + *, + in_channels: int, + out_channels: int, + model_channels: int, + cond_channels: int, + num_heads: int, + mlp_ratio: float, +) -> None: + head_dim = model_channels // num_heads + mlp_channels = int(model_channels * mlp_ratio) + tensors = { + "input_layer.weight": _pattern((model_channels, in_channels), scale=0.02), + "input_layer.bias": mx.zeros((model_channels,), dtype=mx.float32), + "out_layer.weight": _pattern((out_channels, model_channels), scale=0.02), + "out_layer.bias": mx.zeros((out_channels,), dtype=mx.float32), + "t_embedder.mlp.0.weight": _pattern((model_channels, 256), scale=0.01), + "t_embedder.mlp.0.bias": mx.zeros((model_channels,), dtype=mx.float32), + "t_embedder.mlp.2.weight": _pattern((model_channels, model_channels), scale=0.01), + "t_embedder.mlp.2.bias": mx.zeros((model_channels,), dtype=mx.float32), + "adaLN_modulation.1.weight": _pattern((model_channels * 6, model_channels), scale=0.005), + "adaLN_modulation.1.bias": mx.zeros((model_channels * 6,), dtype=mx.float32), + "blocks.0.modulation": mx.zeros((model_channels * 6,), dtype=mx.float32), + "blocks.0.norm2.weight": mx.ones((model_channels,), dtype=mx.float32), + "blocks.0.norm2.bias": mx.zeros((model_channels,), dtype=mx.float32), + "blocks.0.self_attn.to_qkv.weight": _pattern((model_channels * 3, model_channels), scale=0.01), + "blocks.0.self_attn.to_qkv.bias": mx.zeros((model_channels * 3,), dtype=mx.float32), + "blocks.0.self_attn.q_rms_norm.gamma": mx.ones((num_heads, head_dim), dtype=mx.float32), + "blocks.0.self_attn.k_rms_norm.gamma": mx.ones((num_heads, head_dim), dtype=mx.float32), + "blocks.0.self_attn.to_out.weight": _pattern((model_channels, model_channels), scale=0.01), + "blocks.0.self_attn.to_out.bias": mx.zeros((model_channels,), dtype=mx.float32), + "blocks.0.cross_attn.to_q.weight": _pattern((model_channels, model_channels), scale=0.01), + "blocks.0.cross_attn.to_q.bias": mx.zeros((model_channels,), dtype=mx.float32), + "blocks.0.cross_attn.to_kv.weight": _pattern((model_channels * 2, cond_channels), scale=0.01), + "blocks.0.cross_attn.to_kv.bias": mx.zeros((model_channels * 2,), dtype=mx.float32), + "blocks.0.cross_attn.q_rms_norm.gamma": mx.ones((num_heads, head_dim), dtype=mx.float32), + "blocks.0.cross_attn.k_rms_norm.gamma": mx.ones((num_heads, head_dim), dtype=mx.float32), + "blocks.0.cross_attn.to_out.weight": _pattern((model_channels, model_channels), scale=0.01), + "blocks.0.cross_attn.to_out.bias": mx.zeros((model_channels,), dtype=mx.float32), + "blocks.0.mlp.mlp.0.weight": _pattern((mlp_channels, model_channels), scale=0.01), + "blocks.0.mlp.mlp.0.bias": mx.zeros((mlp_channels,), dtype=mx.float32), + "blocks.0.mlp.mlp.2.weight": _pattern((model_channels, mlp_channels), scale=0.01), + "blocks.0.mlp.mlp.2.bias": mx.zeros((model_channels,), dtype=mx.float32), + } + path.parent.mkdir(parents=True, exist_ok=True) + save_safetensors(path, tensors) + + +def _write_sparse_decoder_checkpoint(path: Path, config: SparseStructureDecoderConfig) -> None: + tensors = { + "input_layer.weight": _pattern((config.channels[0], config.latent_channels, 3, 3, 3), scale=0.02), + "input_layer.bias": mx.zeros((config.channels[0],), dtype=mx.float32), + "blocks.0.conv.weight": _pattern((config.channels[1] * 8, config.channels[0], 3, 3, 3), scale=0.01), + "blocks.0.conv.bias": mx.zeros((config.channels[1] * 8,), dtype=mx.float32), + "out_layer.0.weight": mx.ones((config.channels[-1],), dtype=mx.float32), + "out_layer.0.bias": mx.zeros((config.channels[-1],), dtype=mx.float32), + "out_layer.2.weight": _pattern((1, config.channels[-1], 3, 3, 3), scale=0.01), + "out_layer.2.bias": mx.full((1,), 4.0, dtype=mx.float32), + } + path.parent.mkdir(parents=True, exist_ok=True) + save_safetensors(path, tensors) + + +def _write_decoder_checkpoint(path: Path, config: StructuredLatentDecoderConfig) -> None: + first_channels, second_channels = config.model_channels + tensors = { + "from_latent.weight": _pattern((first_channels, config.latent_channels), scale=0.02), + "from_latent.bias": mx.zeros((first_channels,), dtype=mx.float32), + "output_layer.weight": _pattern((config.out_channels, second_channels), scale=0.01), + "output_layer.bias": mx.zeros((config.out_channels,), dtype=mx.float32), + "blocks.0.0.conv.weight": _center_identity_conv(first_channels), + "blocks.0.0.conv.bias": mx.zeros((first_channels,), dtype=mx.float32), + "blocks.0.0.norm.weight": mx.ones((first_channels,), dtype=mx.float32), + "blocks.0.0.norm.bias": mx.zeros((first_channels,), dtype=mx.float32), + "blocks.0.0.mlp.0.weight": _pattern((first_channels * 4, first_channels), scale=0.01), + "blocks.0.0.mlp.0.bias": mx.zeros((first_channels * 4,), dtype=mx.float32), + "blocks.0.0.mlp.2.weight": _pattern((first_channels, first_channels * 4), scale=0.01), + "blocks.0.0.mlp.2.bias": mx.zeros((first_channels,), dtype=mx.float32), + "blocks.0.1.norm1.weight": mx.ones((first_channels,), dtype=mx.float32), + "blocks.0.1.norm1.bias": mx.zeros((first_channels,), dtype=mx.float32), + "blocks.0.1.conv1.weight": _pattern((second_channels * 8, 3, 3, 3, first_channels), scale=0.005), + "blocks.0.1.conv1.bias": mx.zeros((second_channels * 8,), dtype=mx.float32), + "blocks.0.1.conv2.weight": _center_identity_conv(second_channels), + "blocks.0.1.conv2.bias": mx.zeros((second_channels,), dtype=mx.float32), + } + if config.pred_subdiv: + tensors["blocks.0.1.to_subdiv.weight"] = _pattern((8, first_channels), scale=0.005) + tensors["blocks.0.1.to_subdiv.bias"] = mx.full((8,), 8.0, dtype=mx.float32) + tensors["output_layer.bias"] = mx.array([0.0, 0.0, 0.0, 4.0, 4.0, 4.0, 0.0], dtype=mx.float32) + path.parent.mkdir(parents=True, exist_ok=True) + save_safetensors(path, tensors) + + +def _write_dinov3(root: Path, config: DinoV3ModelConfig) -> None: + _write_json( + root / "config.json", + { + "model_type": config.model_type, + "image_size": config.image_size, + "patch_size": config.patch_size, + "hidden_size": config.hidden_size, + "num_hidden_layers": config.num_hidden_layers, + "num_attention_heads": config.num_attention_heads, + "intermediate_size": config.intermediate_size, + "layer_norm_eps": config.layer_norm_eps, + "use_swiglu_ffn": config.use_swiglu_ffn, + "num_register_tokens": config.num_register_tokens, + "rope_theta": config.rope_theta, + "pos_embed_rescale": config.pos_embed_rescale, + }, + ) + hidden = config.hidden_size + intermediate = config.intermediate_size + layer = "layer.0" + tensors = { + "embeddings.cls_token": mx.zeros((1, 1, hidden), dtype=mx.float32), + "embeddings.patch_embeddings.bias": mx.zeros((hidden,), dtype=mx.float32), + "embeddings.patch_embeddings.weight": _pattern( + (hidden, 3, config.patch_size, config.patch_size), + scale=0.002, + ), + "norm.bias": mx.zeros((hidden,), dtype=mx.float32), + "norm.weight": mx.ones((hidden,), dtype=mx.float32), + f"{layer}.attention.k_proj.weight": _pattern((hidden, hidden), scale=0.01), + f"{layer}.attention.o_proj.bias": mx.zeros((hidden,), dtype=mx.float32), + f"{layer}.attention.o_proj.weight": _pattern((hidden, hidden), scale=0.01), + f"{layer}.attention.q_proj.bias": mx.zeros((hidden,), dtype=mx.float32), + f"{layer}.attention.q_proj.weight": _pattern((hidden, hidden), scale=0.01), + f"{layer}.attention.v_proj.bias": mx.zeros((hidden,), dtype=mx.float32), + f"{layer}.attention.v_proj.weight": _pattern((hidden, hidden), scale=0.01), + f"{layer}.layer_scale1.lambda1": mx.full((hidden,), 0.1, dtype=mx.float32), + f"{layer}.layer_scale2.lambda1": mx.full((hidden,), 0.1, dtype=mx.float32), + f"{layer}.mlp.down_proj.bias": mx.zeros((hidden,), dtype=mx.float32), + f"{layer}.mlp.down_proj.weight": _pattern((hidden, intermediate), scale=0.01), + f"{layer}.mlp.up_proj.bias": mx.zeros((intermediate,), dtype=mx.float32), + f"{layer}.mlp.up_proj.weight": _pattern((intermediate, hidden), scale=0.01), + f"{layer}.norm1.bias": mx.zeros((hidden,), dtype=mx.float32), + f"{layer}.norm1.weight": mx.ones((hidden,), dtype=mx.float32), + f"{layer}.norm2.bias": mx.zeros((hidden,), dtype=mx.float32), + f"{layer}.norm2.weight": mx.ones((hidden,), dtype=mx.float32), + } + root.mkdir(parents=True, exist_ok=True) + save_safetensors(root / "model.safetensors", tensors) + + +def _sparse_flow_config_payload(config: SparseStructureFlowConfig) -> dict[str, Any]: + return {"name": config.name, "args": _flow_config_args(config)} + + +def _slat_config_payload(config: SLatFlowConfig, *, resolution: int) -> dict[str, Any]: + args = _flow_config_args(config) + args["resolution"] = resolution + return {"name": config.name, "args": args} + + +def _flow_config_args(config: SparseStructureFlowConfig | SLatFlowConfig) -> dict[str, Any]: + return { + "resolution": config.resolution, + "in_channels": config.in_channels, + "out_channels": config.out_channels, + "model_channels": config.model_channels, + "cond_channels": config.cond_channels, + "num_blocks": config.num_blocks, + "num_heads": config.num_heads, + "mlp_ratio": config.mlp_ratio, + "pe_mode": config.pe_mode, + "share_mod": config.share_mod, + "initialization": config.initialization, + "qk_rms_norm": config.qk_rms_norm, + "qk_rms_norm_cross": config.qk_rms_norm_cross, + "dtype": config.dtype, + } + + +def _sparse_decoder_config_payload(config: SparseStructureDecoderConfig) -> dict[str, Any]: + return { + "name": config.name, + "args": { + "out_channels": config.out_channels, + "latent_channels": config.latent_channels, + "num_res_blocks": config.num_res_blocks, + "channels": list(config.channels), + "num_res_blocks_middle": config.num_res_blocks_middle, + "norm_type": config.norm_type, + "use_fp16": config.use_fp16, + }, + } + + +def _decoder_config_payload(config: StructuredLatentDecoderConfig) -> dict[str, Any]: + args: dict[str, Any] = { + "model_channels": list(config.model_channels), + "latent_channels": config.latent_channels, + "num_blocks": list(config.num_blocks), + "block_type": list(config.block_type), + "up_block_type": list(config.up_block_type), + "block_args": [{} for _ in config.model_channels], + "use_fp16": config.use_fp16, + } + if config.name == "FlexiDualGridVaeDecoder": + args["resolution"] = config.resolution + else: + args["out_channels"] = config.out_channels + args["pred_subdiv"] = config.pred_subdiv + return {"name": config.name, "args": args} + + +def _write_rgba_fixture(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + image = Image.new("RGBA", (16, 16), (0, 0, 0, 0)) + pixels = image.load() + for y in range(3, 13): + for x in range(4, 12): + pixels[x, y] = (40 + x * 4, 30 + y * 5, 120, 255) + image.save(path) + + +def _write_json(path: Path, payload: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def _pattern(shape: tuple[int, ...], *, scale: float) -> mx.array: + size = math.prod(shape) + values = (np.arange(size, dtype=np.float32) % 23.0 - 11.0) / 11.0 + return mx.array((values * scale).reshape(shape), dtype=mx.float32) + + +def _center_identity_conv(channels: int) -> mx.array: + values = np.zeros((channels, 3, 3, 3, channels), dtype=np.float32) + indices = np.arange(channels) + values[indices, 1, 1, 1, indices] = 1.0 + 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=" Date: Fri, 14 Aug 2026 14:49:46 -0400 Subject: [PATCH 2/2] docs(trellis2): update high-resolution validation --- model-cards/trellis2-mlx-8bit/README.md | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/model-cards/trellis2-mlx-8bit/README.md b/model-cards/trellis2-mlx-8bit/README.md index b84bbf5..01035cc 100644 --- a/model-cards/trellis2-mlx-8bit/README.md +++ b/model-cards/trellis2-mlx-8bit/README.md @@ -111,10 +111,13 @@ mlx-spatial-trellis2 generate-textured \ --dino-root weights/trellis2-mlx-8bit/dinov3 \ --rmbg-root weights/trellis2-mlx-8bit/rmbg \ --output outputs/trellis2/object-8bit/model.glb \ - --pipeline-type 512 \ + --pipeline-type 1024_cascade \ --seed 42 ``` +`1024_cascade` is the recommended quality tier; use `512` when lower memory +use or faster iteration matters more. + Do not pass `--slat-steps` for a quality run; the model configuration uses 12 steps. `--slat-steps 1` is intended only for a quick runtime smoke test. @@ -205,19 +208,22 @@ of the quantized checkpoints to the 8-bit repository. recognizable appearance. - The same run completed in 152.23 seconds, observed 3.516 GB peak MLX allocator use, and recorded zero swap growth. +- A separate `1024_cascade`, 12-step run completed the same end-to-end path + and produced a Blender-readable 12,670,448-byte GLB with 199,884 faces and + embedded 1024 x 1024 PBR textures. -The runtime and memory figures are one local Apple Silicon observation, not a -general benchmark. The run establishes executable compatibility and artifact -health; it is not a formal claim of visual equivalence to the source weights. +The 512 runtime and memory figures are one local Apple Silicon observation, +not a general benchmark. The `1024_cascade` run overlapped another MLX workload, +so it establishes compatibility and artifact health rather than performance. +Neither run is a formal claim of visual equivalence to the source weights. ## Limitations - Quantization changes the sampling trajectory. Geometry, pose, topology, texture placement, material values, and unseen surfaces can differ from the source-precision model even with the same seed. -- The verified end-to-end run used the 512 pipeline. The 1024 and cascade - checkpoints passed inventory and runtime tests but have not received an - equivalent end-to-end quality evaluation here. +- End-to-end validation covers `512` and `1024_cascade`. The standalone `1024` + and `1536_cascade` routes have not received an equivalent quality evaluation. - Single-view reconstruction cannot determine unseen geometry with certainty. - Fine detail depends on foreground extraction, cropping, occlusion, reflections, transparency, and thin structures.