From 4d7a29c7c7481feae00a817b7407232ab634e328 Mon Sep 17 00:00:00 2001 From: cmoyates Date: Mon, 9 Mar 2026 17:00:16 -0230 Subject: [PATCH] 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 --- src/corridorkey_mlx/engine.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/corridorkey_mlx/engine.py b/src/corridorkey_mlx/engine.py index 4faa207..d2127f4 100644 --- a/src/corridorkey_mlx/engine.py +++ b/src/corridorkey_mlx/engine.py @@ -102,7 +102,7 @@ class CorridorKeyMLXEngine: # Tiled: model runs at tile_size, input stays full-res self._img_size = self._tile_size self._model: GreenFormer = load_model( - checkpoint, img_size=self._tile_size, compile=False + checkpoint, img_size=self._tile_size, compile=False, slim=True ) logger.info( "Tiled inference: tile_size=%d, overlap=%d", self._tile_size, self._overlap @@ -110,7 +110,7 @@ class CorridorKeyMLXEngine: else: # Full-frame: resize input to img_size self._img_size = img_size - self._model = load_model(checkpoint, img_size=img_size, compile=compile) + self._model = load_model(checkpoint, img_size=img_size, compile=compile, slim=True) def process_frame( self, @@ -213,6 +213,11 @@ class CorridorKeyMLXEngine: alpha_u8 = postprocess_alpha(alpha_out) fg_u8 = postprocess_foreground(fg_out) + # Release all MLX array references, then free Metal buffer cache + del alpha_out, fg_out + gc.collect() + mx.clear_cache() + # -- resize back to original -- if alpha_u8.shape[0] != original_h or alpha_u8.shape[1] != original_w: target = (original_w, original_h)