From 2369eb00e7dea37654ee30e0ff273a5a573def96 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 22 May 2026 19:53:04 -0700 Subject: [PATCH] Route aimdo init through get_all_torch_devices() instead of raw torch.cuda The aimdo init call on worksplit-multigpu was using comfy_aimdo.control.init_devices(range(torch.cuda.device_count())) which required adding `import torch` at the top of main.py (violating the "torch should never be imported before this point" expectation) and an inner is_nvidia() guard added in PR #14068 to defend the raw cuda call on non-NVIDIA systems where --enable-dynamic-vram is explicitly passed. Replace the call with comfy_aimdo.control.init_devices( d.index for d in comfy.model_management.get_all_torch_devices() if d.type == "cuda" and d.index is not None ) comfy_aimdo.control.init_devices accepts any iterable of int-coercible device indices and returns False on an empty iterable, so on non-cuda systems the elif naturally falls through to the existing "No working comfy-aimdo install detected" fallback - no extra vendor gate needed. HIP devices appear as type "cuda" in torch, so ROCm setups (which comfy-aimdo supports via aimdo_rocm.so) are handled correctly too. This lets us drop both the `import torch` at the top of main.py and the inner is_nvidia() guard, leaving a single logical-line divergence from master (init_device(single index) -> init_devices(generator of cuda indices)) for multi-GPU aimdo support. Amp-Thread-ID: https://ampcode.com/threads/T-019e52b4-31ee-72cd-996b-64ecd9420e13 Co-authored-by: Amp --- main.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/main.py b/main.py index 9b22d1304..b1a19c7bf 100644 --- a/main.py +++ b/main.py @@ -200,7 +200,7 @@ import gc if 'torch' in sys.modules: logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.") -import torch + import comfy.utils import execution @@ -216,16 +216,12 @@ 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 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)): + if (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())): + elif comfy_aimdo.control.init_devices( + d.index for d in comfy.model_management.get_all_torch_devices() + if d.type == "cuda" and d.index is not None + ): if args.verbose == 'DEBUG': comfy_aimdo.control.set_log_debug() elif args.verbose == 'CRITICAL':