Implementation of CUDA device id selection (--device-id 0/1/2)

modules/devices.py

@@ -1,18 +1,33 @@
+import sys, os, shlex
 import contextlib
-
 import torch
-
 from modules import errors
 
 # has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
 has_mps = getattr(torch, 'has_mps', False)
 
 cpu = torch.device("cpu")
 
+def extract_device_id(args, name):
+    for x in range(len(args)):
+        if name in args[x]: return args[x+1]
+    return None
 
 def get_optimal_device():
     if torch.cuda.is_available():
-        return torch.device("cuda")
+        # CUDA device selection support:
+        if "shared" not in sys.modules:
+            commandline_args = os.environ.get('COMMANDLINE_ARGS', "") #re-parse the commandline arguments because using the shared.py module creates an import loop.
+            sys.argv += shlex.split(commandline_args)
+            device_id = extract_device_id(sys.argv, '--device-id')
+        else:
+            device_id = shared.cmd_opts.device_id
+
+        if device_id is not None:
+            cuda_device = f"cuda:{device_id}"
+            return torch.device(cuda_device)
+        else:
+            return torch.device("cuda")
 
     if has_mps:
         return torch.device("mps")

modules/shared.py

@@ -79,6 +79,7 @@
 parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)
 parser.add_argument("--api", action='store_true', help="use api=True to launch the api with the webui")
 parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the api instead of the webui")
+parser.add_argument("--device-id", type=str, help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)", default=None)
 parser.add_argument("--browse-all-images", action='store_true', help="Allow browsing all images by Image Browser", default=False)
 
 cmd_opts = parser.parse_args()