onnx ¶

def calc_tilesize(self, clip: vs.VideoNode, **kwargs: Any) -> tuple[tuple[int, int], tuple[int, int]]:
    """Reimplementation of vsmlrt.calc_tilesize helper function"""

    from vsmlrt import calc_tilesize

    kwargs = dict(
        tiles=self.tiles,
        tilesize=self.tilesize,
        width=clip.width,
        height=clip.height,
        multiple=1,
        overlap_w=self.overlap_w,
        overlap_h=self.overlap_h,
    ) | kwargs

    return calc_tilesize(**kwargs)

@classmethod
def ensure_obj(
    cls: type[BaseScalerT], scaler: str | type[BaseScalerT] | BaseScalerT | None = None, /,
    func_except: FuncExceptT | None = None
) -> BaseScalerT:
    return _base_ensure_obj(
        cls, (mro := cls.mro())[mro.index(BaseScaler) - 1], scaler, cls._err_class, [], func_except
    )

@classmethod
def from_param(
    cls: type[BaseScalerT], scaler: str | type[BaseScalerT] | BaseScalerT | None = None, /,
    func_except: FuncExceptT | None = None
) -> type[BaseScalerT]:
    return _base_from_param(
        cls, (mro := cls.mro())[mro.index(BaseScaler) - 1], scaler, cls._err_class, [], func_except
    )

def get_clean_kwargs(self, *funcs: Callable[..., Any] | None) -> KwargsT:
    return _clean_self_kwargs(funcs, self)

def get_implemented_funcs(self) -> tuple[Callable[..., Any], ...]:
    return (self.scale, self.supersample)

@inject_kwargs_params
def get_scale_args(
    self, clip: vs.VideoNode, shift: tuple[TopShift, LeftShift] = (0, 0),
    width: int | None = None, height: int | None = None,
    *funcs: Callable[..., Any], **kwargs: Any
) -> KwargsT:
    return (
        dict(
            src_top=shift[0],
            src_left=shift[1]
        )
        | self.get_clean_kwargs(*funcs)
        | dict(width=width, height=height)
        | kwargs
    )

