init: repo scaffolding, deps, prompts, plan
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
commit
897242d200
16
.gitignore
vendored
Normal file
16
.gitignore
vendored
Normal file
@ -0,0 +1,16 @@
|
||||
# Python-generated files
|
||||
__pycache__/
|
||||
*.py[oc]
|
||||
build/
|
||||
dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
|
||||
# Virtual environments
|
||||
.venv
|
||||
|
||||
# Model checkpoints (large binary files)
|
||||
checkpoints/
|
||||
|
||||
# Large generated fixtures (golden .npz kept small; regenerate with dump script)
|
||||
reference/fixtures/*.npz
|
||||
1
.python-version
Normal file
1
.python-version
Normal file
@ -0,0 +1 @@
|
||||
3.12
|
||||
41
CLAUDE.md
Normal file
41
CLAUDE.md
Normal file
@ -0,0 +1,41 @@
|
||||
# CLAUDE.md — corridorkey-mlx
|
||||
|
||||
MLX inference port of CorridorKey for Apple Silicon.
|
||||
|
||||
## Architecture
|
||||
|
||||
- Input: 4ch (RGB + coarse alpha hint)
|
||||
- Backbone: Hiera (timm, features_only=True) → 4 multiscale features
|
||||
- Decoder heads: alpha (1ch) + foreground (3ch), upsampled to full res
|
||||
- Refiner: CNN over RGB + coarse preds (7ch) → additive delta logits → sigmoid
|
||||
|
||||
## Repo layout
|
||||
|
||||
- `src/corridorkey_mlx/` — main package
|
||||
- `model/` — MLX model definitions
|
||||
- `convert/` — PyTorch→MLX weight conversion
|
||||
- `inference/` — inference pipeline
|
||||
- `io/` — image loading, saving, preprocessing
|
||||
- `utils/` — shared helpers, layout transforms
|
||||
- `scripts/` — CLI tools (dump reference, compare, bench)
|
||||
- `prompts/` — phased port instructions
|
||||
- `reference/` — PyTorch reference harness outputs
|
||||
- `tests/` — parity and unit tests
|
||||
|
||||
## Conventions
|
||||
|
||||
- Python 3.12+, uv for deps
|
||||
- ruff for lint/format, mypy for types, pytest for tests
|
||||
- MLX uses NHWC — centralize layout transforms in `utils/`
|
||||
- All non-trivial changes need a validation path
|
||||
- Inference only — no training code
|
||||
- Preserve PyTorch behavior before optimizing
|
||||
|
||||
## Commands
|
||||
|
||||
```bash
|
||||
uv run pytest # run tests
|
||||
uv run ruff check . # lint
|
||||
uv run ruff format . # format
|
||||
uv run mypy src/ # type check
|
||||
```
|
||||
101
README.md
Normal file
101
README.md
Normal file
@ -0,0 +1,101 @@
|
||||
# corridorkey-mlx
|
||||
|
||||
MLX inference port of [CorridorKey](https://github.com/nikopueringer/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
|
||||
|
||||
```bash
|
||||
uv sync --group dev
|
||||
```
|
||||
|
||||
For PyTorch reference work:
|
||||
```bash
|
||||
uv sync --group reference
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
```bash
|
||||
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:**
|
||||
```bash
|
||||
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.
|
||||
@ -0,0 +1,322 @@
|
||||
---
|
||||
title: "feat: CorridorKey MLX Inference Port"
|
||||
type: feat
|
||||
date: 2026-03-01
|
||||
---
|
||||
|
||||
# CorridorKey MLX Inference Port
|
||||
|
||||
## Overview
|
||||
|
||||
Port CorridorKey's inference pipeline from PyTorch to MLX for native Apple Silicon execution. The port is staged across 5 phases, each with its own parity gate against PyTorch reference outputs. No training code. No UI.
|
||||
|
||||
## Problem Statement
|
||||
|
||||
CorridorKey (alpha matting model) runs on PyTorch. On Apple Silicon, PyTorch uses MPS which is slower and less memory-efficient than MLX's native Metal backend. A clean MLX port enables fast local inference without PyTorch overhead.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
RGB + coarse alpha hint (4ch)
|
||||
|
|
||||
v
|
||||
+------------------+
|
||||
| Hiera backbone | timm, features_only=True
|
||||
| -> 4 multiscale | feature maps
|
||||
+------------------+
|
||||
|
|
||||
+----+----+
|
||||
v v
|
||||
+-------+ +-------+
|
||||
| Alpha | | FG |
|
||||
| head | | head |
|
||||
| (1ch) | | (3ch) |
|
||||
+-------+ +-------+
|
||||
| |
|
||||
+----+----+
|
||||
v
|
||||
+------------------+
|
||||
| CNN Refiner | RGB + coarse preds (7ch)
|
||||
| -> delta logits | -> sigmoid
|
||||
+------------------+
|
||||
|
|
||||
v
|
||||
final alpha + fg
|
||||
```
|
||||
|
||||
**Key model components:**
|
||||
- **Backbone:** Hiera (hierarchical vision transformer) -- 4 multiscale feature outputs
|
||||
- **Decoder heads:** Two parallel heads (alpha 1ch, foreground 3ch) consuming backbone features
|
||||
- **CNN Refiner:** Takes RGB + coarse predictions (7ch), predicts additive delta logits, final sigmoid
|
||||
|
||||
**MLX-specific considerations:**
|
||||
- MLX uses NHWC layout natively vs PyTorch's NCHW
|
||||
- Conv weight transpose needed: `(O,I,H,W)` -> `(O,H,W,I)`
|
||||
- GroupNorm behavior must match PyTorch exactly for parity
|
||||
- First conv is patched to 4 input channels (RGB + alpha hint)
|
||||
|
||||
## Implementation Phases
|
||||
|
||||
### Phase 1: PyTorch Reference Harness
|
||||
|
||||
**Prompt:** `prompts/phase1-backbone.md`
|
||||
|
||||
**Goal:** Deterministic reference pipeline that dumps intermediate tensors for staged MLX parity testing.
|
||||
|
||||
**Deliverables:**
|
||||
|
||||
- [x] `scripts/dump_pytorch_reference.py` -- loads checkpoint via state_dict, runs forward pass, saves intermediates
|
||||
- [x] `reference/fixtures/` -- sample inputs and golden outputs
|
||||
- [x] `tests/test_reference_fixtures.py` -- validates fixture shapes and dtypes
|
||||
|
||||
**Tensors to dump:**
|
||||
1. 4 backbone feature maps (multiscale)
|
||||
2. Alpha coarse logits
|
||||
3. FG coarse logits
|
||||
4. Alpha coarse probs
|
||||
5. FG coarse probs
|
||||
6. Refiner delta logits
|
||||
7. Final alpha
|
||||
8. Final FG
|
||||
|
||||
**Files touched:**
|
||||
|
||||
| File | Action |
|
||||
|------|--------|
|
||||
| `scripts/dump_pytorch_reference.py` | implement |
|
||||
| `reference/fixtures/*.npz` or `*.safetensors` | generate |
|
||||
| `tests/test_reference_fixtures.py` | implement (unskip) |
|
||||
| `README.md` | update fixture format docs |
|
||||
|
||||
**Constraints:**
|
||||
- Load via `state_dict`, not pickle
|
||||
- Deterministic (fixed seed, mode)
|
||||
- One tiny golden example checked in
|
||||
- Print shape report via rich
|
||||
|
||||
**Definition of done:**
|
||||
- Fixture files exist with all 8+ tensor groups
|
||||
- Shape contract tests pass
|
||||
- Script is idempotent
|
||||
|
||||
---
|
||||
|
||||
### Phase 2: MLX Decoder and Refiner Blocks
|
||||
|
||||
**Prompt:** `prompts/phase2-mlx-blocks.md`
|
||||
|
||||
**Goal:** MLX implementations of non-backbone blocks with parity tests against Phase 1 fixtures.
|
||||
|
||||
**Components to implement:**
|
||||
|
||||
- [ ] `MLP` -- feedforward block
|
||||
- [ ] `DecoderHead` -- consumes backbone features, produces coarse predictions
|
||||
- [ ] `RefinerBlock` -- single refiner stage
|
||||
- [ ] `CNNRefinerModule` -- full refiner consuming 7ch input
|
||||
|
||||
**Files touched:**
|
||||
|
||||
| File | Action |
|
||||
|------|--------|
|
||||
| `src/corridorkey_mlx/model/decoder.py` | implement MLP + DecoderHead |
|
||||
| `src/corridorkey_mlx/model/refiner.py` | implement RefinerBlock + CNNRefinerModule |
|
||||
| `src/corridorkey_mlx/utils/layout.py` | implement NCHW to NHWC transforms |
|
||||
| `tests/test_decoder_parity.py` | implement (unskip) |
|
||||
| `tests/test_refiner_parity.py` | implement (unskip) |
|
||||
|
||||
**Constraints:**
|
||||
- NHWC throughout, layout conversions only in `utils/layout.py`
|
||||
- GroupNorm must match PyTorch behavior exactly
|
||||
- Parity tests load saved backbone features (from Phase 1 fixtures) as input
|
||||
- Report max abs error and mean abs error
|
||||
|
||||
**Definition of done:**
|
||||
- Decoder parity test passes (max abs err < 1e-5)
|
||||
- Refiner parity test passes (max abs err < 1e-5)
|
||||
- Modules wired into partial model path for test usage
|
||||
|
||||
---
|
||||
|
||||
### Phase 3: Checkpoint Conversion
|
||||
|
||||
**Prompt:** `prompts/phase3-conversion.md`
|
||||
|
||||
**Goal:** Robust conversion pipeline from PyTorch checkpoint to MLX-compatible weights.
|
||||
|
||||
**Deliverables:**
|
||||
|
||||
- [ ] `src/corridorkey_mlx/convert/converter.py` -- key mapping + weight transforms
|
||||
- [ ] Conversion diagnostic report (source key -> dest key, shapes, transform applied)
|
||||
- [ ] Output as safetensors
|
||||
|
||||
**Files touched:**
|
||||
|
||||
| File | Action |
|
||||
|------|--------|
|
||||
| `src/corridorkey_mlx/convert/converter.py` | implement |
|
||||
| `src/corridorkey_mlx/convert/__init__.py` | export converter |
|
||||
| `tests/test_conversion.py` | implement (unskip) |
|
||||
| `scripts/convert_weights.py` | create (CLI wrapper) |
|
||||
|
||||
**Key mapping concerns:**
|
||||
- Explicit key-by-key mapping, no regex guessing
|
||||
- Conv weights: `(O,I,H,W)` -> `(O,H,W,I)`
|
||||
- Patched 4-channel first conv must be preserved exactly
|
||||
- No silent fallbacks -- all mismatches are errors
|
||||
|
||||
**Diagnostic output per key:**
|
||||
```
|
||||
source_key -> dest_key | src_shape -> dst_shape | transform
|
||||
```
|
||||
|
||||
**Definition of done:**
|
||||
- Converter script exists and runs
|
||||
- Mapping is explicit and auditable
|
||||
- Shape validation passes for all completed modules
|
||||
- No orphan keys (every source key maps or is explicitly skipped with reason)
|
||||
|
||||
---
|
||||
|
||||
### Phase 4: Full Inference Pipeline
|
||||
|
||||
**Prompt:** `prompts/phase4-inference-pipeline.md` (not yet written)
|
||||
|
||||
**Goal:** End-to-end inference matching PyTorch output.
|
||||
|
||||
**Deliverables:**
|
||||
|
||||
- [ ] `src/corridorkey_mlx/model/backbone.py` -- Hiera MLX port
|
||||
- [ ] `src/corridorkey_mlx/model/corridorkey.py` -- full model composition
|
||||
- [ ] `src/corridorkey_mlx/inference/pipeline.py` -- load, preprocess, forward, postprocess, save
|
||||
- [ ] `src/corridorkey_mlx/io/image.py` -- PIL-based image I/O + preprocessing
|
||||
- [ ] End-to-end parity test against PyTorch golden output
|
||||
|
||||
**Files touched:**
|
||||
|
||||
| File | Action |
|
||||
|------|--------|
|
||||
| `src/corridorkey_mlx/model/backbone.py` | implement Hiera MLX |
|
||||
| `src/corridorkey_mlx/model/corridorkey.py` | implement full model |
|
||||
| `src/corridorkey_mlx/inference/pipeline.py` | implement |
|
||||
| `src/corridorkey_mlx/io/image.py` | implement |
|
||||
| `tests/test_e2e_parity.py` | create |
|
||||
| `main.py` | wire CLI via typer |
|
||||
|
||||
**Definition of done:**
|
||||
- `uv run python main.py --image input.png --output output.png` produces correct result
|
||||
- E2e parity test passes within tolerance
|
||||
- Memory usage is reasonable (no unnecessary copies)
|
||||
|
||||
---
|
||||
|
||||
### Phase 5: Optimization and Benchmarking
|
||||
|
||||
**Prompt:** `prompts/phase5-optimization.md` (not yet written)
|
||||
|
||||
**Goal:** Production-quality performance on Apple Silicon.
|
||||
|
||||
**Deliverables:**
|
||||
|
||||
- [ ] `scripts/bench_mlx.py` -- latency, throughput, memory reporting
|
||||
- [ ] `scripts/compare_reference.py` -- side-by-side output comparison
|
||||
- [ ] Performance optimizations (compile, memory layout, batching)
|
||||
|
||||
**Potential optimizations:**
|
||||
- `mx.compile()` on hot paths
|
||||
- Memory-efficient attention if Hiera benefits
|
||||
- Avoid unnecessary `mx.eval()` calls (lazy graph)
|
||||
- Optimal image preprocessing (avoid numpy round-trips)
|
||||
|
||||
**Definition of done:**
|
||||
- Benchmark script reports latency + peak memory
|
||||
- Performance is competitive with PyTorch MPS on same hardware
|
||||
- No correctness regression (parity tests still pass)
|
||||
|
||||
---
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
### Functional Requirements
|
||||
|
||||
- [ ] All 5 phases complete with passing parity gates
|
||||
- [ ] Single-image inference CLI works end-to-end
|
||||
- [ ] Output visually matches PyTorch reference
|
||||
|
||||
### Quality Gates (per phase)
|
||||
|
||||
- [ ] Parity tests pass within documented tolerance
|
||||
- [ ] No skipped tests without explicit phase rationale
|
||||
- [ ] ruff + mypy clean
|
||||
- [ ] Fixtures are deterministic and reproducible
|
||||
|
||||
## Dependencies and Prerequisites
|
||||
|
||||
| Dependency | Phase | Notes |
|
||||
|-----------|-------|-------|
|
||||
| Original CorridorKey checkpoint | 1 | needed for reference dump |
|
||||
| Original CorridorKey source code | 1 | needed to understand architecture |
|
||||
| PyTorch + timm (reference group) | 1 | `uv sync --group reference` |
|
||||
| Phase 1 fixtures | 2, 3 | parity test inputs |
|
||||
| Phase 2 + 3 complete | 4 | blocks + weights needed for full model |
|
||||
|
||||
## Risk Analysis
|
||||
|
||||
| Risk | Impact | Mitigation |
|
||||
|------|--------|------------|
|
||||
| Hiera has no existing MLX port | High -- Phase 4 blocker | Research early; may need full rewrite |
|
||||
| GroupNorm numerical differences | Medium -- parity failures | Test tolerance tuning; document acceptable drift |
|
||||
| 4ch first conv patching | Low -- conversion bug | Explicit test in Phase 3 |
|
||||
| MLX API changes | Low -- version pinned | Pin `mlx>=0.31.0` in pyproject |
|
||||
|
||||
## Current Repo State
|
||||
|
||||
```
|
||||
src/corridorkey_mlx/
|
||||
__init__.py [exists]
|
||||
model/
|
||||
backbone.py [placeholder]
|
||||
decoder.py [placeholder]
|
||||
refiner.py [placeholder]
|
||||
corridorkey.py [placeholder]
|
||||
convert/
|
||||
converter.py [placeholder]
|
||||
inference/
|
||||
pipeline.py [placeholder]
|
||||
io/
|
||||
image.py [placeholder]
|
||||
utils/
|
||||
layout.py [placeholder]
|
||||
|
||||
scripts/
|
||||
dump_pytorch_reference.py [placeholder]
|
||||
compare_reference.py [placeholder]
|
||||
bench_mlx.py [placeholder]
|
||||
|
||||
tests/
|
||||
test_import.py [passes]
|
||||
test_reference_fixtures.py [skipped - Phase 1]
|
||||
test_decoder_parity.py [skipped - Phase 2]
|
||||
test_refiner_parity.py [skipped - Phase 2]
|
||||
test_conversion.py [skipped - Phase 3]
|
||||
|
||||
prompts/
|
||||
phase1-backbone.md [written]
|
||||
phase2-mlx-blocks.md [written]
|
||||
phase3-conversion.md [written]
|
||||
phase4-inference-pipeline.md [empty]
|
||||
phase5-optimization.md [empty]
|
||||
```
|
||||
|
||||
## Open Questions
|
||||
|
||||
- Hiera port strategy: full rewrite vs adapting existing MLX vision transformer code?
|
||||
- Acceptable e2e parity tolerance? (likely max abs < 1e-3 due to float32 Metal vs CUDA differences)
|
||||
- Where is the original CorridorKey checkpoint hosted? (needed for Phase 1)
|
||||
- Original CorridorKey source -- is it the nicehash/CorridorKey repo or a fork?
|
||||
|
||||
## References
|
||||
|
||||
- Original repo: nicehash/CorridorKey on GitHub
|
||||
- MLX: ml-explore/mlx on GitHub
|
||||
- timm Hiera: `timm.create_model("hiera_...", features_only=True)`
|
||||
- Phase prompts: `prompts/phase{1-5}-*.md`
|
||||
6
main.py
Normal file
6
main.py
Normal file
@ -0,0 +1,6 @@
|
||||
def main():
|
||||
print("Hello from corridorkey-mlx!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
39
prompts/phase1-backbone.md
Normal file
39
prompts/phase1-backbone.md
Normal file
@ -0,0 +1,39 @@
|
||||
Work only on the PyTorch reference harness.
|
||||
|
||||
Goal:
|
||||
Create a deterministic reference pipeline that loads the original CorridorKey checkpoint and dumps intermediate tensors needed for staged MLX parity.
|
||||
|
||||
Deliverables:
|
||||
|
||||
- scripts/dump_pytorch_reference.py
|
||||
- reference/fixtures/ sample inputs and outputs
|
||||
- tests that validate fixture generation shape contracts
|
||||
- README updates describing the fixture format
|
||||
|
||||
Requirements:
|
||||
|
||||
- Load the model via state_dict, not entire-model pickle semantics.
|
||||
- Save:
|
||||
- 4 backbone feature maps
|
||||
- alpha coarse logits
|
||||
- fg coarse logits
|
||||
- alpha coarse probs
|
||||
- fg coarse probs
|
||||
- delta logits
|
||||
- final alpha
|
||||
- final fg
|
||||
- Make fixture generation deterministic where practical.
|
||||
- Keep one tiny golden example checked in.
|
||||
- Print a concise shape report.
|
||||
|
||||
Do not:
|
||||
|
||||
- start MLX implementation
|
||||
- refactor unrelated files
|
||||
- add training code
|
||||
|
||||
Before editing:
|
||||
|
||||
- inspect the original CorridorKey model code carefully
|
||||
- summarize exact tensors that will be dumped
|
||||
- define file naming and serialization format first
|
||||
29
prompts/phase2-mlx-blocks.md
Normal file
29
prompts/phase2-mlx-blocks.md
Normal file
@ -0,0 +1,29 @@
|
||||
Work only on the MLX implementations of the custom non-backbone blocks.
|
||||
|
||||
Goal:
|
||||
Implement MLX versions of:
|
||||
|
||||
- MLP
|
||||
- DecoderHead
|
||||
- RefinerBlock
|
||||
- CNNRefinerModule
|
||||
|
||||
Requirements:
|
||||
|
||||
- Use MLX idioms and explicit NHWC handling.
|
||||
- Centralize tensor layout transforms in one utility module.
|
||||
- Use pytorch-compatible GroupNorm behavior where needed for parity.
|
||||
- Write parity tests that use saved PyTorch backbone features and saved coarse predictions.
|
||||
- Report max abs error and mean abs error in test output or helper scripts.
|
||||
|
||||
Do not:
|
||||
|
||||
- port Hiera yet
|
||||
- optimize prematurely
|
||||
- spread layout conversions across many files
|
||||
|
||||
Definition of done:
|
||||
|
||||
- decoder parity test exists
|
||||
- refiner parity test exists
|
||||
- modules are wired into a partial model path for test usage
|
||||
31
prompts/phase3-conversion.md
Normal file
31
prompts/phase3-conversion.md
Normal file
@ -0,0 +1,31 @@
|
||||
Work only on checkpoint conversion.
|
||||
|
||||
Goal:
|
||||
Create a robust conversion pipeline from the PyTorch CorridorKey checkpoint to MLX-compatible weights.
|
||||
|
||||
Requirements:
|
||||
|
||||
- Inspect state_dict keys and map them explicitly.
|
||||
- Convert conv weights from PyTorch layout to MLX layout.
|
||||
- Preserve the patched 4-channel first conv behavior exactly.
|
||||
- Write converter diagnostics:
|
||||
- source key
|
||||
- destination key
|
||||
- source shape
|
||||
- destination shape
|
||||
- transform applied
|
||||
- Save output as safetensors or npz.
|
||||
- Validate load with MLX strict loading wherever possible.
|
||||
|
||||
Do not:
|
||||
|
||||
- attempt full end-to-end model parity yet if Hiera is incomplete
|
||||
- hide key mismatches
|
||||
- use silent fallbacks
|
||||
|
||||
Definition of done:
|
||||
|
||||
- converter script exists
|
||||
- mapping file exists
|
||||
- shape validation passes for completed modules
|
||||
- conversion report is readable and auditable
|
||||
0
prompts/phase4-inference-pipeline.md
Normal file
0
prompts/phase4-inference-pipeline.md
Normal file
0
prompts/phase5-optimization.md
Normal file
0
prompts/phase5-optimization.md
Normal file
51
pyproject.toml
Normal file
51
pyproject.toml
Normal file
@ -0,0 +1,51 @@
|
||||
[project]
|
||||
name = "corridorkey-mlx"
|
||||
version = "0.1.0"
|
||||
description = "MLX inference port of CorridorKey for Apple Silicon"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"mlx>=0.31.0",
|
||||
"numpy>=2.4.2",
|
||||
"pillow>=12.1.1",
|
||||
"pydantic>=2.12.5",
|
||||
"rich>=14.3.3",
|
||||
"safetensors>=0.7.0",
|
||||
"typer>=0.24.1",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"mypy>=1.19.1",
|
||||
"pytest>=9.0.2",
|
||||
"pytest-xdist>=3.8.0",
|
||||
"ruff>=0.15.4",
|
||||
]
|
||||
reference = [
|
||||
"timm>=1.0.25",
|
||||
"torch>=2.10.0",
|
||||
"torchvision>=0.25.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/corridorkey_mlx"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py312"
|
||||
line-length = 99
|
||||
src = ["src"]
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "UP", "B", "SIM", "TCH"]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.12"
|
||||
strict = true
|
||||
mypy_path = "src"
|
||||
3
src/corridorkey_mlx/__init__.py
Normal file
3
src/corridorkey_mlx/__init__.py
Normal file
@ -0,0 +1,3 @@
|
||||
"""CorridorKey MLX — inference port for Apple Silicon."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
1
src/corridorkey_mlx/convert/__init__.py
Normal file
1
src/corridorkey_mlx/convert/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""PyTorch to MLX weight conversion utilities."""
|
||||
9
src/corridorkey_mlx/convert/converter.py
Normal file
9
src/corridorkey_mlx/convert/converter.py
Normal file
@ -0,0 +1,9 @@
|
||||
"""PyTorch → MLX weight converter (not yet implemented).
|
||||
|
||||
Responsibilities:
|
||||
- Map state_dict keys between PyTorch and MLX naming
|
||||
- Transpose conv weights from NCHW → NHWC
|
||||
- Handle patched 4-channel first conv
|
||||
- Emit diagnostic report (src key, dst key, shapes, transform)
|
||||
- Save as safetensors
|
||||
"""
|
||||
1
src/corridorkey_mlx/inference/__init__.py
Normal file
1
src/corridorkey_mlx/inference/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""Inference pipeline for CorridorKey MLX."""
|
||||
4
src/corridorkey_mlx/inference/pipeline.py
Normal file
4
src/corridorkey_mlx/inference/pipeline.py
Normal file
@ -0,0 +1,4 @@
|
||||
"""Inference pipeline (not yet implemented).
|
||||
|
||||
Orchestrates: load image → preprocess → model forward → postprocess → save.
|
||||
"""
|
||||
1
src/corridorkey_mlx/io/__init__.py
Normal file
1
src/corridorkey_mlx/io/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""Image I/O, preprocessing, and postprocessing."""
|
||||
1
src/corridorkey_mlx/io/image.py
Normal file
1
src/corridorkey_mlx/io/image.py
Normal file
@ -0,0 +1 @@
|
||||
"""Image loading, saving, and preprocessing (not yet implemented)."""
|
||||
1
src/corridorkey_mlx/model/__init__.py
Normal file
1
src/corridorkey_mlx/model/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""Model definitions for CorridorKey MLX."""
|
||||
5
src/corridorkey_mlx/model/backbone.py
Normal file
5
src/corridorkey_mlx/model/backbone.py
Normal file
@ -0,0 +1,5 @@
|
||||
"""Hiera backbone — MLX port (not yet implemented).
|
||||
|
||||
Original: timm Hiera with features_only=True.
|
||||
Emits 4 multiscale feature maps.
|
||||
"""
|
||||
4
src/corridorkey_mlx/model/corridorkey.py
Normal file
4
src/corridorkey_mlx/model/corridorkey.py
Normal file
@ -0,0 +1,4 @@
|
||||
"""Top-level CorridorKey model — MLX port (not yet implemented).
|
||||
|
||||
Composes backbone + decoder heads + refiner into full pipeline.
|
||||
"""
|
||||
5
src/corridorkey_mlx/model/decoder.py
Normal file
5
src/corridorkey_mlx/model/decoder.py
Normal file
@ -0,0 +1,5 @@
|
||||
"""Decoder heads — MLX port (not yet implemented).
|
||||
|
||||
Two heads: alpha (1ch) and foreground (3ch).
|
||||
Consume multiscale backbone features, upsample to full resolution.
|
||||
"""
|
||||
5
src/corridorkey_mlx/model/refiner.py
Normal file
5
src/corridorkey_mlx/model/refiner.py
Normal file
@ -0,0 +1,5 @@
|
||||
"""CNN refiner — MLX port (not yet implemented).
|
||||
|
||||
Input: RGB + coarse predictions (7ch total).
|
||||
Output: additive delta logits → final sigmoid.
|
||||
"""
|
||||
1
src/corridorkey_mlx/utils/__init__.py
Normal file
1
src/corridorkey_mlx/utils/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""Shared utilities and tensor layout helpers."""
|
||||
4
src/corridorkey_mlx/utils/layout.py
Normal file
4
src/corridorkey_mlx/utils/layout.py
Normal file
@ -0,0 +1,4 @@
|
||||
"""Tensor layout conversion utilities (not yet implemented).
|
||||
|
||||
Centralizes NCHW ↔ NHWC transforms. All layout conversions go through here.
|
||||
"""
|
||||
Loading…
Reference in New Issue
Block a user