diff --git a/docs/plans/2026-03-03-refactor-deep-modules-plan.md b/docs/plans/2026-03-03-refactor-deep-modules-plan.md new file mode 100644 index 0000000..d1c68dc --- /dev/null +++ b/docs/plans/2026-03-03-refactor-deep-modules-plan.md @@ -0,0 +1,295 @@ +--- +title: Deep Modules Refactor +type: refactor +date: 2026-03-03 +--- + +# Deep Modules Refactor + +## Overview + +Restructure corridorkey-mlx into "deep modules" with explicit public interfaces (`__all__` gateways in `__init__.py`), `_testing` subpackages for parity test access to internals, and `import-linter` enforcing cross-package boundaries. + +**Guiding principle:** each subpackage exposes a simple interface via `__init__.py`. Implementation files are internal. Tests access internals through `_testing` re-export modules. Import-linter enforces boundaries between top-level packages only (within a package, free imports). + +## Problem Statement + +All subpackage `__init__.py` files (except `convert/`) are empty. Consumers must know internal filenames (`from corridorkey_mlx.model.corridorkey import GreenFormer`). No tooling enforces boundaries. An AI agent dropping into this codebase has no way to discover a module's API without reading every file. + +## Proposed Module Interfaces + +### `model/` — Deep Module + +```python +# model/__init__.py +from corridorkey_mlx.model.corridorkey import GreenFormer + +__all__ = ["GreenFormer"] +``` + +Internal: `backbone.py`, `hiera.py`, `decoder.py`, `refiner.py`, `corridorkey.py` + +```python +# model/_testing.py +"""Unstable re-exports for parity tests. Not part of public API.""" +from corridorkey_mlx.model.backbone import HieraBackbone +from corridorkey_mlx.model.hiera import HieraPatchEmbed +from corridorkey_mlx.model.decoder import DecoderHead, MLP +from corridorkey_mlx.model.refiner import CNNRefinerModule, RefinerBlock +``` + +### `inference/` — Deep Module + +```python +# inference/__init__.py +from corridorkey_mlx.inference.pipeline import ( + load_model, + compile_model, + infer, + infer_and_save, +) +from corridorkey_mlx.inference.tiling import tiled_inference +from corridorkey_mlx.inference.selective_refine import ( + selective_refine, + SelectiveRefineConfig, + SelectiveRefineResult, +) + +__all__ = [ + "load_model", + "compile_model", + "infer", + "infer_and_save", + "tiled_inference", + "SelectiveRefineConfig", + "SelectiveRefineResult", + "selective_refine", +] +``` + +Internal: `pipeline.py`, `tiling.py`, `selective_refine.py` + +```python +# inference/_testing.py +"""Unstable re-exports for unit tests. Not part of public API.""" +from corridorkey_mlx.inference.tiling import ( + _compute_tile_coords, + _make_blend_weights_2d, +) +from corridorkey_mlx.inference.selective_refine import ( + _binary_dilation, + _build_input, + _gaussian_blur, + _gaussian_kernel_1d, + _resize_to_square, +) +``` + +### `io/` — Thin Module + +```python +# io/__init__.py +from corridorkey_mlx.io.image import ( + load_image, + load_alpha_hint, + normalize_rgb, + preprocess, + postprocess_alpha, + postprocess_foreground, + save_alpha, + save_foreground, +) + +__all__ = [ + "load_image", + "load_alpha_hint", + "normalize_rgb", + "preprocess", + "postprocess_alpha", + "postprocess_foreground", + "save_alpha", + "save_foreground", +] +``` + +### `utils/` — Shared Base Layer + +```python +# utils/__init__.py +from corridorkey_mlx.utils.layout import ( + nchw_to_nhwc, + nhwc_to_nchw, + conv_weight_pt_to_mlx, + nchw_to_nhwc_np, + nhwc_to_nchw_np, +) +from corridorkey_mlx.utils.profiling import ( + TimingResult, + time_fn, + warmup_and_bench, +) + +__all__ = [ + "nchw_to_nhwc", + "nhwc_to_nchw", + "conv_weight_pt_to_mlx", + "nchw_to_nhwc_np", + "nhwc_to_nchw_np", + "TimingResult", + "time_fn", + "warmup_and_bench", +] +``` + +### `convert/` — Already Done + +Keep as-is. Already has proper `__all__`. + +### Root-level files — Stay Flat + +- `engine.py` — add `__all__ = ["CorridorKeyMLXEngine"]` +- `weights.py` — add `__all__ = ["download_weights"]` +- `weights_cli.py` — add `__all__ = ["main", "build_parser"]` +- `__init__.py` — already has `__all__` +- `__main__.py` — no changes needed + +## Technical Approach + +### Phase 1: Add Gateways (No Breaking Changes) + +1. Populate all `__init__.py` files with re-exports and `__all__` +2. Create `_testing.py` modules in `model/` and `inference/` +3. Add `__all__` to `engine.py`, `weights.py`, `weights_cli.py` +4. Run full test suite — everything must still pass (imports unchanged) + +**Success criteria:** `uv run pytest` passes, no import changes yet. + +### Phase 2: Migrate Imports + +Update all consumers to use the gateway imports: + +#### Tests + +| File | Before | After | +|------|--------|-------| +| `test_decoder_parity.py` | `from corridorkey_mlx.model.decoder import DecoderHead` | `from corridorkey_mlx.model._testing import DecoderHead` | +| `test_refiner_parity.py` | `from corridorkey_mlx.model.refiner import CNNRefinerModule` | `from corridorkey_mlx.model._testing import CNNRefinerModule` | +| `test_hiera_stage_parity.py` | `from corridorkey_mlx.model.hiera import HieraBackbone` | `from corridorkey_mlx.model._testing import HieraBackbone` | +| `test_hiera_stage_shapes.py` | `from corridorkey_mlx.model.hiera import HieraBackbone, HieraPatchEmbed` | `from corridorkey_mlx.model._testing import HieraBackbone, HieraPatchEmbed` | +| `test_greenformer_forward.py` | `from corridorkey_mlx.model.corridorkey import GreenFormer` | `from corridorkey_mlx.model import GreenFormer` | +| `test_end_to_end_smoke.py` | `from corridorkey_mlx.io.image import ...` | `from corridorkey_mlx.io import ...` | +| `test_end_to_end_parity.py` | `from corridorkey_mlx.model.corridorkey import GreenFormer` | `from corridorkey_mlx.model import GreenFormer` | +| `test_compiled_vs_eager_consistency.py` | `from corridorkey_mlx.model.corridorkey import GreenFormer` | `from corridorkey_mlx.model import GreenFormer` | +| `test_tiling_consistency.py` | `from corridorkey_mlx.inference.tiling import _compute_tile_coords, ...` | `from corridorkey_mlx.inference._testing import _compute_tile_coords, ...` | +| `test_selective_refine.py` | `from corridorkey_mlx.inference.selective_refine import _binary_dilation, ...` | `from corridorkey_mlx.inference._testing import _binary_dilation, ...` | +| `test_engine.py` | `from corridorkey_mlx.engine import _validate_image, _validate_mask` | Keep as-is (engine is flat, linter won't enforce within-file) | +| `test_weights.py` | `from corridorkey_mlx.weights import _parse_hash_line, ...` | Keep as-is (weights is flat) | +| `test_conversion.py` | `from corridorkey_mlx.convert.converter import ...` | `from corridorkey_mlx.convert import ...` | + +#### Scripts + +| File | Before | After | +|------|--------|-------| +| `infer.py` | `from corridorkey_mlx.inference.pipeline import load_model, infer_and_save` | `from corridorkey_mlx.inference import load_model, infer_and_save` | +| `bench_mlx.py` | `from corridorkey_mlx.model.corridorkey import GreenFormer` | `from corridorkey_mlx.model import GreenFormer` | +| `bench_mlx.py` | `from corridorkey_mlx.utils.profiling import warmup_and_bench` | `from corridorkey_mlx.utils import warmup_and_bench` | +| `compare_reference.py` | `from corridorkey_mlx.model.corridorkey import GreenFormer` | `from corridorkey_mlx.model import GreenFormer` | +| `convert_weights.py` | `from corridorkey_mlx.convert.converter import convert_checkpoint` | `from corridorkey_mlx.convert import convert_checkpoint` | +| `experiment_selective_refine.py` | `from corridorkey_mlx.inference.pipeline import load_model` + `selective_refine.*` | `from corridorkey_mlx.inference import load_model, selective_refine, ...` | + +#### Internal cross-package imports + +| File | Before | After | +|------|--------|-------| +| `engine.py` | `from corridorkey_mlx.model.corridorkey import GreenFormer` (TYPE_CHECKING) | `from corridorkey_mlx.model import GreenFormer` | +| `engine.py` | `from corridorkey_mlx.inference.pipeline import load_model, compile_model` | `from corridorkey_mlx.inference import load_model, compile_model` | +| `engine.py` | `from corridorkey_mlx.io.image import preprocess, ...` | `from corridorkey_mlx.io import preprocess, ...` | +| `inference/pipeline.py` | `from corridorkey_mlx.model.corridorkey import GreenFormer` | `from corridorkey_mlx.model import GreenFormer` | +| `inference/pipeline.py` | `from corridorkey_mlx.io.image import ...` | `from corridorkey_mlx.io import ...` | +| `inference/selective_refine.py` | `from corridorkey_mlx.io.image import normalize_rgb` | `from corridorkey_mlx.io import normalize_rgb` | + +**Success criteria:** `uv run pytest` passes with all new import paths. + +### Phase 3: Import Linter + +Add `import-linter` to enforce cross-package boundaries. + +```toml +# pyproject.toml additions +[tool.importlinter] +root_packages = ["corridorkey_mlx"] +include_external_packages = false + +[[tool.importlinter.contracts]] +name = "Package layering" +type = "layers" +layers = [ + "corridorkey_mlx.engine", + "corridorkey_mlx.inference", + "corridorkey_mlx.model | corridorkey_mlx.io | corridorkey_mlx.convert", + "corridorkey_mlx.utils", +] +# engine → inference → model/io/convert → utils + +[[tool.importlinter.contracts]] +name = "Convert independence" +type = "independence" +modules = [ + "corridorkey_mlx.convert", + "corridorkey_mlx.model", + "corridorkey_mlx.io", +] +# convert, model, io should not import from each other + +[[tool.importlinter.contracts]] +name = "Weights independence" +type = "independence" +modules = [ + "corridorkey_mlx.weights", + "corridorkey_mlx.model", + "corridorkey_mlx.inference", + "corridorkey_mlx.io", +] +# weights module is standalone +``` + +Add to dev dependencies and CLAUDE.md commands: + +```bash +uv run lint-imports # import boundary check +``` + +**Success criteria:** `uv run lint-imports` passes. All cross-package imports go through gateways. + +### Phase 4: Documentation & CLAUDE.md Updates + +1. Update CLAUDE.md repo layout section to document public interfaces +2. Add a "Module Boundaries" section explaining the gateway pattern +3. Add `_testing.py` convention to CLAUDE.md +4. Update commands section with `lint-imports` + +## Acceptance Criteria + +- [ ] Every subpackage `__init__.py` has `__all__` defining its public API +- [ ] `model/_testing.py` and `inference/_testing.py` exist with internal re-exports +- [ ] All tests/scripts use gateway imports (no direct `.corridorkey`, `.decoder`, `.pipeline` imports cross-package) +- [ ] `import-linter` passes with layering + independence contracts +- [ ] Full test suite passes (`uv run pytest`) +- [ ] Lint passes (`uv run ruff check .`) +- [ ] CLAUDE.md documents module boundaries and gateway convention + +## Dependencies & Risks + +**Risk: circular imports from re-exports.** Mitigated by the existing acyclic dependency graph — no cycles detected. + +**Risk: import-linter false positives.** `TYPE_CHECKING` imports should be excluded. `import-linter` supports `--exclude-type-checking-imports`. Verify this works. + +**Risk: breaking scripts outside the repo.** Only the root `CorridorKeyMLXEngine` export matters for external consumers. Internal restructuring doesn't affect the public API. + +**Dependency:** `import-linter` package added to dev deps. + +## Unresolved Questions + +- `utils.layout` funcs used inside `convert/converter.py` — does `convert` importing from `utils` violate any intended boundary? (current plan: `utils` is the base layer, everyone can import it) +- Should `_testing.py` modules have their own `__all__`? (leaning yes for explicitness) +- `weights.py` has module-level constants (`DEFAULT_ASSET_NAME`, etc.) used by `test_weights.py` — add to `__all__` or keep test accessing flat file directly?