From c7b8556856286aa35f152349d520db31834c49a0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=95=D0=B2=D0=B3=D0=B5=D0=BD=D0=B8=D0=B9=20=D0=93=D1=83?=
 =?UTF-8?q?=D1=80=D1=8C=D0=B5=D0=B2=20=7C=20Eugene=20Gourieff=20=7C=20?=
 =?UTF-8?q?=E5=8F=A4=E4=BB=81?= <gourieff@gmail.com>
Date: Tue, 8 Oct 2024 03:39:49 +0700
Subject: [PATCH] FIX.1: onnx, onnxruntime-gpu stable versions

Issues: #483 #486 #502 #536 #538 #539
---
 install.py       | 27 +++++++++------------------
 requirements.txt |  2 +-
 2 files changed, 10 insertions(+), 19 deletions(-)

diff --git a/install.py b/install.py
index 70d34e7..596596c 100644
--- a/install.py
+++ b/install.py
@@ -21,20 +21,6 @@ req_file = os.path.join(BASE_PATH, "requirements.txt")
 
 models_dir = os.path.join(models_path, "insightface")
 
-# DEPRECATED:
-# models_dir_old = os.path.join(models_path, "roop")
-# if os.path.exists(models_dir_old):
-#     if not os.listdir(models_dir_old) and (not os.listdir(models_dir) or not os.path.exists(models_dir)):
-#         os.rename(models_dir_old, models_dir)
-#     else:
-#         import shutil
-#         for file in os.listdir(models_dir_old):
-#             shutil.move(os.path.join(models_dir_old, file), os.path.join(models_dir, file))
-#         try:
-#             os.rmdir(models_dir_old)
-#         except Exception as e:
-#             print(f"OSError: {e}")
-            
 model_url = "https://huggingface.co/datasets/Gourieff/ReActor/resolve/main/models/inswapper_128.onnx"
 model_name = os.path.basename(model_url)
 model_path = os.path.join(models_dir, model_name)
@@ -117,12 +103,17 @@ with open(req_file) as file:
                 last_device = "CPU"
         with open(os.path.join(BASE_PATH, "last_device.txt"), "w") as txt:
             txt.write(last_device)
-        if cuda_version is not None and float(cuda_version)>=12: # CU12
-            if not is_installed(ort,"1.17.1",False):
+        if cuda_version is not None:
+            if float(cuda_version)>=12: # CU12.x
+                extra_index_url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/"
+            else: # CU11.8
+                extra_index_url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-11/pypi/simple"
+            if not is_installed(ort,"1.17.1",True):
                 install_count += 1
+                ort = "onnxruntime-gpu==1.17.1"
                 pip_uninstall("onnxruntime", "onnxruntime-gpu")
-                pip_install(ort,"--extra-index-url","https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/")
-        elif not is_installed(ort,"1.16.1",False):
+                pip_install(ort,"--extra-index-url",extra_index_url)
+        elif not is_installed(ort,"1.18.1",False):
             install_count += 1
             pip_install(ort, "-U")
     except Exception as e:
diff --git a/requirements.txt b/requirements.txt
index 14c8871..0e1bd70 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
 albumentations==1.4.3
 insightface==0.7.3
-onnx>=1.14.0
+onnx==1.16.1
 opencv-python>=4.7.0.72