c2e7967538
* feat(phase4a): MLX Hiera backbone with unroll/reroll parity Port timm hiera_base_plus_224 to MLX: PatchEmbed, MaskUnitAttention, HieraBlock, unroll/reroll, HieraBackbone. 4 NHWC feature maps at strides 4/8/16/32. backbone.py now re-exports from hiera.py. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat(phase4c): weight loading + pos_embed bicubic interpolation HieraBackbone.load_checkpoint() loads safetensors, strips encoder.model. prefix, bicubic-interpolates pos_embed from 512x512 to 128x128 tokens. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat(phase4d): shape + parity tests, fix attention transpose Fix MaskUnitAttention output transpose (0,2,3,1,4 → 0,3,2,1,4) to match PyTorch transpose(1,3) token ordering for windowed attention. Parity results (4 stages): Stage 0: max_abs 2.9e-4, mean 1.6e-5 Stage 1: max_abs 1.3e-4, mean 1.3e-5 Stage 2: max_abs 1.1e-2, mean 5.2e-5 (16 blocks, expected drift) Stage 3: max_abs 5.8e-4, mean 2.6e-5 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs: clarify pos_embed parameter intent in HieraBackbone Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|---|---|---|
| docs/plans | ||
| prompts | ||
| scripts | ||
| src/corridorkey_mlx | ||
| tests | ||
| .gitignore | ||
| .python-version | ||
| CLAUDE.md | ||
| main.py | ||
| pyproject.toml | ||
| README.md | ||
| uv.lock | ||
corridorkey-mlx
MLX inference port of CorridorKey for Apple Silicon.
Architecture
RGB image + coarse alpha hint (4ch)
│
▼
┌──────────────────┐
│ Hiera backbone │ (timm, features_only)
│ → 4 multiscale │
│ feature maps │
└──────────────────┘
│
┌────┴────┐
▼ ▼
┌───────┐ ┌───────┐
│ Alpha │ │ FG │
│ head │ │ head │
│ (1ch) │ │ (3ch) │
└───────┘ └───────┘
│ │
└────┬────┘
▼
┌──────────────────┐
│ CNN Refiner │ RGB + coarse preds (7ch)
│ → delta logits │ → sigmoid
└──────────────────┘
│
▼
final alpha + fg
Phased Roadmap
| Phase | Scope | Status |
|---|---|---|
| 1 | PyTorch reference harness + fixture dump | In progress |
| 2 | MLX decoder/refiner blocks + parity tests | Not started |
| 3 | Checkpoint conversion (PyTorch → MLX) | Not started |
| 4 | Full inference pipeline | Not started |
| 5 | Optimization + benchmarking | Not started |
See prompts/ for detailed phase instructions.
Setup
uv sync --group dev
For PyTorch reference work:
uv sync --group reference
Development
uv run pytest # tests
uv run ruff check . # lint
uv run ruff format . # format
uv run mypy src/ # type check
Reference Fixtures
Phase 1 generates golden reference tensors from PyTorch for MLX parity testing.
Format: single reference/fixtures/golden.npz (numpy compressed archive)
Generate:
uv run --group reference python scripts/dump_pytorch_reference.py \
--checkpoint checkpoints/CorridorKey_v1.0.pth
Contents (all float32, NCHW, batch=1, img_size=512):
| Key | Shape | Description |
|---|---|---|
input |
(1, 4, 512, 512) | Random input (seed=42) |
encoder_feature_0 |
(1, 112, 128, 128) | Backbone stride-4 |
encoder_feature_1 |
(1, 224, 64, 64) | Backbone stride-8 |
encoder_feature_2 |
(1, 448, 32, 32) | Backbone stride-16 |
encoder_feature_3 |
(1, 896, 16, 16) | Backbone stride-32 |
alpha_logits |
(1, 1, 128, 128) | Alpha decoder output (H/4) |
fg_logits |
(1, 3, 128, 128) | FG decoder output (H/4) |
alpha_logits_up |
(1, 1, 512, 512) | Alpha logits upsampled |
fg_logits_up |
(1, 3, 512, 512) | FG logits upsampled |
alpha_coarse |
(1, 1, 512, 512) | sigmoid(alpha_logits_up) |
fg_coarse |
(1, 3, 512, 512) | sigmoid(fg_logits_up) |
delta_logits |
(1, 4, 512, 512) | Refiner output (10x scaled) |
alpha_final |
(1, 1, 512, 512) | Final alpha prediction |
fg_final |
(1, 3, 512, 512) | Final FG prediction |
Current Status
Phase 1 in progress — reference harness and fixture dump implemented.