diff --git a/main.py b/main.py index 9686d93..8e6c1eb 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ import collections import configparser import ctypes +import itertools import locale import os import re @@ -613,6 +614,7 @@ def getConfigParams(self) -> param.REConfigParams: return param.REConfigParams( self.varstrModel.get(), self.modelFactors[self.varstrModel.get()], + self.config['Config'].get('ModelDir') or os.path.join(define.APP_PATH, 'models'), self.varintResizeMode.get(), resizeModeValue, self.downsample[self.varintDownsampleIndex.get()][1], @@ -645,7 +647,7 @@ def getOutputPath(self, p: str) -> str: # Because for the WarningNotFoundRE warning message app language # must be initialized and for that config must be initialized # and for that models variable needs to be set -def init_config_and_model_paths() -> tuple[configparser.ConfigParser, set[str], list[str]]: +def init_config_and_model_paths() -> tuple[configparser.ConfigParser, list[str]]: config = configparser.ConfigParser({ 'Upscaler': '', 'ModelDir': '', @@ -669,21 +671,40 @@ def init_config_and_model_paths() -> tuple[configparser.ConfigParser, set[str], config['Config'] = {} config.read(define.APP_CONFIG_PATH) + if config['Config'].get('Upscaler'): + define.RE_PATH = config['Config'].get('Upscaler') + try: - modelFiles = set(os.listdir(config['Config'].get('ModelDir') or os.path.join(define.APP_PATH, 'models'))) - models = sorted( - x for x in set(os.path.splitext(y)[0] for y in modelFiles) - if f'{x}.bin' in modelFiles and f'{x}.param' in modelFiles - ) + modelDir = config['Config'].get('ModelDir') or os.path.join(define.APP_PATH, 'models') + if os.path.splitext(os.path.split(define.RE_PATH)[1])[0] == 'realcugan-ncnn-vulkan': + # 兼容Real-CUGAN的模型文件名格式 + # https://github.com/nihui/realcugan-ncnn-vulkan/blob/395302c5c70f1bff604c974e92e0a87e45c9f9ee/src/main.cpp#L733 + # -m model-path + # -s scale + # -n noise-level + # /upx-conservative.{param,bin} + # /upx-no-denoise.{param,bin} + # /upx-denoisex.{param,bin} + models = [] + for name, scale, noise in itertools.product( + sorted(x for x in os.listdir(modelDir) if os.path.isdir(os.path.join(modelDir, x))), + range(2, 5), + ('conservative', 'no-denoise', *(f'denoise{i}x' for i in range(1, 4))), + ): + if all(os.path.exists(os.path.join(modelDir, name, f'up{scale}x-{noise}.{ext}')) for ext in ('bin', 'param')): + models.append(f'{name}#up{scale}x-{noise}') + else: + modelFiles = set(x for x in os.listdir(modelDir) if os.path.isfile(os.path.join(modelDir, x))) + models = sorted( + x for x in set(os.path.splitext(y)[0] for y in modelFiles) + if f'{x}.bin' in modelFiles and f'{x}.param' in modelFiles + ) except FileNotFoundError: # in case of FileNotFoundError exception, return empty modelFiles and models. # This does not change any behabiour because in this case # we will be showing a warning message and terminate app models = [] - if config['Config'].get('Upscaler'): - define.RE_PATH = config['Config'].get('Upscaler') - i18n.set_current_language(config['Config'].get('AppLanguage')) return config, models diff --git a/param.py b/param.py index 3f3621d..5eb6d36 100644 --- a/param.py +++ b/param.py @@ -10,6 +10,7 @@ class ResizeMode(enum.IntEnum): class REConfigParams(typing.NamedTuple): model: str modelFactor: int + modelDir: str resizeMode: ResizeMode resizeModeValue: int downsample: 'Image._Resample' diff --git a/task.py b/task.py index bcfa556..21ddee8 100644 --- a/task.py +++ b/task.py @@ -79,17 +79,38 @@ def run(self) -> None: for i in range(len(files) - 1): inputPath, outputPath = files[i:(i + 2)] alphaOverridePath = None - cmd = ( - define.RE_PATH, - '-v', - '-i', inputPath, - '-o', outputPath, - '-s', str(self.config.modelFactor), - '-t', str(self.config.tileSize), - '-n', self.config.model, - '-g', 'auto' if self.config.gpuID < 0 else str(self.config.gpuID), - ('-x' if self.config.useTTA else ''), - ) + if os.path.splitext(os.path.split(define.RE_PATH)[1])[0] == 'realcugan-ncnn-vulkan': + model, modelFilename = self.config.model.split('#', 1) + denoiseLevel = { + 'conservative': -1, + 'no-denoise': 0, + **{f'denoise{i}x': i for i in range(1, 4)}, + }[modelFilename.split('-', 1)[1]] + cmd = ( + define.RE_PATH, + '-v', + '-i', inputPath, + '-o', outputPath, + '-s', str(self.config.modelFactor), + '-t', str(self.config.tileSize), + '-m', os.path.join(self.config.modelDir, model), + '-n', str(denoiseLevel), + '-g', 'auto' if self.config.gpuID < 0 else str(self.config.gpuID), + '-c', '1', # accurate sync + *(('-x', ) if self.config.useTTA else ()), + ) + else: + cmd = ( + define.RE_PATH, + '-v', + '-i', inputPath, + '-o', outputPath, + '-s', str(self.config.modelFactor), + '-t', str(self.config.tileSize), + '-n', self.config.model, + '-g', 'auto' if self.config.gpuID < 0 else str(self.config.gpuID), + *(('-x', ) if self.config.useTTA else ()), + ) with subprocess.Popen( cmd, stderr=subprocess.PIPE,