implement CUDA device selection, --device-id arg

This commit is contained in:
Extraltodeus 2022-10-22 00:12:46 +02:00 committed by GitHub
parent 57eb54b838
commit 29bfacd63c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -80,6 +80,7 @@ parser.add_argument('--vae-path', type=str, help='Path to Variational Autoencode
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)
cmd_opts = parser.parse_args()
restricted_opts = [