From 4a0f581ecb53ced370d21b0eb347466891316123 Mon Sep 17 00:00:00 2001
From: TNTwise <tntwise3@gmail.com>
Date: Fri, 20 Dec 2024 20:33:55 -0600
Subject: [PATCH] fix install all models

---
 REAL-Video-Enhancer.py |  2 +-
 src/Util.py            |  1 -
 src/ui/DownloadTab.py  | 20 +++++++++++---------
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/REAL-Video-Enhancer.py b/REAL-Video-Enhancer.py
index 2e61d295..d5fb1b3a 100644
--- a/REAL-Video-Enhancer.py
+++ b/REAL-Video-Enhancer.py
@@ -156,7 +156,7 @@ def __init__(self):
             parent=self,
         )
         self.homeTab = HomeTab(parent=self)
-        self.downloadTab = DownloadTab(parent=self)
+        self.downloadTab = DownloadTab(parent=self, backends=self.backends)
         self.settingsTab = SettingsTab(
             parent=self, halfPrecisionSupport=halfPrecisionSupport
         )
diff --git a/src/Util.py b/src/Util.py
index 0cc0c6e8..756f7982 100644
--- a/src/Util.py
+++ b/src/Util.py
@@ -1,4 +1,3 @@
-import cv2
 import os
 import warnings
 import requests
diff --git a/src/ui/DownloadTab.py b/src/ui/DownloadTab.py
index c173fb76..56ff7f73 100644
--- a/src/ui/DownloadTab.py
+++ b/src/ui/DownloadTab.py
@@ -8,6 +8,7 @@
     ncnnUpscaleModels,
     pytorchUpscaleModels,
 )
+from ..constants import MODELS_PATH
 
 
 def downloadModelsBasedOnInstalledBackend(installed_backends: list):
@@ -16,30 +17,28 @@ def downloadModelsBasedOnInstalledBackend(installed_backends: list):
             match backend:
                 case "ncnn":
                     for model in ncnnInterpolateModels:
-                        DownloadModel(model, ncnnInterpolateModels[model][1], "ncnn")
+                        DownloadModel(model, ncnnInterpolateModels[model][1], MODELS_PATH)
                     for model in ncnnUpscaleModels:
-                        DownloadModel(model, ncnnUpscaleModels[model][1], "ncnn")
+                        DownloadModel(model, ncnnUpscaleModels[model][1], MODELS_PATH)
                 case "pytorch":  # no need for tensorrt as it uses pytorch models
                     for model in pytorchInterpolateModels:
                         DownloadModel(
-                            model, pytorchInterpolateModels[model][1], "pytorch"
+                            model, pytorchInterpolateModels[model][1], MODELS_PATH
                         )
                     for model in pytorchUpscaleModels:
-                        DownloadModel(model, pytorchUpscaleModels[model][1], "pytorch")
-        """case "directml":
-            for model in onnxInterpolateModels:
-                DownloadModel(model, onnxInterpolateModels[model][1], "onnx")
-            for model in onnxUpscaleModels:
-                DownloadModel(model, onnxUpscaleModels[model][1], "onnx")"""
+                        DownloadModel(model, pytorchUpscaleModels[model][1], MODELS_PATH)
+
 
 
 class DownloadTab:
     def __init__(
         self,
         parent: QMainWindow,
+        backends: list,
     ):
         self.parent = parent
         self.downloadDeps = DownloadDependencies()
+        self.backends = backends
         self.QButtonConnect()
 
     def QButtonConnect(self):
@@ -61,6 +60,9 @@ def QButtonConnect(self):
                 ["ncnn", "pytorch", "tensorrt", "directml"]
             )
         )
+        self.parent.downloadSomeModelsBasedOnInstalledBackendbtn.clicked.connect(
+            lambda: downloadModelsBasedOnInstalledBackend(self.backends)
+        )
         self.parent.uninstallNCNNBtn.clicked.connect(
             lambda: self.download("ncnn", False)
         )