Method process_frame

CorridorKeyModule/backend.py:378–432

Delegate to MLX engine, then normalize output to Torch contract. ``screen_channel`` is accepted for API parity with the Torch engine but only ``1`` (green) is supported here — the MLX backend has no blue checkpoint yet, so the despill in ``_wrap_mlx_output`` is hard-wired to

(
        self,
        image,
        mask_linear,
        refiner_scale=1.0,
        input_is_linear=False,
        fg_is_straight=True,
        despill_strength=1.0,
        auto_despeckle=True,
        despeckle_size=400,
        screen_channel: int = 1,
        **_kwargs,
    )

        logger.info("MLX adapter active: despill and despeckle are handled by the adapter layer, not native MLX")

    def process_frame(
        self,
        image,
        mask_linear,
        refiner_scale=1.0,
        input_is_linear=False,
        fg_is_straight=True,
        despill_strength=1.0,
        auto_despeckle=True,
        despeckle_size=400,
        screen_channel: int = 1,
        **_kwargs,
    ):
        """Delegate to MLX engine, then normalize output to Torch contract.

        ``screen_channel`` is accepted for API parity with the Torch engine but
        only ``1`` (green) is supported here — the MLX backend has no blue
        checkpoint yet, so the despill in ``_wrap_mlx_output`` is hard-wired to
        the green channel. Calling with ``screen_channel != 1`` is a programmer
        error (the public ``create_engine`` rejects MLX + blue earlier); we
        raise instead of silently returning a green-keyed result.
        """
        if screen_channel != 1:
            raise NotImplementedError(
                f"_MLXEngineAdapter does not support screen_channel={screen_channel}. "
                "MLX has no blue-screen checkpoint yet; use the Torch backend with "
                "--screen-color blue, or wait for the MLX blue release."
            )
        # MLX engine expects uint8 input — convert if float
        if image.dtype != np.uint8:
            image_u8 = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
        else:
            image_u8 = image

        if mask_linear.dtype != np.uint8:
            mask_u8 = (np.clip(mask_linear, 0.0, 1.0) * 255).astype(np.uint8)
        else:
            mask_u8 = mask_linear

        # Squeeze mask to 2D for MLX (it validates [H,W] or [H,W,1])
        if mask_u8.ndim == 3:
            mask_u8 = mask_u8[:, :, 0]

        raw = self._engine.process_frame(
            image_u8,
            mask_u8,
            refiner_scale=refiner_scale,
            input_is_linear=input_is_linear,
            fg_is_straight=fg_is_straight,
            despill_strength=0.0,  # disable MLX stubs — adapter applies these
            auto_despeckle=False,
            despeckle_size=despeckle_size,
        )

        return _wrap_mlx_output(raw, despill_strength, auto_despeckle, despeckle_size)


DEFAULT_MLX_TILE_SIZE = 512