From c3e5316b935f3a93ce0879a87222c5ded4450ab9 Mon Sep 17 00:00:00 2001
From: NevermindNilas <nilascontact@gmail.com>
Date: Sun, 17 Mar 2024 17:43:17 +0200
Subject: [PATCH] fixed rife-ncnn

---
 TheAnimeScripter.jsx     |  2 +-
 build.py                 |  4 ++--
 main.py                  |  2 +-
 requirements.txt         |  2 +-
 src/initializeModels.py  |  7 ++++++-
 src/rifencnn/rifencnn.py | 10 +++++++---
 6 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/TheAnimeScripter.jsx b/TheAnimeScripter.jsx
index fd6d61f1..37fa583f 100644
--- a/TheAnimeScripter.jsx
+++ b/TheAnimeScripter.jsx
@@ -2,7 +2,7 @@ var panelGlobal = this;
 var TheAnimeScripter = (function() {
 
     var scriptName = "TheAnimeScripter";
-    var scriptVersion = "v1.4.7";
+    var scriptVersion = "v1.4.8";
 
     // Default Values for the settings
     var outputFolder = app.settings.haveSetting(scriptName, "outputFolder") ? app.settings.getSetting(scriptName, "outputFolder") : "undefined";
diff --git a/build.py b/build.py
index 71837ab9..50d8340d 100644
--- a/build.py
+++ b/build.py
@@ -38,7 +38,7 @@ def create_executable():
     main_path = os.path.join(base_dir, "main.py")
     icon_path = os.path.join(base_dir, "demos", "icon.ico")
     rife_ncnn_models_path = os.path.join(
-        distribution("rife_ncnn_vulkan_python_TAS").locate_file("rife_ncnn_vulkan_python"),
+        distribution("rife_ncnn_vulkan_python").locate_file("rife_ncnn_vulkan_python"),
         "models",
     )
 
@@ -62,7 +62,7 @@ def create_executable():
             "--add-data",
             f"{universal_ncnn_models_path};upscale_ncnn_py/models",
             "--hidden-import",
-            "rife_ncnn_vulkan_python.rife_ncnn_vulkan_python_wrapper",
+            "rife_ncnn_vulkan_python_tas.rife_ncnn_vulkan_python_wrapper",
             "--hidden-import",
             "upscale_ncnn_py.upscale_ncnn_py_wrapper",
             "--collect-all",
diff --git a/main.py b/main.py
index 0e47bf4c..3dfc0164 100644
--- a/main.py
+++ b/main.py
@@ -39,7 +39,7 @@
 else:
     mainPath = os.path.dirname(os.path.abspath(__file__))
 
-scriptVersion = "1.4.7"
+scriptVersion = "1.4.8"
 warnings.filterwarnings("ignore")
 
 
diff --git a/requirements.txt b/requirements.txt
index 18022424..76672a60 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,5 +17,5 @@ realcugan-ncnn-py
 pytorch-lightning
 torch 
 torchvision 
-https://github.com/NevermindNilas/TAS-Modes-Host/releases/download/main/rife_ncnn_vulkan_python_TAS-2.0.6-cp312-cp312-win_amd64.whl
+https://github.com/NevermindNilas/TAS-Modes-Host/releases/download/main/rife_ncnn_vulkan_python-1.2.1-cp312-cp312-win_amd64.whl
 https://github.com/NevermindNilas/TAS-Modes-Host/releases/download/main/upscale_ncnn_py-1.2.0-cp312-none-win_amd64.whl
diff --git a/src/initializeModels.py b/src/initializeModels.py
index a3102171..071c78ce 100644
--- a/src/initializeModels.py
+++ b/src/initializeModels.py
@@ -77,7 +77,12 @@ def intitialize_models(self):
                 from src.rifencnn.rifencnn import rifeNCNN
 
                 interpolate_process = rifeNCNN(
-                    UHD, self.interpolate_method, self.ensemble
+                    UHD,
+                    self.interpolate_method,
+                    self.ensemble,
+                    self.nt,
+                    new_width,
+                    new_height,
                 )
             case "gmfss":
                 from src.gmfss.gmfss_fortuna_union import GMFSS
diff --git a/src/rifencnn/rifencnn.py b/src/rifencnn/rifencnn.py
index 13d5623b..5633e3ed 100644
--- a/src/rifencnn/rifencnn.py
+++ b/src/rifencnn/rifencnn.py
@@ -1,9 +1,12 @@
 from rife_ncnn_vulkan_python import Rife
 
 class rifeNCNN:
-    def __init__(self, UHD, interpolate_method, ensemble=False):
+    def __init__(self, UHD, interpolate_method, ensemble=False, nt=1, width=1920, height=1080):
         self.UHD = UHD
         self.interpolate_method = interpolate_method
+        self.nt = nt
+        self.height = height
+        self.width = width
 
         # Since the built in models folder use the default naming without
         # Ncnn or other suffixes, we need to change the name to match the
@@ -30,13 +33,14 @@ def __init__(self, UHD, interpolate_method, ensemble=False):
             tta_mode=False,
             tta_temporal_mode=False,
             uhd_mode=self.UHD,
-            num_threads=1,
+            num_threads=self.nt,
         )
         
         self.frame1 = None
+        self.shape = (self.height, self.width)
         
     def make_inference(self, timestep):
-        output = self.rife.process(self.frame1, self.frame2, timestep=timestep)
+        output = self.rife.process(self.frame1, self.frame2, timestep=timestep, shape=self.shape)
 
         return output