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

feat: add tiled inference support to engine

Browse Source 
tile_size/overlap params on CorridorKeyMLXEngine — model loads at
tile_size, input stays full-res, forward uses tiled_inference w/ GC.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
cmoyates 2026-03-08 22:41:45 -02:30
parent f30ab9255b
commit ee23f27735
No known key found for this signature in database
GPG Key ID: F65E08480FDC22D2
1 changed files with 78 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

114
src/corridorkey_mlx/engine.py
View File

@ -17,6 +17,10 @@ import numpy as np
from PIL import Image from PIL import Image
from corridorkey_mlx.inference.pipeline import load_model from corridorkey_mlx.inference.pipeline import load_model
from corridorkey_mlx.inference.tiling import (
DEFAULT_OVERLAP,
tiled_inference,
)
from corridorkey_mlx.io.image import ( from corridorkey_mlx.io.image import (
postprocess_alpha, postprocess_alpha,
postprocess_foreground, postprocess_foreground,
@ -40,17 +44,27 @@ class CorridorKeyMLXEngine:
device: Ignored on MLX (Apple Silicon uses unified memory). device: Ignored on MLX (Apple Silicon uses unified memory).
Accepted for API compatibility with the Torch engine. Accepted for API compatibility with the Torch engine.
img_size: Internal model resolution (square). The model was trained img_size: Internal model resolution (square). The model was trained
at 2048. Use 512 for fast dev iteration. at 2048. Use 512 for fast dev iteration. When tiling is enabled,
this is ignored and the model runs at ``tile_size``.
use_refiner: If True, return refined alpha/fg. If False, return use_refiner: If True, return refined alpha/fg. If False, return
coarse predictions (skips refiner in postprocessing, not forward pass). coarse predictions (skips refiner in postprocessing, not forward pass).
compile: If True, wrap model forward with ``mx.compile`` for faster compile: If True, wrap model forward with ``mx.compile`` for faster
repeated inference at the same resolution. repeated inference at the same resolution.
tile_size: If set, enable tiled inference — split full-res input into
overlapping tiles of this size. Model loads at tile_size instead of
img_size. Set to 0 or None to disable.
overlap: Overlap in pixels between adjacent tiles (default 64).
Example:: Example::
from corridorkey_mlx import CorridorKeyMLXEngine from corridorkey_mlx import CorridorKeyMLXEngine
# Full-frame (default)
engine = CorridorKeyMLXEngine("/path/to/corridorkey_mlx.safetensors") engine = CorridorKeyMLXEngine("/path/to/corridorkey_mlx.safetensors")
# Tiled — 12x less memory at 2048x2048
engine = CorridorKeyMLXEngine("/path/to/ckpt.safetensors", tile_size=512, overlap=64)
result = engine.process_frame(rgb_uint8, mask_uint8) result = engine.process_frame(rgb_uint8, mask_uint8)
# result["alpha"] — (H, W) uint8 # result["alpha"] — (H, W) uint8
# result["fg"] — (H, W, 3) uint8 # result["fg"] — (H, W, 3) uint8
@ -68,6 +82,8 @@ class CorridorKeyMLXEngine:
img_size: int = PRODUCTION_IMG_SIZE, img_size: int = PRODUCTION_IMG_SIZE,
use_refiner: bool = True, use_refiner: bool = True,
compile: bool = True, compile: bool = True,
tile_size: int | None = None,
overlap: int = DEFAULT_OVERLAP,
) -> None: ) -> None:
checkpoint = Path(checkpoint_path) checkpoint = Path(checkpoint_path)
if not checkpoint.exists(): if not checkpoint.exists():
@ -77,9 +93,24 @@ class CorridorKeyMLXEngine:
if device is not None: if device is not None:
logger.info("device=%r ignored on MLX (unified memory)", device) logger.info("device=%r ignored on MLX (unified memory)", device)
self._img_size = img_size
self._use_refiner = use_refiner self._use_refiner = use_refiner
self._model: GreenFormer = load_model(checkpoint, img_size=img_size, compile=compile) self._tiled = bool(tile_size)
self._tile_size = tile_size or img_size
self._overlap = overlap
if self._tiled:
# 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
)
logger.info(
"Tiled inference: tile_size=%d, overlap=%d", self._tile_size, self._overlap
)
else:
# Full-frame: resize input to img_size
self._img_size = img_size
self._model = load_model(checkpoint, img_size=img_size, compile=compile)
def process_frame( def process_frame(
self, self,
@ -125,48 +156,59 @@ class CorridorKeyMLXEngine:
if mask_f32.ndim == 2: if mask_f32.ndim == 2:
mask_f32 = mask_f32[:, :, np.newaxis] mask_f32 = mask_f32[:, :, np.newaxis]
# -- resize to model resolution -- if self._tiled:
if rgb_f32.shape[0] != self._img_size or rgb_f32.shape[1] != self._img_size: # Tiled: keep full-res, preprocess, run tiled_inference
rgb_pil = Image.fromarray(image).resize( x = preprocess(rgb_f32, mask_f32)
(self._img_size, self._img_size), Image.Resampling.BICUBIC outputs = tiled_inference(
self._model, x, tile_size=self._tile_size, overlap=self._overlap
) )
rgb_f32 = np.asarray(rgb_pil, dtype=np.float32) / 255.0 # NOTE: mx.eval is MLX array materialization, not Python eval()
mx.eval(outputs) # noqa: S307
else:
# Full-frame: resize to model resolution
if rgb_f32.shape[0] != self._img_size or rgb_f32.shape[1] != self._img_size:
rgb_pil = Image.fromarray(image).resize(
(self._img_size, self._img_size), Image.Resampling.BICUBIC
)
rgb_f32 = np.asarray(rgb_pil, dtype=np.float32) / 255.0
mask_u8 = mask_linear if mask_linear.ndim == 2 else mask_linear[:, :, 0] mask_u8 = mask_linear if mask_linear.ndim == 2 else mask_linear[:, :, 0]
mask_pil = Image.fromarray(mask_u8, mode="L").resize( mask_pil = Image.fromarray(mask_u8, mode="L").resize(
(self._img_size, self._img_size), Image.Resampling.BICUBIC (self._img_size, self._img_size), Image.Resampling.BICUBIC
) )
mask_f32 = np.asarray(mask_pil, dtype=np.float32)[:, :, np.newaxis] / 255.0 mask_f32 = np.asarray(mask_pil, dtype=np.float32)[:, :, np.newaxis] / 255.0
# -- preprocess (ImageNet norm + concat) -> (1, H, W, 4) NHWC -- x = preprocess(rgb_f32, mask_f32)
x = preprocess(rgb_f32, mask_f32) outputs = self._model(x)
# NOTE: mx.eval is MLX array materialization, not Python eval()
# -- forward -- mx.eval(outputs) # noqa: S307
outputs = self._model(x)
mx.eval(outputs) # noqa: S307 — mx.eval materializes lazy MLX arrays, not Python eval
# -- select coarse vs refined -- # -- select coarse vs refined --
alpha_coarse = outputs["alpha_coarse"] if self._tiled:
fg_coarse = outputs["fg_coarse"] # tiled_inference only returns final (refined) outputs
alpha_refined = outputs["alpha_final"] alpha_out = outputs["alpha_final"]
fg_refined = outputs["fg_final"] fg_out = outputs["fg_final"]
else:
alpha_coarse = outputs["alpha_coarse"]
fg_coarse = outputs["fg_coarse"]
alpha_refined = outputs["alpha_final"]
fg_refined = outputs["fg_final"]
# Release unused intermediate tensors (logits, delta_logits, etc.) if not self._use_refiner or refiner_scale == 0.0:
alpha_out = alpha_coarse
fg_out = fg_coarse
elif refiner_scale == 1.0:
alpha_out = alpha_refined
fg_out = fg_refined
else:
s = refiner_scale
alpha_out = alpha_coarse * (1.0 - s) + alpha_refined * s
fg_out = fg_coarse * (1.0 - s) + fg_refined * s
# Release unused intermediate tensors
del outputs del outputs
gc.collect() gc.collect()
if not self._use_refiner or refiner_scale == 0.0:
alpha_out = alpha_coarse
fg_out = fg_coarse
elif refiner_scale == 1.0:
alpha_out = alpha_refined
fg_out = fg_refined
else:
# output-space lerp
s = refiner_scale
alpha_out = alpha_coarse * (1.0 - s) + alpha_refined * s
fg_out = fg_coarse * (1.0 - s) + fg_refined * s
# -- postprocess to uint8 -- # -- postprocess to uint8 --
alpha_u8 = postprocess_alpha(alpha_out) alpha_u8 = postprocess_alpha(alpha_out)
fg_u8 = postprocess_foreground(fg_out) fg_u8 = postprocess_foreground(fg_out)