Commit Graph

44 Commits

Author SHA1 Message Date
cmoyates
31d83e218b
feat: GPU-side ImageNet normalize + concat for full-frame path
PIL bicubic vs MLX cubic diff too large (0.22-0.59) — resize stays
on CPU. Only normalize+concat moved to MLX. Tiled path unchanged.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-09 17:06:47 -02:30
cmoyates
f80bec92b6
feat: replace manual attention with mx.fast.scaled_dot_product_attention
Fold window dim into batch (transpose before reshape), call SDPA,
unfold back. All parity tests pass with 0 regression.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-09 17:04:51 -02:30
cmoyates
6787a7fd21
feat: stage-boundary GC between backbone/decoder/refiner
mx.eval + gc + clear_cache at backbone→decoder and decoder→refiner
boundaries in eager mode. _compiled flag skips GC under mx.compile.
Compilation tests updated to use compile_model() API.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-09 17:02:38 -02:30
cmoyates
ce30eddcff
feat: slim forward mode — drop intermediate tensor refs
slim=True returns 4-key dict, engine uses it, pipeline.infer keeps full.
2 new tests: key count + bit-exact match vs full output.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-09 17:00:21 -02:30
cmoyates
4d7a29c7c7
feat: add mx.clear_cache() to engine cleanup after postprocessing
Del MLX array refs + gc + cache clear after numpy extraction.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-09 17:00:16 -02:30
cmoyates
ee23f27735
feat: add tiled inference support to engine
tile_size/overlap params on CorridorKeyMLXEngine — model loads at
tile_size, input stays full-res, forward uses tiled_inference w/ GC.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-08 22:41:45 -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
f73aabf5af
chore: default bf16+fused decode on, add benchmark results to plan
load_model() now defaults to dtype=bf16, fused_decode=True. Both are
free (zero parity regression, bit-exact fused path). Backbone/sigmoid
stay fp32.

Plan updated with benchmark results: tiled+GC = 12x peak memory
reduction at 2048x2048 (27.6GB → 2.3GB), all acceptance criteria met.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-08 21:27:56 -02:30
cmoyates
b3b7df271d
feat: MLX memory optimizations — bf16, fused decode, deterministic GC
Phase 1: selective bf16 mixed precision (backbone fp32, decoders bf16,
sigmoid always fp32). dtype param on GreenFormer, all outputs guaranteed
fp32.

Phase 2: FusedDecoderPair batches 3 upsample ops by concatenating
alpha+fg projections along channel axis. Bit-exact parity with unfused
path.

Phase 3: deterministic GC pipeline in tile loop (del + gc.collect +
mx.clear_cache) and engine.py intermediate tensor cleanup.

7 new tests: bf16 dtype/range/nan, fused-vs-unfused parity, fp32
default unchanged. 80 passed, 0 failures.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-08 21:21:01 -02:30
cmoyates
830e88f050
docs: add brainstorm + plan for MLX memory optimizations
bf16 mixed precision, batched decoder upsample, deterministic GC pipeline.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-08 21:12:59 -02:30
cmoyates
838bd0dd39
refactor(tests): restructure suite — 15 files to 8, consolidate parity
- Add conftest.py: shared paths, tolerances, skip markers, fixtures
- Consolidate imports/shapes/smoke/forward → test_model_contract.py
- Consolidate all parity (decoder/refiner/backbone/e2e) → test_parity.py
- Fold 2048 slow test into test_engine.py
- Rework test_conversion.py to use public convert_checkpoint() API
- Simplify test_weights.py: drop CLI/env var tests, keep checksum+config
- Rename tiling/compilation test files for consistency
- Delete 10 redundant files

80 tests collected (75 pass, 4 skip, 1 deselected)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-03 09:58:53 -03:30
cmoyates
62645a2e93
chore: add TDD skill definitions
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-03 09:45:43 -03:30
cmoyates
88a0cd9cb8
docs: add deep modules refactor plan
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-03 09:45:43 -03:30
cmoyates
1a992a853d
chore: ignore safetensors files in gitignore
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-03 09:45:43 -03:30
cmoyates
e95bef4716
feat(weights): add CLI for downloading weights from GitHub Releases
Streaming download w/ rich progress, SHA256 verification, platformdirs
caching. Supports --tag, --asset, --force, --print-path flags + env
var overrides. Console script: corridorkey-weights.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-03 09:45:43 -03: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
db2a8cecf0
chore: migrate type checker from mypy to ty
Replace mypy with Astral's ty for faster type checking. Fix
Image.BICUBIC deprecation and inline **kwargs for ty compatibility.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 22:25:14 -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
9cd800d37b
fix: reverse decoder concat order to match Torch (c4,c3,c2,c1)
Torch decoder concatenates feature projections as [c4,c3,c2,c1] but
MLX was using [c1,c2,c3,c4]. The linear_fuse conv was trained with
the Torch order, producing scrambled features. Verified fix gives
correlation=1.0 against Torch on identical inputs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 09:44:41 -03:30
Cristopher Yates
81967ccc7c
Merge pull request #7 from cmoyates/feat/smoke-2048
feat: add 2048 smoke test
2026-03-01 07:08:22 -03:30
cmoyates
98e02db630
chore: gitignore output/ directory
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 07:06:01 -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
Cristopher Yates
6e427fde99
Merge pull request #6 from cmoyates/feat/engine-integration-surface
feat: add CorridorKeyMLXEngine integration surface
2026-03-01 06:47:48 -03:30
cmoyates
6c711e85bd
fix: remove double normalize_rgb in engine preprocessing
preprocess() already calls normalize_rgb() internally; calling it
before passing to preprocess() applied ImageNet normalization twice.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 06:47:22 -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
Cristopher Yates
b0c7d8086c
Merge pull request #5 from cmoyates/feat/phase6-optimization
feat(phase6): optimization, benchmarking, and tiled inference
2026-03-01 06:27:57 -03:30
cmoyates
3d330d0e0b
fix: guard against overlap >= tile_size in tiled inference
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 06:27:31 -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
Cristopher Yates
a14f8a74e8
Merge pull request #4 from cmoyates/feat/phase5-assembly-parity
feat(phase5): full model assembly + e2e parity
2026-03-01 06:12:04 -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
fd6ca78935
Merge pull request #3 from cmoyates/feat/phase4a-hiera-backbone
feat(phase4): MLX Hiera backbone port
2026-03-01 05:59:19 -03:30
Cristopher Yates
c2e7967538
feat(phase4): MLX Hiera backbone port (#2)
* 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>
2026-03-01 05:55:24 -03:30
cmoyates
ac4393fcb4
docs: clarify pos_embed parameter intent in HieraBackbone
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:55:01 -03:30
cmoyates
cc2d31b435
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>
2026-03-01 05:51:36 -03:30
cmoyates
0d6d41cb45
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>
2026-03-01 05:48:40 -03:30
cmoyates
853efbe2cf
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>
2026-03-01 05:45:43 -03:30
cmoyates
7315a082df
chore: add phase 4-6 prompts, fix naming, note zoxide cd issue
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:20:40 -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
cmoyates
897242d200
init: repo scaffolding, deps, prompts, plan
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 04:48:51 -03:30