diff --git a/modules/CUGAN.py b/modules/CUGAN.py index c7aa0d4d..fc48c8b2 100644 --- a/modules/CUGAN.py +++ b/modules/CUGAN.py @@ -46,7 +46,7 @@ def modelOptions(self): pass self.ui.Rife_Model.setEnabled(True) self.ui.RifeStart.clicked.connect( - lambda: upscale.start_upscale(self, "realcugan-ncnn-vulkan") + lambda: upscale.start_upscale(self, "realcugan-ncnn-python") ) self.ui.Rife_Times.clear() self.ui.Rife_Times.currentIndexChanged.connect(self.cuganDenoiseLevel) diff --git a/src/runAI/workers.py b/src/runAI/workers.py index 146984f6..51dd741d 100644 --- a/src/runAI/workers.py +++ b/src/runAI/workers.py @@ -925,6 +925,8 @@ def realESRGAN(self): bool(settings.HalfPrecision), method=self.main.AI, threads=int(settings.VRAM), + ncnn_gpu=self.main.ui.gpuIDSpinBox.value(), + ) @@ -944,6 +946,26 @@ def realESRGAN(self): bool(settings.HalfPrecision), method=self.main.AI, threads=int(settings.VRAM), + ncnn_gpu=self.main.ui.gpuIDSpinBox.value(), + ) + if self.main.AI == "realcugan-ncnn-python": + output_file = returnOutputFile( + self.main, self.main.videoName, self.main.encoder + ) + + #model = os.path.join(f"{settings.ModelDir}","realsr",f"models-{self.main.ui.Rife_Model.currentText()}","x4") + + self.main.renderAI = RenderCUDA.Upscaling( + self.main, + self.main.input_file, + output_file, + int(self.main.ui.Rife_Times.currentText()[0]), + + half=bool(settings.HalfPrecision), + method=self.main.AI, + threads=int(settings.VRAM), + ncnn_gpu=self.main.ui.gpuIDSpinBox.value(), + cugan_noise=str(self.main.ui.denoiseLevelSpinBox.value()) ) if self.main.AI == "custom-models-ncnn-python": diff --git a/src/torch/RenderCUDA.py b/src/torch/RenderCUDA.py index 20fb936b..372bfef0 100644 --- a/src/torch/RenderCUDA.py +++ b/src/torch/RenderCUDA.py @@ -14,9 +14,10 @@ try: from .rife.rife import * from .UpscaleImage import UpscaleCUDA - from .UpscaleImageNCNN import UpscaleNCNN + except: - from .UpscaleImageNCNN import UpscaleNCNN + pass +from .UpscaleImageNCNN import UpscaleNCNN, UpscaleCuganNCNN try: from src.torch.gmfss.gmfss_fortuna_union import GMFSS except: @@ -299,10 +300,12 @@ def __init__( input_file, output_file, resIncrease, - model_path, - half, + model_path="", + half=True, method="cuda", threads=2, + cugan_noise=0, + ncnn_gpu=0 ): super(Upscaling, self).__init__( main, @@ -317,15 +320,25 @@ def __init__( self.resIncrease = resIncrease self.threads = threads self.frame = 0 + self.cugan_noise = cugan_noise + self.ncnn_gpu=ncnn_gpu self.handleModel() def handleModel(self): if "cuda" in self.method and "ncnn" not in self.method: self.upscaleMethod = UpscaleCUDA( self.originalWidth, self.originalHeight, self.model_path, self.half ) - if "ncnn" in self.method: + if "ncnn" in self.method and not "cugan" in self.method: self.upscaleMethod = UpscaleNCNN( - model=self.model_path, num_threads=self.threads, scale=self.resIncrease + gpuid=self.ncnn_gpu,model=self.model_path, num_threads=self.threads, scale=self.resIncrease + ) + if "cugan" in self.method and "ncnn" in self.method: + self.upscaleMethod = UpscaleCuganNCNN( + gpuid=self.ncnn_gpu, + models_path = os.path.join(f"{thisdir}","models","realcugan","models-se"), + model="models-se", + num_threads=self.threads, + scale=self.resIncrease ) def procUpscaleThread(self): diff --git a/src/torch/UpscaleImageNCNN.py b/src/torch/UpscaleImageNCNN.py index 38c49cd3..670621fd 100644 --- a/src/torch/UpscaleImageNCNN.py +++ b/src/torch/UpscaleImageNCNN.py @@ -1,18 +1,29 @@ import numpy as np from upscale_ncnn_py import UPSCALE +from realcugan_ncnn_py import Realcugan import cv2 class UpscaleNCNN: - def __init__(self, model, num_threads, scale): + def __init__(self, model, num_threads, scale, gpuid=0): self.model = UPSCALE( - gpuid=0, model_str=model, num_threads=num_threads, scale=scale + gpuid=gpuid, model_str=model, num_threads=num_threads, scale=scale ) def UpscaleImage(self, image): return self.model.process_cv2(image) - +class UpscaleCuganNCNN: + def __init__(self,model="models-se",models_path="",num_threads=2,scale=2, gpuid=0,noise=0): + self.model = Realcugan(gpuid=gpuid, + models_path=models_path, + model=model, + scale=scale, + num_threads=num_threads, + noise=noise) + + def UpscaleImage(self,image): + return self.model.process_cv2(image) if __name__ == "__main__": from PIL import Image import matplotlib.pyplot as plt diff --git a/src/torch/__pycache__/RenderCUDA.cpython-311.pyc b/src/torch/__pycache__/RenderCUDA.cpython-311.pyc index 9552791e..e7efa120 100644 Binary files a/src/torch/__pycache__/RenderCUDA.cpython-311.pyc and b/src/torch/__pycache__/RenderCUDA.cpython-311.pyc differ