From 7cd71536bf81f3baf58e7a5b4ccf99c512d7f351 Mon Sep 17 00:00:00 2001 From: Yao Chi <43831444@qq.com> Date: Sun, 25 Aug 2024 21:07:13 +0800 Subject: [PATCH] support two gpus for pipeline parallelism --- comfy/cli_args.py | 2 +- comfy/model_management.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index c98d4dfaf..2885a4e68 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -49,7 +49,7 @@ parser.add_argument("--temp-directory", type=str, default=None, help="Set the Co parser.add_argument("--input-directory", type=str, default=None, help="Set the ComfyUI input directory.") parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.") parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.") -parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.") +parser.add_argument("--cuda-device", type=str, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.") cm_group = parser.add_mutually_exclusive_group() cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).") cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Disable cudaMallocAsync.") diff --git a/comfy/model_management.py b/comfy/model_management.py index 91e692ba2..f7630b62b 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -691,7 +691,9 @@ def text_encoder_offload_device(): return torch.device("cpu") def text_encoder_device(): - if args.gpu_only: + if "," in args.cuda_device: + return torch.device("cuda", 1) + elif args.gpu_only: return get_torch_device() elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM: if should_use_fp16(prioritize_performance=False):