From 2342761fa10b6004341414a120736ef76cfd12d2 Mon Sep 17 00:00:00 2001 From: lllyasviel Date: Mon, 6 Nov 2023 00:44:38 -0800 Subject: [PATCH] allow set device --- launch.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/launch.py b/launch.py index bb02b09f..e3247448 100644 --- a/launch.py +++ b/launch.py @@ -5,11 +5,10 @@ import sys print('[System ARGV] ' + str(sys.argv)) root = os.path.dirname(os.path.abspath(__file__)) -sys.path.append(root) -os.chdir(root) backend_path = os.path.join(root, 'backend', 'headless') -if backend_path not in sys.path: - sys.path.append(backend_path) +sys.path += [root, backend_path] + +os.chdir(root) os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" @@ -87,14 +86,21 @@ def download_models(): return -def ini_cbh_args(): +def ini_fcbh_args(): from args_manager import args return args prepare_environment() build_launcher() -ini_cbh_args() +args = ini_fcbh_args() + + +if args.cuda_device is not None: + os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) + print("Set device to:", args.cuda_device) + + download_models() from webui import *