Add shared cuda→mps→cpu device-selection helper (#166) - #254
Open
MarcBarnaba wants to merge 2 commits into
Open
Conversation
BaseGenerator subclasses currently hardcode torch.cuda.is_available(), so the four bundled models only run on Windows/Linux with an NVIDIA GPU. Add select_device()/select_dtype() to base.py so extensions can detect Apple Silicon (MPS) and fall back to CPU in one shared place instead of duplicating the check per extension. unload() now also releases MPS memory via torch.mps.empty_cache(). Windows/CUDA behavior is unchanged. Ref: lightningpixel#166
torch.backends.mps hard-crashes with NotImplementedError on ops that have no Metal kernel yet (e.g. 3D pooling), instead of transparently falling back to CPU like on other backends. The opt-in env var that enables the fallback must be set before the process's first `import torch` — setting it inside select_device() is too late for extensions that import torch before calling it (confirmed by testing TripoSG end to end). Set it once, for every extension subprocess, in ExtensionProcess._build_env() on macOS. Ref: lightningpixel#166
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of #166 (FEATURE 5 — macOS support for all models).
BaseGeneratorsubclasses currently hardcodetorch.cuda.is_available(),so model extensions only run on Windows/Linux with an NVIDIA GPU. This adds
two small helpers to
api/services/generators/base.py:select_device()— pickscuda→mps(Apple Silicon) →cpu.select_dtype(device)—float16on CUDA,float32on MPS/CPU (MPS hasincomplete fp16 kernel coverage for attention/layer-norm/interpolation).
BaseGenerator.unload()now also releases MPS memory viatorch.mps.empty_cache()(it already did this for CUDA).Extensions import these from
services.generators.basethe same way theyalready import
smooth_progress/GenerationCancelled, so device logiclives in one place instead of being duplicated per extension, per the
ticket's implementation notes.
Companion extension PR: lightningpixel/modly-triposg-extension uses
these helpers — see that PR for the per-model changes.
Status of the other three models in the ticket, for visibility:
main(PR feat(update): add update available modal with version check #11 on modly-hunyuan3d-mini-extension, and an equivalent commiton modly-hunyuan3d-mini-turbo-extension) — no changes needed there.
spconvfor its sparse convolutions, which hasno CPU or Metal/MPS backend at all (only an unsupported Linux-only debug
CPU build) — full MPS support isn't achievable without replacing that
dependency, so it's out of scope for this pass.
No frontend/store/API surface changes. Windows/CUDA behavior is
unchanged (
torch.cuda.is_available()is still checked first everywhere).