def inference(self, clip: ConstantFormatVideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    from vsmlrt import ArtCNN as mlrt_ArtCNN
    from vsmlrt import ArtCNNModel

    return mlrt_ArtCNN(clip, self.tiles, self.tilesize, self.overlap, ArtCNNModel(self._model), self.backend)

@inject_self.cached.property
def kernel_radius(self) -> int:
    if hasattr(self, '_static_kernel_radius'):
        return ceil(self._static_kernel_radius)
    raise CustomNotImplementedError('kernel_radius is not implemented!', self.__class__)

@inject_self.cached
def multi(
    self, clip: vs.VideoNode, multi: float = 2.0, shift: tuple[TopShift, LeftShift] = (0, 0), **kwargs: Any
) -> vs.VideoNode:

    import warnings

    warnings.warn('The "multi" method is deprecated. Use "supersample" instead.', DeprecationWarning)

    return self.supersample(clip, multi, shift, **kwargs)

def postprocess_clip(self, clip: vs.VideoNode, input_clip: vs.VideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    """Handles postprocessing of the model's output after inference."""

    return depth(
        clip, input_clip, dither_type=DitherType.ORDERED if 0 in {clip.width, clip.height} else DitherType.AUTO
    )

def preprocess_clip(self, clip: vs.VideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    return super().preprocess_clip(get_y(clip), **kwargs)

@inject_self.cached.property
def pretty_string(self) -> str:
    return self._pretty_string()

@inject_self.cached
def scale(
    self,
    clip: vs.VideoNode,
    width: int | None = None,
    height: int | None = None,
    shift: tuple[float, float] = (0, 0),
    **kwargs: Any
) -> ConstantFormatVideoNode:
    """
    Scale the given clip using the ONNX model.

    :param clip:        The input clip to be scaled.
    :param width:       The target width for scaling. If None, the width of the input clip will be used.
    :param height:      The target height for scaling. If None, the height of the input clip will be used.
    :param shift:       A tuple representing the shift values for the x and y axes.
    :param **kwargs:    Additional arguments to be passed to the `preprocess_clip`, `postprocess_clip`,
                        `inference`, and `_final_scale` methods.
                        Use the prefix `preprocess_` or `postprocess_` to pass an argument to the respective method.
                        Use the prefix `inference_` to pass an argument to the inference method.

    :return:            The scaled clip.
    """
    from vsmlrt import Backend

    assert check_variable_format(clip, self.__class__)

    width, height = self._wh_norm(clip, width, height)

    preprocess_kwargs = dict[str, Any]()
    postprocess_kwargs = dict[str, Any]()
    inference_kwargs = dict[str, Any]()

    for k in kwargs.copy():
        for prefix, ckwargs in zip(
            ("preprocess_", "postprocess_", "inference_"),
            (preprocess_kwargs, postprocess_kwargs, inference_kwargs)
        ):
            if k.startswith(prefix):
                ckwargs[k.removeprefix(prefix)] = kwargs.pop(k)
                break

    wclip = self.preprocess_clip(clip, **preprocess_kwargs)

    if 0 not in {clip.width, clip.height}:
        scaled = self.inference(wclip, **inference_kwargs)
    else:
        if not isinstance(self.backend, Backend.TRT):
            raise CustomValueError(
                "Variable resolution clips can only be processed with TRT Backend!", self.__class__, self.backend
            )

        warning(f"{self.__class__.__name__}: Variable resolution clip detected!")

        if self.backend.static_shape:
            warning("static_shape is True, setting it to False...")
            self.backend.static_shape = False

        if not self.backend.max_shapes:
            warning("max_shapes is None, setting it to (1936, 1088). You may want to adjust it...")
            self.backend.max_shapes = (1936, 1088)

        if not self.backend.opt_shapes:
            warning("opt_shapes is None, setting it to (64, 64). You may want to adjust it...")
            self.backend.opt_shapes = (64, 64)

        scaled = ProcessVariableResClip[ConstantFormatVideoNode].from_func(
            wclip, lambda c: self.inference(c, **inference_kwargs), False, wclip.format, self.max_instances
        )

    scaled = self.postprocess_clip(scaled, clip, **postprocess_kwargs)

    return self._finish_scale(scaled, clip, width, height, shift, **kwargs)

@inject_self.cached
def supersample(
    self, clip: vs.VideoNode, rfactor: float = 2.0, shift: tuple[TopShift, LeftShift] = (0, 0), **kwargs: Any
) -> vs.VideoNode:
    assert check_variable_resolution(clip, self.multi)

    dst_width, dst_height = ceil(clip.width * rfactor), ceil(clip.height * rfactor)

    if max(dst_width, dst_height) <= 0.0:
        raise CustomValueError(
            'Multiplying the resolution by "rfactor" must result in a positive resolution!', self.supersample, rfactor
        )

    return self.scale(clip, dst_width, dst_height, shift, **kwargs)

def calc_tilesize(self, clip: vs.VideoNode, **kwargs: Any) -> tuple[tuple[int, int], tuple[int, int]]:
    """Reimplementation of vsmlrt.calc_tilesize helper function"""

    from vsmlrt import calc_tilesize

    kwargs = dict(
        tiles=self.tiles,
        tilesize=self.tilesize,
        width=clip.width,
        height=clip.height,
        multiple=1,
        overlap_w=self.overlap_w,
        overlap_h=self.overlap_h,
    ) | kwargs

    return calc_tilesize(**kwargs)

@classmethod
def ensure_obj(
    cls: type[BaseScalerT], scaler: str | type[BaseScalerT] | BaseScalerT | None = None, /,
    func_except: FuncExceptT | None = None
) -> BaseScalerT:
    return _base_ensure_obj(
        cls, (mro := cls.mro())[mro.index(BaseScaler) - 1], scaler, cls._err_class, [], func_except
    )

@classmethod
def from_param(
    cls: type[BaseScalerT], scaler: str | type[BaseScalerT] | BaseScalerT | None = None, /,
    func_except: FuncExceptT | None = None
) -> type[BaseScalerT]:
    return _base_from_param(
        cls, (mro := cls.mro())[mro.index(BaseScaler) - 1], scaler, cls._err_class, [], func_except
    )

def get_clean_kwargs(self, *funcs: Callable[..., Any] | None) -> KwargsT:
    return _clean_self_kwargs(funcs, self)

def get_implemented_funcs(self) -> tuple[Callable[..., Any], ...]:
    return (self.scale, self.supersample)

@inject_kwargs_params
def get_scale_args(
    self, clip: vs.VideoNode, shift: tuple[TopShift, LeftShift] = (0, 0),
    width: int | None = None, height: int | None = None,
    *funcs: Callable[..., Any], **kwargs: Any
) -> KwargsT:
    return (
        dict(
            src_top=shift[0],
            src_left=shift[1]
        )
        | self.get_clean_kwargs(*funcs)
        | dict(width=width, height=height)
        | kwargs
    )

def inference(self, clip: ConstantFormatVideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    from vsmlrt import ArtCNN as mlrt_ArtCNN
    from vsmlrt import ArtCNNModel

    return mlrt_ArtCNN(clip, self.tiles, self.tilesize, self.overlap, ArtCNNModel(self._model), self.backend)

@inject_self.cached.property
def kernel_radius(self) -> int:
    if hasattr(self, '_static_kernel_radius'):
        return ceil(self._static_kernel_radius)
    raise CustomNotImplementedError('kernel_radius is not implemented!', self.__class__)

@inject_self.cached
def multi(
    self, clip: vs.VideoNode, multi: float = 2.0, shift: tuple[TopShift, LeftShift] = (0, 0), **kwargs: Any
) -> vs.VideoNode:

    import warnings

    warnings.warn('The "multi" method is deprecated. Use "supersample" instead.', DeprecationWarning)

    return self.supersample(clip, multi, shift, **kwargs)

def postprocess_clip(self, clip: vs.VideoNode, input_clip: vs.VideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    """Handles postprocessing of the model's output after inference."""

    return depth(
        clip, input_clip, dither_type=DitherType.ORDERED if 0 in {clip.width, clip.height} else DitherType.AUTO
    )

def preprocess_clip(self, clip: vs.VideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    assert check_variable_format(clip, self.__class__)

    if clip.format.subsampling_h != 0 or clip.format.subsampling_w != 0:
        chroma_scaler = Kernel.ensure_obj(kwargs.pop("chroma_scaler", Bilinear))

        clip = chroma_scaler.resample(
            clip, clip.format.replace(
                subsampling_h=0, subsampling_w=0,
                sample_type=vs.FLOAT, bits_per_sample=16 if self.backend.fp16 else 32
            )
        )
        return limiter(clip, func=self.__class__)

    return super().preprocess_clip(clip, **kwargs)

@inject_self.cached.property
def pretty_string(self) -> str:
    return self._pretty_string()

@inject_self.cached
def scale(
    self,
    clip: vs.VideoNode,
    width: int | None = None,
    height: int | None = None,
    shift: tuple[float, float] = (0, 0),
    **kwargs: Any
) -> ConstantFormatVideoNode:
    """
    Scale the given clip using the ONNX model.

    :param clip:        The input clip to be scaled.
    :param width:       The target width for scaling. If None, the width of the input clip will be used.
    :param height:      The target height for scaling. If None, the height of the input clip will be used.
    :param shift:       A tuple representing the shift values for the x and y axes.
    :param **kwargs:    Additional arguments to be passed to the `preprocess_clip`, `postprocess_clip`,
                        `inference`, and `_final_scale` methods.
                        Use the prefix `preprocess_` or `postprocess_` to pass an argument to the respective method.
                        Use the prefix `inference_` to pass an argument to the inference method.

    :return:            The scaled clip.
    """
    from vsmlrt import Backend

    assert check_variable_format(clip, self.__class__)

    width, height = self._wh_norm(clip, width, height)

    preprocess_kwargs = dict[str, Any]()
    postprocess_kwargs = dict[str, Any]()
    inference_kwargs = dict[str, Any]()

    for k in kwargs.copy():
        for prefix, ckwargs in zip(
            ("preprocess_", "postprocess_", "inference_"),
            (preprocess_kwargs, postprocess_kwargs, inference_kwargs)
        ):
            if k.startswith(prefix):
                ckwargs[k.removeprefix(prefix)] = kwargs.pop(k)
                break

    wclip = self.preprocess_clip(clip, **preprocess_kwargs)

    if 0 not in {clip.width, clip.height}:
        scaled = self.inference(wclip, **inference_kwargs)
    else:
        if not isinstance(self.backend, Backend.TRT):
            raise CustomValueError(
                "Variable resolution clips can only be processed with TRT Backend!", self.__class__, self.backend
            )

        warning(f"{self.__class__.__name__}: Variable resolution clip detected!")

        if self.backend.static_shape:
            warning("static_shape is True, setting it to False...")
            self.backend.static_shape = False

        if not self.backend.max_shapes:
            warning("max_shapes is None, setting it to (1936, 1088). You may want to adjust it...")
            self.backend.max_shapes = (1936, 1088)

        if not self.backend.opt_shapes:
            warning("opt_shapes is None, setting it to (64, 64). You may want to adjust it...")
            self.backend.opt_shapes = (64, 64)

        scaled = ProcessVariableResClip[ConstantFormatVideoNode].from_func(
            wclip, lambda c: self.inference(c, **inference_kwargs), False, wclip.format, self.max_instances
        )

    scaled = self.postprocess_clip(scaled, clip, **postprocess_kwargs)

    return self._finish_scale(scaled, clip, width, height, shift, **kwargs)

@inject_self.cached
def supersample(
    self, clip: vs.VideoNode, rfactor: float = 2.0, shift: tuple[TopShift, LeftShift] = (0, 0), **kwargs: Any
) -> vs.VideoNode:
    assert check_variable_resolution(clip, self.multi)

    dst_width, dst_height = ceil(clip.width * rfactor), ceil(clip.height * rfactor)

    if max(dst_width, dst_height) <= 0.0:
        raise CustomValueError(
            'Multiplying the resolution by "rfactor" must result in a positive resolution!', self.supersample, rfactor
        )

    return self.scale(clip, dst_width, dst_height, shift, **kwargs)