From 5218e5d596ce4f097f41d2f0d304e53f00aff423 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 29 Mar 2023 23:28:21 -0400 Subject: [PATCH] Command line option to set CUDA device. --- main.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/main.py b/main.py index d82d1d6e..c9809137 100644 --- a/main.py +++ b/main.py @@ -18,6 +18,7 @@ if __name__ == "__main__": print("\t--use-split-cross-attention\tUse the split cross attention optimization instead of the sub-quadratic one.\n\t\t\t\t\tIgnored when xformers is used.") print("\t--use-pytorch-cross-attention\tUse the new pytorch 2.0 cross attention function.") print("\t--disable-xformers\t\tdisables xformers") + print("\t--cuda-device 1\t\tSet the id of the cuda device this instance will use.") print() print("\t--highvram\t\t\tBy default models will be unloaded to CPU memory after being used.\n\t\t\t\t\tThis option keeps them in GPU memory.\n") print("\t--normalvram\t\t\tUsed to force normal vram use if lowvram gets automatically enabled.") @@ -31,6 +32,14 @@ if __name__ == "__main__": print("disabling upcasting of attention") os.environ['ATTN_PRECISION'] = "fp16" + try: + index = sys.argv.index('--cuda-device') + device = sys.argv[index + 1] + os.environ['CUDA_VISIBLE_DEVICES'] = device + print("Set cuda device to:", device) + except: + pass + import execution import server import folder_paths