wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-25 16:37:23 +08:00
Code Issues Actions 15 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: if 'torch' in sys.modules:
logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.") logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
import torch
import comfy.utils import comfy.utils
import execution import execution
@ -216,16 +216,12 @@ import comfy.memory_management
import comfy.model_patcher 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 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(): if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)):
# 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") 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': if args.verbose == 'DEBUG':
comfy_aimdo.control.set_log_debug() comfy_aimdo.control.set_log_debug()
elif args.verbose == 'CRITICAL': elif args.verbose == 'CRITICAL':