monster/corridorkey-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add mx.clear_cache() to engine cleanup after postprocessing

Browse Source 
Del MLX array refs + gc + cache clear after numpy extraction.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
cmoyates 2026-03-09 17:00:16 -02:30
parent ee23f27735
commit 4d7a29c7c7
No known key found for this signature in database
GPG Key ID: F65E08480FDC22D2
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/corridorkey_mlx/engine.py
View File

@ -102,7 +102,7 @@ class CorridorKeyMLXEngine:
# Tiled: model runs at tile_size, input stays full-res # Tiled: model runs at tile_size, input stays full-res
self._img_size = self._tile_size self._img_size = self._tile_size
self._model: GreenFormer = load_model( 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( logger.info(
"Tiled inference: tile_size=%d, overlap=%d", self._tile_size, self._overlap "Tiled inference: tile_size=%d, overlap=%d", self._tile_size, self._overlap
@ -110,7 +110,7 @@ class CorridorKeyMLXEngine:
else: else:
# Full-frame: resize input to img_size # Full-frame: resize input to img_size
self._img_size = 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( def process_frame(
self, self,
@ -213,6 +213,11 @@ class CorridorKeyMLXEngine:
alpha_u8 = postprocess_alpha(alpha_out) alpha_u8 = postprocess_alpha(alpha_out)
fg_u8 = postprocess_foreground(fg_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 -- # -- resize back to original --
if alpha_u8.shape[0] != original_h or alpha_u8.shape[1] != original_w: if alpha_u8.shape[0] != original_h or alpha_u8.shape[1] != original_w:
target = (original_w, original_h) target = (original_w, original_h)