wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-24 07:57:29 +08:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

Route aimdo init through get_all_torch_devices() instead of raw torch.cuda

Browse Source 
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 <amp@ampcode.com>
This commit is contained in:
Jedrzej Kosinski 2026-05-22 19:53:04 -07:00
parent 403ff49647
commit 2369eb00e7
1 changed files with 6 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
main.py
View File

@ -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':