mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-07-23 09:10:51 -05:00
5ffea26de7
Two fixes for single-GPU users on non-NVIDIA backends; multi-GPU
non-CUDA support is intentionally out of scope here (tracked separately).
1. get_all_torch_devices: add AMD/ROCm, MLU, and a generic fallback arm.
Previously the function only enumerated NVIDIA, Intel XPU, and Ascend
NPU when cpu_state==GPU; on AMD/ROCm (which exposes its GPU through
torch.cuda.*) and DirectML it fell through to an empty list. The
biggest user-visible regression: unload_all_models() iterates this
list, so it became a silent no-op on AMD/ROCm. /free, manager
unloads, and shutdown stopped releasing VRAM.
- is_amd() now shares the torch.cuda.* arm with is_nvidia(), since
ROCm reuses the CUDA API surface.
- is_mlu() gets its own arm using torch.mlu.device_count().
- A final fallback appends get_torch_device() for any GPU backend
the explicit arms miss (notably DirectML), so callers see at
least the current device and unload_all_models works.
MPS users are unaffected: cpu_state==MPS already routes to the
else branch which appends get_torch_device() returning mps.
2. main.py DynamicVRAM init: guard the comfy_aimdo branch with an
explicit is_nvidia() check.
The outer condition allows entering the DynamicVRAM init block when
the user passes --enable-dynamic-vram explicitly, bypassing the
implicit is_nvidia() gate. On non-NVIDIA backends this then runs
comfy_aimdo.control.init_devices(range(torch.cuda.device_count())),
which is comfy-aimdo-only territory and may crash at startup. Add a
leading is_nvidia() check that logs a clean warning and falls back
to the legacy ModelPatcher path.