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 = {
        "tiles": self.tiles,
        "tilesize": self.tilesize,
        "width": clip.width,
        "height": clip.height,
        "multiple": self.multiple,
        "overlap_w": self.overlap_w,
        "overlap_h": self.overlap_h,
    } | kwargs

    return calc_tilesize(**kwargs)

@classmethod
def ensure_obj(
    cls,
    scaler: str | type[Self] | Self | None = None,
    /,
    func_except: FuncExcept | None = None,
) -> Self:
    """
    Ensure that the input is a scaler instance, resolving it if necessary.

    Args:
        scaler: Scaler identifier (string, class, or instance).
        func_except: Function returned for custom error handling.

    Returns:
        Scaler instance.
    """
    return _base_ensure_obj(cls, scaler, func_except)

@classmethod
def from_param(
    cls,
    scaler: str | type[Self] | Self | None = None,
    /,
    func_except: FuncExcept | None = None,
) -> type[Self]:
    """
    Resolve and return a scaler type from a given input (string, type, or instance).

    Args:
        scaler: Scaler identifier (string, class, or instance).
        func_except: Function returned for custom error handling.

    Returns:
        Resolved scaler type.
    """
    return _base_from_param(cls, scaler, cls._err_class, func_except)

def get_scale_args(
    self,
    clip: vs.VideoNode,
    shift: tuple[TopShift, LeftShift] = (0, 0),
    width: int | None = None,
    height: int | None = None,
    **kwargs: Any,
) -> dict[str, Any]:
    """
    Generate the keyword arguments used for scaling.

    Args:
        clip: The source clip.
        shift: Subpixel shift (top, left).
        width: Target width.
        height: Target height.
        **kwargs: Extra parameters to merge.

    Returns:
        Final dictionary of keyword arguments for the scale function.
    """
    return {"width": width, "height": height, "src_top": shift[0], "src_left": shift[1]} | self.kwargs | kwargs

@classproperty
@classmethod
def implemented_funcs(cls) -> frozenset[str]:
    """
    Returns a set of function names that are implemented in the current class and the parent classes.

    These functions determine which keyword arguments will be extracted from the __init__ method.

    Returns:
        Frozen set of function names.
    """
    return frozenset(func for klass in cls.mro() for func in getattr(klass, "_implemented_funcs", ()))

def inference(self, clip: ConstantFormatVideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    """
    Runs inference on the given video clip using the configured model and backend.
    """

    from vsmlrt import inference

    tilesize, overlaps = self.calc_tilesize(clip)

    debug(f"{self}: Passing clip to inference: {clip.format!r}")
    debug(f"{self}: Passing model: {self.model}")
    debug(f"{self}: Passing tiles size: {tilesize}")
    debug(f"{self}: Passing overlaps: {overlaps}")
    debug(f"{self}: Passing extra kwargs: {kwargs}")

    return inference(clip, self.model, overlaps, tilesize, self.backend, **kwargs)

@BaseScalerMeta.cachedproperty
def kernel_radius(self) -> int:
    """
    Return the effective kernel radius for the scaler.

    Raises:
        CustomNotImplementedError: If no kernel radius is defined.

    Returns:
        Kernel radius.
    """
    ...

def postprocess_clip(self, clip: vs.VideoNode, input_clip: vs.VideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    """
    Handles postprocessing of the model's output after inference.
    """
    debug(f"{self}.post: Before pp; Clip format is {clip.format!r}")

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

    debug(f"{self}.post: After pp; Clip format is {clip.format!r}")

    return clip

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

@BaseScalerMeta.cachedproperty
def pretty_string(self) -> str:
    """
    Cached property returning a user-friendly string representation.

    Returns:
        Pretty-printed string with arguments.
    """
    return self._pretty_string()

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.

    Args:
        clip: The input clip to be scaled.
        width: The target width for scaling. If None, the width of the input clip will be used.
        height: The target height for scaling. If None, the height of the input clip will be used.
        shift: A tuple representing the shift values for the x and y axes.
        **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.

    Returns:
        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

    debug(f"{self}: Preprocess kwargs: {preprocess_kwargs}")
    debug(f"{self}: Postprocess kwargs: {postprocess_kwargs}")
    debug(f"{self}: Inference kwargs: {inference_kwargs}")

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

    if 0 not in {clip.width, clip.height}:
        scaled = self.inference(wclip, **inference_kwargs)
    else:
        debug(f"{self}: Variable resolution clip detected!")

        if not isinstance(self.backend, (Backend.TRT, Backend.TRT_RTX)):
            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)

def supersample(
    self, clip: VideoNodeT, rfactor: float = 2.0, shift: tuple[TopShift, LeftShift] = (0, 0), **kwargs: Any
) -> VideoNodeT:
    """
    Supersample a clip by a given scaling factor.

    Keyword arguments passed during initialization are automatically injected here,
    unless explicitly overridden by the arguments provided at call time.
    Only arguments that match named parameters in this method are injected.

    Args:
        clip: The source clip.
        rfactor: Scaling factor for supersampling.
        shift: Subpixel shift (top, left) applied during scaling.
        **kwargs: Additional arguments forwarded to the scale function.

    Raises:
        CustomValueError: If resulting resolution is non-positive.

    Returns:
        The supersampled clip.
    """
    assert check_variable_resolution(clip, self.supersample)

    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)  # type: ignore[return-value]

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 = {
        "tiles": self.tiles,
        "tilesize": self.tilesize,
        "width": clip.width,
        "height": clip.height,
        "multiple": self.multiple,
        "overlap_w": self.overlap_w,
        "overlap_h": self.overlap_h,
    } | kwargs

    return calc_tilesize(**kwargs)

