Skip to content

fix: add Apple Silicon (MPS) support (was silently falling back to CPU) - #2

Open
gaaaork wants to merge 1 commit into
RealRebelAI:mainfrom
GaaaorkOrga:fix/apple-silicon-mps-support
Open

fix: add Apple Silicon (MPS) support (was silently falling back to CPU)#2
gaaaork wants to merge 1 commit into
RealRebelAI:mainfrom
GaaaorkOrga:fix/apple-silicon-mps-support

Conversation

@gaaaork

@gaaaork gaaaork commented Jul 18, 2026

Copy link
Copy Markdown

Problem

Every device check in nodes.py and encoder_loader.py was torch.cuda.is_available(),
with no MPS branch. On a Mac, that resolves to False, so the pack silently falls back to
CPU for the whole 5B transformer — no error, no warning, just very slow generation.

Confirmed on real Apple Silicon hardware (M-series, 48GB unified memory, PyTorch 2.13):
~1432s/step (24 min) on CPU vs ~4-6s/step once routed to MPS — roughly a 250-400x
difference, same weights, same output.

Fix

Added device_compat.py (new, small, self-contained):

  • best_device()cuda > mps > cpu, replaces the "cuda" if torch.cuda.is_available() else "cpu" pattern used at the one spot that actually decides where the transformer runs.
  • empty_cache(device) — MPS/CUDA-aware cache clear, replaces the scattered
    if torch.cuda.is_available(): torch.cuda.empty_cache() calls.

Left untouched on purpose: the CUDA-specific VRAM-budget block-residency planner
(_plan_residency, using torch.cuda.mem_get_info()). MPS's unified memory model doesn't
have a direct equivalent, and the existing "not CUDA -> keep everything resident" fallback
already does the right thing once device correctly resolves to "mps" instead of "cpu"
no new logic needed there.

Testing

Real end-to-end runs on Apple Silicon, standard Rebels SeFi Loader -> Rebels SeFi Sampler ->
SaveImage graph, weight_dtype=bf16 (fp8_e4m3fn is CUDA-oriented, not exercised here):

  • 4-step smoke test: completed cleanly, confirmed Device: mps in the ComfyUI log, no
    unsupported-op errors.
  • Real 50-step Base generation (SeFi-5B-Base, 1024x1024, real prompt): 349s total, correct
    image, no errors.

No behavior change on CUDA or CPU-only machines — best_device() returns the exact same
value torch.cuda.is_available() would have on those platforms.

Scope

Deliberately minimal — 3 files, one new. Doesn't touch the folder-structure issue (#1),
unrelated to this fix.

torch.cuda.is_available() was the only device check throughout nodes.py and
encoder_loader.py, so on any Mac this pack ran the 5B transformer on CPU
instead of the GPU, with no warning. Confirmed on real hardware: ~1432s/step
(24 min) on Apple Silicon CPU vs ~4-6s/step once routed to MPS — roughly a
250-400x difference for the exact same output.

New device_compat.py centralizes the platform check (best_device(), used for
picking cuda > mps > cpu; empty_cache(device), the MPS/CUDA-aware equivalent
of the scattered torch.cuda.empty_cache() calls). Block-residency streaming
(CUDA-specific VRAM budgeting via torch.cuda.mem_get_info()) is left CUDA-only
on purpose — MPS's unified memory model doesn't need it, the existing
non-CUDA fallback (whole transformer resident) already does the right thing
once device correctly resolves to "mps" instead of "cpu".

Tested end-to-end on Apple Silicon (M-series, 48GB unified memory, PyTorch
2.13, macOS): 4-step smoke test then a real 50-step Base generation, both via
the standard Rebels SeFi Loader -> Rebels SeFi Sampler -> SaveImage graph,
weight_dtype=bf16 (fp8_e4m3fn is CUDA-oriented and untested here), 1024x1024,
real prompt. 50 steps completed in 349s with no errors and a correct image.
No behavior change on CUDA or CPU-only machines.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant