Skip to content

Device and Backend Selection

Device Selection

By default, CorridorKey auto-detects the best available compute device in this priority order:

CUDA → MPS → CPU

Override via CLI Flag

uv run python clip_manager.py --action wizard --win_path "V:\..." --device mps
uv run python clip_manager.py --action run_inference --device cpu

Override via Environment Variable

export CORRIDORKEY_DEVICE=cpu
uv run python clip_manager.py --action wizard --win_path "V:\..."

Resolution order

--device flag > CORRIDORKEY_DEVICE env var > auto-detect.

Apple Silicon (MPS / MLX)

CorridorKey runs on Apple Silicon Macs using unified memory. Two backend options are available:

  • MPS — PyTorch's Metal Performance Shaders backend. Works out of the box but some operators may require the CPU fallback flag:

    export PYTORCH_ENABLE_MPS_FALLBACK=1

  • MLX — Native Apple Silicon acceleration via the MLX framework. Avoids PyTorch's MPS layer entirely and typically runs faster. Requires installing the MLX extras (uv sync --extra mlx) and obtaining .safetensors weights.

Because Apple Silicon shares memory between the CPU and GPU, the full system RAM is available to the model — no separate VRAM budget applies.

Backend Selection

CorridorKey supports two inference backends:

Backend Platforms Notes
Torch (default on Linux / Windows) CUDA, MPS, CPU Standard PyTorch inference.
MLX (Apple Silicon) Metal Native Apple Silicon acceleration — avoids PyTorch's MPS layer entirely and typically runs faster.

Override via CLI Flag

uv run python corridorkey_cli.py --action wizard --win_path "/path/to/clips" --backend mlx
uv run python corridorkey_cli.py --action run_inference --backend torch

Override via Environment Variable

export CORRIDORKEY_BACKEND=mlx
uv run python corridorkey_cli.py --action run_inference

Resolution order

--backend flag > CORRIDORKEY_BACKEND env var > auto-detect.

Auto mode prefers MLX on Apple Silicon when the package is installed.

MLX Setup (Apple Silicon)

Follow these steps to use the native MLX backend on an M1+ Mac.

1. Install the MLX Extras

uv sync --extra mlx

2. Obtain MLX Weights (.safetensors)

# Download weights from GitHub Releases into a local cache directory
uv run python -m corridorkey_mlx weights download

# Print the cached path, then copy to the checkpoints folder
WEIGHTS=$(uv run python -m corridorkey_mlx weights download --print-path)
cp "$WEIGHTS" CorridorKeyModule/checkpoints/corridorkey_mlx.safetensors

# Clone the MLX repo (contains the conversion script)
git clone https://github.com/nikopueringer/corridorkey-mlx.git
cd corridorkey-mlx
uv sync

# Convert (point --checkpoint at your CorridorKey.pth)
uv run python scripts/convert_weights.py \
    --checkpoint ../CorridorKeyModule/checkpoints/CorridorKey_v1.0.pth \
    --output ../CorridorKeyModule/checkpoints/corridorkey_mlx.safetensors
cd ..

Either way, the final file must be at:

CorridorKeyModule/checkpoints/corridorkey_mlx.safetensors

3. Run

CORRIDORKEY_BACKEND=mlx uv run python clip_manager.py --action run_inference

MLX uses img_size=2048 by default (same as Torch).

Troubleshooting

MPS (PyTorch Metal)

Confirm MPS is active — run with verbose logging to see which device was selected:

uv run python clip_manager.py --action list 2>&1 | grep -i "device\|backend\|mps"

MPS operator errors (NotImplementedError: ... not implemented for 'MPS'): Some PyTorch operations are not yet supported on MPS. Enable CPU fallback:

export PYTORCH_ENABLE_MPS_FALLBACK=1
uv run python corridorkey_cli.py --action wizard --win_path "/path/to/clips"

Make the fallback permanent

Add export PYTORCH_ENABLE_MPS_FALLBACK=1 to your shell profile (~/.zshrc) so it is always active. Without it, MPS may silently fall back to CPU, making runs much slower.

Use native MLX instead of PyTorch MPS — MLX avoids PyTorch's MPS layer entirely and typically runs faster on Apple Silicon. See the MLX Setup section above.

MLX

Symptom Fix
No .safetensors checkpoint found Place MLX weights in CorridorKeyModule/checkpoints/.
corridorkey_mlx not installed Run uv sync --extra mlx.
MLX requires Apple Silicon MLX only works on M1+ Macs.
Auto picked Torch unexpectedly Set CORRIDORKEY_BACKEND=mlx explicitly.