diff --git a/comfy/model_management.py b/comfy/model_management.py index 051062a90..c146eee11 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -214,7 +214,10 @@ def get_all_torch_devices(exclude_current=False): global cpu_state devices = [] if cpu_state == CPUState.GPU: - if is_nvidia(): + # NVIDIA + AMD/ROCm both expose their GPUs through torch.cuda.*; + # without the AMD arm, single-GPU ROCm users get an empty list + # which silently turns unload_all_models() into a no-op. + if is_nvidia() or is_amd(): for i in range(torch.cuda.device_count()): devices.append(torch.device("cuda", i)) elif is_intel_xpu(): @@ -223,6 +226,14 @@ def get_all_torch_devices(exclude_current=False): elif is_ascend_npu(): for i in range(torch.npu.device_count()): devices.append(torch.device("npu", i)) + elif is_mlu(): + for i in range(torch.mlu.device_count()): + devices.append(torch.device("mlu", i)) + else: + # Fallback for unhandled GPU backends (e.g. DirectML): at least + # report the current device so callers like unload_all_models() + # do not silently no-op. + devices.append(get_torch_device()) else: devices.append(get_torch_device()) if exclude_current: diff --git a/main.py b/main.py index 9933d11ee..9b22d1304 100644 --- a/main.py +++ b/main.py @@ -216,7 +216,14 @@ import comfy.memory_management import comfy.model_patcher if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia() and not comfy.model_management.is_wsl()): - if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)): + if not comfy.model_management.is_nvidia(): + # The implicit auto-enable path is already gated by is_nvidia(); + # this guard handles users who pass --enable-dynamic-vram explicitly + # on a non-NVIDIA system, where torch.cuda.device_count() below would + # either return 0 (silently disabling) or crash on backends that + # raise without CUDA. Be explicit and disable cleanly. + logging.warning("DynamicVRAM was requested but no NVIDIA GPU was detected. Falling back to legacy ModelPatcher.") + elif (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)): logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows") elif comfy_aimdo.control.init_devices(range(torch.cuda.device_count())): if args.verbose == 'DEBUG':