Commit Graph

14 Commits

Author SHA1 Message Date
cmoyates
d13ba1639d
feat: add optimization toggle flags + benchmark matrix script
use_sdpa, stage_gc flags propagated through backbone/model/pipeline.
bench_optimizations.py runs exhaustive/ablation/key sweep of all 6
toggles (slim, stage_gc, sdpa, bf16, fused_decode, gpu_preprocess)
with latency + peak memory reporting.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-09 17:11:28 -02:30
cmoyates
f30ab9255b
chore: fix lint, format, type errors
- ruff format: 2 files reformatted
- ruff check: fix import sort in scripts/infer_pytorch.py
- ty: suppress tree_flatten arg-type (MLX stub imprecision) and
  optional torch import

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-08 21:29:06 -02:30
cmoyates
7f1c2ee221
fix: match original CorridorKey decoder concat order [c4,c3,c2,c1]
PyTorch reference script had wrong concat order (forward) vs the
original nikopueringer/CorridorKey which uses torch.cat([c4,c3,c2,c1]).
Golden fixtures regenerated. Relax 3 tight e2e tolerances for Metal
vs CPU float32 drift.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 22:58:08 -03:30
cmoyates
b95c1f9e31
fix: pass img_size to GreenFormer + gitignore .DS_Store
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 09:46:31 -03:30
cmoyates
8be8d003e7
feat: add 2048 smoke test for native-resolution validation
Script + pytest tests to verify MLX inference works at CorridorKey's
training resolution. Uses samples/ by default, synthetic fallback.
Reports timing, peak memory, output diagnostics.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 07:04:28 -03:30
cmoyates
a3e723791f
add PyTorch reference inference script + gitignore samples/
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 06:56:04 -03:30
cmoyates
4bad586112
feat: add CorridorKeyMLXEngine integration surface
Drop-in MLX backend for main CorridorKey repo. Engine class wraps
existing model/inference with process_frame() API returning alpha,
fg, comp, processed as uint8 numpy arrays. Lowers Python to >=3.11,
removes unused deps, adds smoke script and 16 contract tests.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 06:43:07 -03:30
cmoyates
57c3aff913
feat(phase6): optimization, benchmarking, and tiled inference
- Cache nn.Upsample instances in DecoderHead/GreenFormer __init__
  (eliminated ~7 allocations per forward pass)
- Add mx.compile() support via load_model(compile=True)
- Benchmark harness: eager vs compiled, multi-resolution, parity checks
- Tiled inference with overlap blending for large images
- Profiling utilities with forced mx.eval for accurate timing
- Reference comparison script (scripts/compare_reference.py)
- 12 new tests (compiled consistency + tiling)
- README performance section with Apple Silicon guidance

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 06:24:10 -03:30
cmoyates
329db3fc43
fix: resize inputs to model img_size, use infer_and_save in CLI
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 06:11:57 -03:30
cmoyates
0481955cfd
feat(phase5): full model assembly + e2e parity
Wire GreenFormer (backbone + decoders + refiner), add image I/O,
inference pipeline, CLI entry point, and 3 test files (23 new tests).

E2e parity vs PyTorch: max abs err < 1.5e-4 across all outputs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 06:07:43 -03:30
Cristopher Yates
150c974eed
feat(phase3): PyTorch→MLX weight converter (#1)
* feat(phase3): PyTorch→MLX weight converter + safetensors output

365 keys mapped (367 source - 2 num_batches_tracked skipped).
Conv weights transposed (O,I,H,W→O,H,W,I), refiner stem remapped,
4ch patch embed preserved. 12 conversion tests + diagnostic report.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix(phase3): conv allowlist, unused var, test fixture

- Replace _is_conv_weight heuristic with explicit CONV_WEIGHT_KEYS frozenset (15 keys)
- Remove unused `skipped` list from convert_state_dict
- Module-scoped pytest fixture eliminates 11 redundant checkpoint loads

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:16:56 -03:30
cmoyates
091261ccff
feat(phase2): MLX decoder + refiner with parity tests
- MLP, DecoderHead (SegFormer-style), RefinerBlock, CNNRefinerModule
- All NHWC native, GroupNorm with pytorch_compatible=True
- Layout utils: NCHW<->NHWC, conv weight transpose (OIHW->OHWI)
- Dump script extended to export decoder/refiner weights
- Parity: decoder max_abs ~5e-6, refiner max_abs ~1e-5

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:00:00 -03:30
cmoyates
393dd0d5a7
chore: add remaining test stubs and script placeholders
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 04:49:09 -03:30
cmoyates
9b07961619
feat(phase1): PyTorch reference harness + fixture tests
Dump script loads CorridorKey checkpoint, runs deterministic forward
pass (seed=42, 512x512), saves 14 intermediate tensors to golden.npz.
Includes pos_embed bicubic interpolation for resolution mismatch.
32 shape/dtype/range validation tests all passing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 04:49:01 -03:30