@classmethod
def ensure_obj(
    cls,
    scaler: str | type[Self] | Self | None = None,
    /,
    func_except: FuncExcept | None = None,
) -> Self:
    """
    Ensure that the input is a scaler instance, resolving it if necessary.

    Args:
        scaler: Scaler identifier (string, class, or instance).
        func_except: Function returned for custom error handling.

    Returns:
        Scaler instance.
    """
    return _base_ensure_obj(cls, scaler, func_except)

@classmethod
def from_param(
    cls,
    scaler: str | type[Self] | Self | None = None,
    /,
    func_except: FuncExcept | None = None,
) -> type[Self]:
    """
    Resolve and return a scaler type from a given input (string, type, or instance).

    Args:
        scaler: Scaler identifier (string, class, or instance).
        func_except: Function returned for custom error handling.

    Returns:
        Resolved scaler type.
    """
    return _base_from_param(cls, scaler, cls._err_class, func_except)

def get_scale_args(
    self,
    clip: vs.VideoNode,
    shift: tuple[TopShift, LeftShift] = (0, 0),
    width: int | None = None,
    height: int | None = None,
    **kwargs: Any,
) -> dict[str, Any]:
    """
    Generate the keyword arguments used for scaling.

    Args:
        clip: The source clip.
        shift: Subpixel shift (top, left).
        width: Target width.
        height: Target height.
        **kwargs: Extra parameters to merge.

    Returns:
        Final dictionary of keyword arguments for the scale function.
    """
    return {"width": width, "height": height, "src_top": shift[0], "src_left": shift[1]} | self.kwargs | kwargs

@classproperty
@classmethod
def implemented_funcs(cls) -> frozenset[str]:
    """
    Returns a set of function names that are implemented in the current class and the parent classes.

    These functions determine which keyword arguments will be extracted from the __init__ method.

    Returns:
        Frozen set of function names.
    """
    return frozenset(func for klass in cls.mro() for func in getattr(klass, "_implemented_funcs", ()))

def inference(self, clip: ConstantFormatVideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    from vsmlrt import flexible_inference

    tilesize, overlaps = self.calc_tilesize(clip)

    debug(f"{self}: Passing clip to inference: {clip.format!r}")
    debug(f"{self}: Passing model: {self.model}")
    debug(f"{self}: Passing tiles size: {tilesize}")
    debug(f"{self}: Passing overlaps: {overlaps}")
    debug(f"{self}: Passing extra kwargs: {kwargs}")

    u, v = flexible_inference(clip, self.model, overlaps, tilesize, self.backend, **kwargs)

    debug(f"{self}: Inferenced clip: {u.format!r}")
    debug(f"{self}: Inferenced clip: {v.format!r}")

    return core.std.ShufflePlanes([clip, u, v], [0, 0, 0], vs.YUV, clip)

@BaseScalerMeta.cachedproperty
def kernel_radius(self) -> int:
    """
    Return the effective kernel radius for the scaler.

    Raises:
        CustomNotImplementedError: If no kernel radius is defined.

    Returns:
        Kernel radius.
    """
    ...

def postprocess_clip(self, clip: vs.VideoNode, input_clip: vs.VideoNode, **kwargs: Any) -> ConstantFormatVideoNode:
    clip = norm_expr(clip, "x 0.5 -", [1, 2], func=self.__class__)
    return super().postprocess_clip(clip, input_clip, **kwargs)

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

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

        format = clip.format.replace(
            subsampling_h=0,
            subsampling_w=0,
            sample_type=vs.FLOAT,
            bits_per_sample=self._pick_precision(16, 32),
        )
        dither_type = DitherType.ORDERED if DitherType.should_dither(clip.format, format) else DitherType.NONE

        debug(f"{self}.pre: Before pp; Clip format is {clip.format!r}")

        clip = limiter(
            chroma_scaler.resample(clip, **dict[str, Any](format=format, dither_type=dither_type) | kwargs),
            func=self.__class__,
        )

        debug(f"{self}.pre: Before pp; Clip format is {clip.format!r}")

        return norm_expr(clip, "x 0.5 +", [1, 2], func=self.__class__)

    return norm_expr(super().preprocess_clip(clip, **kwargs), "x 0.5 +", [1, 2], func=self.__class__)

@BaseScalerMeta.cachedproperty
def pretty_string(self) -> str:
    """
    Cached property returning a user-friendly string representation.

    Returns:
        Pretty-printed string with arguments.
    """
    return self._pretty_string()

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.

    Args:
        clip: The input clip to be scaled.
        width: The target width for scaling. If None, the width of the input clip will be used.
        height: The target height for scaling. If None, the height of the input clip will be used.
        shift: A tuple representing the shift values for the x and y axes.
        **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.

    Returns:
        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

    debug(f"{self}: Preprocess kwargs: {preprocess_kwargs}")
    debug(f"{self}: Postprocess kwargs: {postprocess_kwargs}")
    debug(f"{self}: Inference kwargs: {inference_kwargs}")

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

    if 0 not in {clip.width, clip.height}:
        scaled = self.inference(wclip, **inference_kwargs)
    else:
        debug(f"{self}: Variable resolution clip detected!")

        if not isinstance(self.backend, (Backend.TRT, Backend.TRT_RTX)):
            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)