diff --git a/scripts/vsmlrt.py b/scripts/vsmlrt.py index 2061cdc..ecab3ba 100644 --- a/scripts/vsmlrt.py +++ b/scripts/vsmlrt.py @@ -1,10 +1,11 @@ -__version__ = "3.4.2" +__version__ = "3.5.0" __all__ = [ "Backend", "Waifu2x", "Waifu2xModel", "DPIR", "DPIRModel", - "RealESRGANv2", "RealESRGANv2Model" + "RealESRGANv2", "RealESRGANv2Model", + "CUGAN" ] import copy @@ -397,6 +398,91 @@ def RealESRGANv2( return clip +def CUGAN( + clip: vs.VideoNode, + noise: typing.Literal[-1, 0, 1, 2, 3] = -1, + scale: typing.Literal[2, 3, 4] = 2, + tiles: typing.Optional[typing.Union[int, typing.Tuple[int, int]]] = None, + tilesize: typing.Optional[typing.Union[int, typing.Tuple[int, int]]] = None, + overlap: typing.Optional[typing.Union[int, typing.Tuple[int, int]]] = None, + backend: backendT = Backend.OV_CPU(), + preprocess: bool = True +) -> vs.VideoNode: + + func_name = "vsmlrt.CUGAN" + + if not isinstance(clip, vs.VideoNode): + raise TypeError(f'{func_name}: "clip" must be a clip!') + + if clip.format.sample_type != vs.FLOAT or clip.format.bits_per_sample != 32: + raise ValueError(f"{func_name}: only constant format 32 bit float input supported") + + if not isinstance(noise, int) or noise not in range(-1, 4): + raise ValueError(f'{func_name}: "noise" must be -1, 0, 1, 2, or 3') + + if not isinstance(scale, int) or scale not in (2, 3, 4): + raise ValueError(f'{func_name}: "scale" must be 2, 3 or 4') + + if scale != 2 and noise in [1, 2]: + raise ValueError( + f'{func_name}: "scale={scale}" model' + f' does not support noise reduction level {noise}' + ) + + if clip.format.id != vs.RGBS: + raise ValueError(f'{func_name}: "clip" must be of RGBS format') + + if overlap is None: + overlap_w = overlap_h = 4 + elif isinstance(overlap, int): + overlap_w = overlap_h = overlap + else: + overlap_w, overlap_h = overlap + + multiple = 2 + + width, height = clip.width, clip.height + + (tile_w, tile_h), (overlap_w, overlap_h) = calc_tilesize( + tiles=tiles, tilesize=tilesize, + width=clip.width, height=clip.height, + multiple=multiple, + overlap_w=overlap_w, overlap_h=overlap_h + ) + + if tile_w % multiple != 0 or tile_h % multiple != 0: + raise ValueError( + f'{func_name}: tile size must be divisible by {multiple} ({tile_w}, {tile_h})' + ) + + channels = 3 + + backend = init_backend( + backend=backend, + channels=channels, + trt_max_shapes=(tile_w, tile_h) + ) + + folder_path = os.path.join(models_path, "cugan") + + if noise == -1: + model_name = f"up{scale}x-latest-no-denoise.onnx" + elif noise == 0: + model_name = f"up{scale}x-latest-conservative.onnx" + else: + model_name = f"up{scale}x-latest-denoise{noise}x.onnx" + + network_path = os.path.join(folder_path, model_name) + + clip = inference( + clips=[clip], network_path=network_path, + overlap=(overlap_w, overlap_h), tilesize=(tile_w, tile_h), + backend=backend + ) + + return clip + + def get_engine_path( network_path: str, opt_shapes: typing.Tuple[int, int],