This commit is contained in:
Haoming 2025-12-06 23:27:32 +08:00
parent 324434a3f7
commit c2cd85e4f1
3 changed files with 12 additions and 6 deletions

View File

@ -63,6 +63,7 @@ parser.add_argument("--sage", action="store_true", help="install sageattention")
parser.add_argument("--flash", action="store_true", help="install flash_attn")
parser.add_argument("--nunchaku", action="store_true", help="install nunchaku for SVDQ inference")
parser.add_argument("--bnb", action="store_true", help="install bitsandbytes for 4-bit inference")
parser.add_argument("--onnxruntime-gpu", action="store_true", help="install nightly onnxruntime-gpu with cu130 support")
parser.add_argument("--disable-xformers", action="store_true")
parser.add_argument("--disable-sage", action="store_true")

View File

@ -10,7 +10,6 @@ parser.add_argument("-f", action="store_true", help=argparse.SUPPRESS)
parser.add_argument("--update-all-extensions", action="store_true", help="launch.py argument: download updates for all extensions when starting the program")
parser.add_argument("--skip-python-version-check", action="store_true", help="launch.py argument: do not check python version")
parser.add_argument("--skip-torch-cuda-test", action="store_true", help="launch.py argument: do not check if CUDA is able to work properly")
parser.add_argument("--onnxruntime-gpu", action="store_true", help="launch.py argument: install the appropriate version of onnxruntime-gpu with cu128 support")
parser.add_argument("--reinstall-xformers", action="store_true", help="launch.py argument: install the appropriate version of xformers even if you have some version already installed")
parser.add_argument("--reinstall-torch", action="store_true", help="launch.py argument: install the appropriate version of torch even if you have some version already installed")
parser.add_argument("--log-startup", action="store_true", help="launch.py argument: print a detailed log of what's happening at startup")

View File

@ -124,9 +124,12 @@ def run_pip(command, desc=None, live=default_command_live):
return run(f'"{python}" -m pip {command} --prefer-binary{index_url_line}', desc=f"Installing {desc}", errdesc=f"Couldn't install {desc}", live=live)
def check_run_python(code: str) -> bool:
def check_run_python(code: str, *, return_error: bool = False) -> bool | tuple[bool, str]:
result = subprocess.run([python, "-c", code], capture_output=True, shell=False)
return result.returncode == 0
if return_error:
return result.returncode == 0, result.stderr
else:
return result.returncode == 0
def git_fix_workspace(*args, **kwargs):
@ -311,7 +314,10 @@ def prepare_environment():
startup_timer.record("install torch")
if not args.skip_torch_cuda_test:
if not check_run_python("import torch; assert torch.cuda.is_available()"):
success, err = check_run_python("import torch; assert torch.cuda.is_available()", return_error=True)
if not success:
if "older driver" in str(err).lower():
raise SystemError("Please update your GPU driver to support cu130 ; or manually install older PyTorch")
raise RuntimeError("PyTorch is not able to access CUDA")
startup_timer.record("torch GPU test")
@ -417,8 +423,8 @@ def prepare_environment():
startup_timer.record("install requirements")
if args.onnxruntime_gpu and not is_installed("onnxruntime-gpu"):
# https://onnxruntime.ai/docs/install/#install-onnx-runtime-gpu-cuda-12x
onnxruntime_package = os.environ.get("ONNX_PACKAGE", "onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/")
# https://onnxruntime.ai/docs/install/#nightly-for-cuda-13x
onnxruntime_package = os.environ.get("ONNX_PACKAGE", "onnxruntime-gpu --pre --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-13-nightly/pypi/simple/")
run_pip(f"install {onnxruntime_package}", "onnxruntime-gpu")
startup_timer.record("install onnxruntime-gpu")