Skip to content

idol: tf-idol-generate, one picture to a 3DGS PLY end to end on Metal (stacked on #1773, #1774, #1775) - #1780

Draft
jayleaton wants to merge 21 commits into
ashhart:mainfrom
jayleaton:idol-e2e
Draft

jayleaton wants to merge 21 commits into
ashhart:mainfrom
jayleaton:idol-e2e

Conversation

@jayleaton

@jayleaton jayleaton commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

Stacked on #1773, #1774 and #1775 (and through them #1766). This branch merges those three; its own change is the last two commits (25b6e6f50, 883353480). Review after them.

What this changes

tf-idol-generate: one SDK-preprocessed picture to a 3DGS PLY and camera, natively on M5 Metal. The steps are DINOv3's feature1 (#1775), the FLUX.2 VAE's feature2 (#1774), the denoiser's condition and 50 guided steps (#1766), then the octree and Gaussian decoders (#1773). Each model loads when its stage starts and is freed when it ends. For #1757; with this, the issue's scope is covered on Metal except background removal. Details: the "End to end" part of docs/recipes/idol.md.

  • Background removal and cropping (BiRefNet, the SDK's preprocess) are not ported. The input is the picture preprocess returns, as the case tools write it.
  • Starting noise: torch.Generator("cpu").manual_seed(seed)'s randn, reimplemented (mt19937 and Box-Muller). The camera's five values are bit-equal to torch's; 28 of the 131,072 latent values are 1 ulp off, from torch's float cos and sin. The VAE posterior's noise and the octree's draws come from native seeded generators. --noise FILE injects recorded noise for parity.
  • The merge also moves the VAE case into tools/zig/idol_vae_case.py, as DINOv3's is, so the merged case tool stays under 600 lines.

CLI: extended by #1787

tf-idol-generate's arguments, flags, defaults and outputs here are the interface integrations build on. #1787 extends them additively and changes none of them:

  • PICTURE may also be a .png, .jpg, .jpeg or .webp photo, preprocessed natively (BiRefNet from SNAPSHOT_DIR/birefnet.safetensors);
  • new flags --mode synthetic|real (default synthetic) and --crop-margin X;
  • a photo also writes OUT_DIR/preprocessed.png;
  • a photo's tf-idol-generate.json adds preprocess_s and birefnet_s, and its load_s includes BiRefNet's load.

A safetensors PICTURE behaves exactly as in this PR.

On CUDA (#1862) the tool takes the same arguments, flags, defaults and outputs; its JSON's peak key is cuda_peak_allocated_gb where Metal's is metal_peak_allocated_gb.

Kernel contract

There are no new kernels. All of Idol's together, on the merged tree:

SUMMARY 57 kernels: 57 pass gates 1-4, 0 fail, 0 known findings; 57 controls: 57 red as planted, 0 missed; 0 limit problems

Receipt

  • Environment: Mac Studio M5 Ultra 256 GB, macOS 27; this branch; idol-sdk 4702a7f0 IdolPipeline.__call__ on torch 2.13.0 (MPS), each stage timed by wrapping its methods with torch.mps.synchronize(); the avocado example, seed 42, 50 steps, guidance 3, 262,144 Gaussians. The SDK's PLY is byte-identical to an earlier run of it.

  • Stage timings (ours: two fresh processes, weights in the page cache both times):

    stage ours SDK, MPS
    load (all models) 0.67 s 2.11 s (and 2.75 s of imports)
    preprocess (BiRefNet) not ported 0.51 s
    DINOv3 0.155 s 0.275 s
    VAE 0.241 s 0.119 s
    condition 0.071 s in the sampling
    denoise, 50 steps x 2 27.50 s 93.89 s
    octree 0.143 s 0.612 s
    Gaussian decoder 0.223 s 0.504 s
    PLY and camera 0.004 s 0.018 s
    total 29.1 s 100.8 s (95.4 s without preprocess)
  • Memory (/usr/bin/time -l, the whole process): ours 2.71 GB peak footprint (2.40 GB Metal allocated at most); the SDK 12.88 GB (5.01 GB MPS allocated, 12.64 GB driver).

  • Quality: the PLYs cannot be equal, because the SDK draws both its starting noise and the VAE posterior's noise from MPS's generator. Symmetric nearest-neighbour distance between 4,096-point samples, where two samples of the SDK's own PLY sit 0.00896 apart:

    run NN opacity mean
    the SDK 0.00906 0.885
    ours, native noise 0.0462 0.423
    ours, the SDK's VAE noise only 0.0463 0.425
    ours, the SDK's VAE and starting noise 0.00922 0.887
    the SDK itself, same features, CPU-seeded starting noise 0.0462 0.422

    The starting noise, not the port, makes the gap. The SDK's own run with CPU-seeded noise lands as far from its MPS-noise run as ours does, and its sample has cosine 0.479 to that run's. Given the SDK's noises, ours is as close as sampling resolves: sample cosine 0.99835, camera yaw/pitch/fov 1.6956/0.5397/0.1094 against 1.7006/0.5413/0.1096. Every PLY has 262,144 finite Gaussians, and repeated runs of ours are byte-identical. I did not render them side by side.

  • Parity chain, with the SDK's recorded features and CPU-seeded noise: the sample has cosine 0.999991 to the SDK's from the same inputs, and the camera is within 0.0022. The PLY sits at NN 0.00981 (baseline 0.00944) from the SDK's decode of its own sample. From the picture through our encoders, the sample moves to cosine 0.9944 with native VAE noise and 0.9995 with the SDK's.

  • Tests: zig build test test-golden -Dcpu=apple_m1: 585 of 585 steps, 8,958 of 9,011 tests pass, 53 skipped, 0 fail. zig fmt --check build.zig zig: clean. lean_check.py: 424, the same as main.

  • Not run: CUDA (only the denoiser has a CUDA port, idol: Idol's denoiser on CUDA (stacked on #1766; not yet run on a GPU) #1767, not yet run on a GPU). Metal M1 to M4: refused at load (no tensor units).

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (No token path.)
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (n/a.)
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (No new kernels; the combined SUMMARY is above.)
  • Prompt processing is no slower than the last release. (Not touched.)
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

Every pipeline of idol.metal and idol_attention.metal is registered: id_linear's eight forms (plain and GELU),
id_split's two, id_gemm32, id_norm (plain, affine, modulated), id_norm32 (affine, modulated), id_norm_split,
id_repo, id_qk, id_gate_add, id_embed, id_cond, id_bias32 and the attention's four forms. Each output is checked
at idol-sdk's fp16 points as the range its fp32 steps allow; the NAX projections against the tensor unit's
16-deep step bound, the attention against attention_ref's formula with P at the tensor unit's fp32. Inputs at the
large served shapes repeat at periods prime to every tile, so float64 stays cheap and every output is checked.

19 controls: rounding points dropped or added (gate 3), fp16 sums in id_norm32, id_repo, id_qk, id_embed and
id_gemm32, and the NAX fp32 path in id_gemm32's place (gate 4), and the attention's tail keys unmasked.
Gate 5 holds the 24 launches, each reserved by its source.

Known finding: id_qk gives a head holding a NaN +-inf in its other outputs where F.normalize gives NaN, as
Metal's max drops the NaN norm that clamp_min keeps.
…f-idol-vae and its case

Builds the denoiser's feature2 [4101, 128] from the preprocessed picture as idol-sdk's ImageEncoder does: the
encoder, quant_conv, the posterior's sample with the SDK's noise as an input, the 2 x 2 patches, the batch norm's
statistics and five zero rows. fp32 throughout on the SDK's bf16 weights (FLUX.2-klein-4B, or FLUX.2-dev rounded):
Winograd 3x3s, implicit-GEMM downsamples and 1x1s, two-pass GroupNorm, single-head attention.

idol_case.py vae records the SDK's MPS rows, its noise and a float64 ruler; tf-idol-vae gates worst row and
1 - cosine within twice the SDK's distance from float64 (ours 5.1e-5 and 5.0e-11 against 2.7e-2 and 5.1e-5), and
repeat bytes. All eight kernels register with the contract: 13 forms pass gates 1-4, 12 controls red.
The octree decoder (4 blocks) and Gaussian decoder (16 blocks) run natively on
Metal from the released checkpoint. Host-side pieces: systematic resampling with
a pluggable RNG (native PRNG, or the SDK's recorded draws), compaction, leaves
and jitter, activations, the PLY writer and the camera JSON. tf-idol-decode
decodes a latent and gates it against an idol_case.py `decode` case.

New kernels live in idol_decoder.metal: the point embedding, K/V and head
packing, an fp32 attention for <= 8 query rows, and a copy of the idol
attention whose P.V uses an fp16 hi/lo split. This removes the -2^-12 per-pass
bias that fp32 P picks up on the tensor unit.

Gates are measured against fp32, allowing twice the worst of: the SDK's own
distance, the fp16 floor, and seeded nudged SDK reruns. Further gates:
per-block teacher-forced checks, counts identical given the SDK's
probabilities and draws, activated fields, the PLY's columns, the camera, and
byte-identical repeats. The comparators are covered by plant-a-mismatch unit
tests.

idol_case.py gains `decode`, which pins decoder.py, gaussian.py and camera.py by
sha256 and records levels, draws, counts, fields, the PLY and timings.
metal_idol_decoder*.zig registers dec_embed, dec_half, dec_kv, dec_heads,
dec_attention_few and dec_attention_{00,01,10,11} as "metal idol decoder
...", against float64 at idol-sdk's fp32 and fp16 points. Fifteen controls
each go red at their gate: an added or dropped rounding point, fp16 sums,
P's lo half dropped (caught on heads built so O is P's fp16 remainder
alone), keys past Lk unmasked, an input written, and the two NaN rules
below undone. Gate-5 records are in idol_records.zig; the common Lib is now
generic over a source and its launchers.

Gate 1 found two real rule breaks, fixed with finite outputs unchanged
(the 262k and 32k decode cases give byte-identical PLYs and gate lines):
dec_heads clamped its norm by max, dropping a NaN; and the tiled attention
gave NaN for an infinite value, because an fp16 lo of 0 or of the opposite
sign met the infinity. The lo half now meets the values clamped to fp16's
finite range (about 2% at 262k, 7% at 1M).

docs/recipes/idol.md gains a Decoders part with the precision, gates,
run-a-case commands, measurements and the contract SUMMARY.
…idol_case.py dino

The denoiser's feature1 [4101, 1280] from the SDK's 1024 x 1024 image, as idol-sdk's ImageEncoder makes it:
transformers' DINOv3ViTModel in bf16, its last_hidden_state, then F.layer_norm in fp32. Weights from the ungated
timm/vit_huge_plus_patch16_dinov3.lvd1689m (DINOv3 License), rounded to bf16 at load. Patch conv as one GEMM, class
token and four registers, 2D RoPE on the patches (transformers' fp32 periods, tables bit-equal to torch's), 32 blocks of
20 heads of 64 with LayerScale and SwiGLU, rounded where torch on MPS rounds (measured): one rounding a linear or norm,
two for SiLU times up and LayerScale with the residual, three for the rotary. The attention rounds P to bf16 itself.

idol_case.py dino pins idol-sdk's conditioning.py, transformers 5.22.1's DINOv3 files and the weights by sha256, and
writes the SDK's feature1 on MPS, a float32 CPU ruler, every block's rows and each block in float32. With --official it
proves timm's weights Meta's: 550 tensors byte for byte, feature1 bit for bit.

On the avocado: feature1 worst row against float32 0.1353 (torch 0.1686), 1 - cos 5.37e-5 (5.57e-5); every block on
torch's own input within the 2x gate; repeat byte-identical; 0.120 s an encode against the SDK's 0.261 s on MPS.
…ngelog line

Every pipeline of idol_dino.metal is registered: dn_linear_0 and dn_linear_1, dn_norm, dn_final, dn_heads,
dn_scale_add, dn_swiglu and dn_attention, each checked at transformers' bf16 points. 11 controls: a bf16 rounding
added or dropped (gate 3), fp16 sums in dn_norm and dn_final and P V rounded to bf16 each key block (gate 4), the
attention's tail keys unmasked. Gate 5 holds the 8 launches, each reserved by its source.

SUMMARY 8 kernels: 8 pass gates 1-4, 0 fail, 0 known findings; 11 controls: 11 red as planted, 0 missed; 0 limit problems
# Conflicts:
#	CHANGELOG.md
#	docs/kernel-contract.md
#	docs/recipes/idol.md
#	tools/zig/idol_case.py
#	zig/build/metal_targets.zig
#	zig/kernels/metal/sources.zig
#	zig/src/core/contract/idol_records.zig
#	zig/src/families/idol/idol.zig
#	zig/tests/contract/metal_limits.zig
#	zig/tests/contract/metal_main.zig
# Conflicts:
#	CHANGELOG.md
#	docs/kernel-contract.md
#	docs/recipes/idol.md
#	tools/zig/idol_case.py
#	zig/build/metal_targets.zig
#	zig/kernels/metal/sources.zig
#	zig/src/families/idol/idol.zig
#	zig/tests/contract/metal_limits.zig
#	zig/tests/contract/metal_main.zig
…erged case tool stays under 600 lines

The records file's header comment is one line again.
…oders, the denoiser and the decoders to a PLY

Each stage is timed and each model freed after it. The starting noise is torch's CPU-seeded randn, reimplemented;
the VAE's posterior noise and the starting noise can be given for parity with the SDK's MPS draws. The recipe's
End to end section has the stage timings, memory and PLY comparison against the SDK's own IdolPipeline run.
@ashhart

ashhart commented Nov 4, 2026

Copy link
Copy Markdown
Owner

Thanks @jayleaton. The noise table settles the PLY question: with the SDK's noise injected, ours lands as close as sampling resolves. It waits on #1773, #1774 and #1775. Once they're on main, please rebase onto main, and we'll run tf-idol-generate end to end on an Apple M5 Ultra.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants