From 7ff1ef77dd22f7b38612f91b389237a5dbef2474 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 23 Jan 2023 17:17:31 +0300
Subject: add a message about new torch/xformers version and a way to upgrade
 by specifying a commandline flag

---
 launch.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'launch.py')

diff --git a/launch.py b/launch.py
index 51c322c0..e7a0b50c 100644
--- a/launch.py
+++ b/launch.py
@@ -208,6 +208,7 @@ def prepare_environment():
     sys.argv, _ = extract_arg(sys.argv, '-f')
     sys.argv, skip_torch_cuda_test = extract_arg(sys.argv, '--skip-torch-cuda-test')
     sys.argv, reinstall_xformers = extract_arg(sys.argv, '--reinstall-xformers')
+    sys.argv, reinstall_torch = extract_arg(sys.argv, '--reinstall-torch')
     sys.argv, update_check = extract_arg(sys.argv, '--update-check')
     sys.argv, run_tests, test_dir = extract_opt(sys.argv, '--tests')
     sys.argv, skip_install = extract_arg(sys.argv, '--skip-install')
@@ -219,7 +220,7 @@ def prepare_environment():
     print(f"Python {sys.version}")
     print(f"Commit hash: {commit}")
     
-    if not is_installed("torch") or not is_installed("torchvision"):
+    if reinstall_torch or not is_installed("torch") or not is_installed("torchvision"):
         run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch")
 
     if not skip_torch_cuda_test:
-- 
cgit v1.2.3