diff --git a/src/corridorkey_mlx/inference/pipeline.py b/src/corridorkey_mlx/inference/pipeline.py index 2e8c979..be94813 100644 --- a/src/corridorkey_mlx/inference/pipeline.py +++ b/src/corridorkey_mlx/inference/pipeline.py @@ -68,6 +68,7 @@ def compile_model(model: GreenFormer, shapeless: bool = False) -> GreenFormer: same resolution. Shapeless compile is experimental — the backbone uses shape-dependent reshapes that may trigger recompilation. """ + model._compiled = True model.__call__ = mx.compile(model.__call__, shapeless=shapeless) # type: ignore[method-assign] return model diff --git a/src/corridorkey_mlx/model/corridorkey.py b/src/corridorkey_mlx/model/corridorkey.py index a41a45f..979dda2 100644 --- a/src/corridorkey_mlx/model/corridorkey.py +++ b/src/corridorkey_mlx/model/corridorkey.py @@ -6,6 +6,7 @@ All internal operations use NHWC layout. from __future__ import annotations +import gc from typing import TYPE_CHECKING import mlx.core as mx @@ -42,6 +43,7 @@ class GreenFormer(nn.Module): self._compute_dtype = dtype self._fused_decode = fused_decode self._slim = slim + self._compiled = False self.backbone = HieraBackbone(img_size=img_size) self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1) self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3) @@ -69,6 +71,13 @@ class GreenFormer(nn.Module): # Backbone always runs in fp32 features = self.backbone(x) + # Materialize backbone output so MLX can free intermediate graph nodes. + # NOTE: mx.eval is MLX array materialization, not Python eval() + if not self._compiled: + mx.eval(features) # noqa: S307 + gc.collect() + mx.clear_cache() + # Cast features to compute dtype for decoders (bf16 saves memory) if self._compute_dtype != mx.float32: features = [f.astype(self._compute_dtype) for f in features] @@ -93,6 +102,13 @@ class GreenFormer(nn.Module): alpha_coarse = mx.sigmoid(alpha_logits_up) fg_coarse = mx.sigmoid(fg_logits_up) + # Materialize decoder output so MLX can free decoder graph nodes. + # NOTE: mx.eval is MLX array materialization, not Python eval() + if not self._compiled: + mx.eval(alpha_coarse, fg_coarse, alpha_logits_up, fg_logits_up) # noqa: S307 + gc.collect() + mx.clear_cache() + # Refiner receives fp32 coarse predictions rgb = x[:, :, :, :3] # (B, H, W, 3) coarse_pred = mx.concatenate([alpha_coarse, fg_coarse], axis=-1) # (B, H, W, 4) diff --git a/tests/test_compilation.py b/tests/test_compilation.py index 221b8cd..e19631f 100644 --- a/tests/test_compilation.py +++ b/tests/test_compilation.py @@ -5,6 +5,7 @@ from __future__ import annotations import mlx.core as mx import pytest +from corridorkey_mlx.inference.pipeline import compile_model from corridorkey_mlx.model.corridorkey import GreenFormer IMG_SIZE = 256 @@ -32,11 +33,11 @@ def dummy_input() -> mx.array: def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> None: """Fixed-shape compiled output matches eager output within tolerance.""" eager_out = model(dummy_input) - mx.eval(eager_out) # noqa: S307 + mx.eval(eager_out) # noqa: S307 # NOTE: MLX materialization, not Python eval - compiled_fn = mx.compile(model.__call__) - compiled_out = compiled_fn(dummy_input) - mx.eval(compiled_out) # noqa: S307 + compiled_model = compile_model(model) + compiled_out = compiled_model(dummy_input) + mx.eval(compiled_out) # noqa: S307 # NOTE: MLX materialization for key in OUTPUT_KEYS: diff = float(mx.max(mx.abs(eager_out[key] - compiled_out[key]))) @@ -45,12 +46,12 @@ def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> No def test_compiled_deterministic(model: GreenFormer, dummy_input: mx.array) -> None: """Compiled model produces identical results across consecutive calls.""" - compiled_fn = mx.compile(model.__call__) + compiled_model = compile_model(model) - out1 = compiled_fn(dummy_input) - mx.eval(out1) # noqa: S307 - out2 = compiled_fn(dummy_input) - mx.eval(out2) # noqa: S307 + out1 = compiled_model(dummy_input) + mx.eval(out1) # noqa: S307 # NOTE: MLX materialization + out2 = compiled_model(dummy_input) + mx.eval(out2) # noqa: S307 # NOTE: MLX materialization for key in OUTPUT_KEYS: diff = float(mx.max(mx.abs(out1[key] - out2[key])))