allow set device

This commit is contained in:
lllyasviel
2023-11-06 00:44:38 -08:00
parent 10574f1cc2
commit 2342761fa1
+12 -6
View File
@@ -5,11 +5,10 @@ import sys
print('[System ARGV] ' + str(sys.argv)) print('[System ARGV] ' + str(sys.argv))
root = os.path.dirname(os.path.abspath(__file__)) root = os.path.dirname(os.path.abspath(__file__))
sys.path.append(root)
os.chdir(root)
backend_path = os.path.join(root, 'backend', 'headless') backend_path = os.path.join(root, 'backend', 'headless')
if backend_path not in sys.path: sys.path += [root, backend_path]
sys.path.append(backend_path)
os.chdir(root)
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
@@ -87,14 +86,21 @@ def download_models():
return return
def ini_cbh_args(): def ini_fcbh_args():
from args_manager import args from args_manager import args
return args return args
prepare_environment() prepare_environment()
build_launcher() 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() download_models()
from webui import * from webui import *