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

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.cached_property
def kernel_radius(self) -> int:
    """
    Return the effective kernel radius for the scaler.

    Raises:
        CustomNotImplementedError: If no kernel radius is defined.

    Returns:
        Kernel radius.
    """
    ...

@deprecated('The "multi" method is deprecated. Use "supersample" instead.', category=DeprecationWarning)
def multi(
    self, clip: VideoNodeT, multi: float = 2.0, shift: tuple[TopShift, LeftShift] = (0, 0), **kwargs: Any
) -> VideoNodeT:
    """
    Deprecated alias for `supersample`.

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

    Returns:
        The supersampled clip.
    """
    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.
    """
    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.cached_property
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.

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

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.cached_property
def kernel_radius(self) -> int:
    """
    Return the effective kernel radius for the scaler.

    Raises:
        CustomNotImplementedError: If no kernel radius is defined.

    Returns:
        Kernel radius.
    """
    ...

@deprecated('The "multi" method is deprecated. Use "supersample" instead.', category=DeprecationWarning)
def multi(
    self, clip: VideoNodeT, multi: float = 2.0, shift: tuple[TopShift, LeftShift] = (0, 0), **kwargs: Any
) -> VideoNodeT:
    """
    Deprecated alias for `supersample`.

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

    Returns:
        The supersampled clip.
    """
    return self.supersample(clip, multi, shift, **kwargs)

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 super().preprocess_clip(norm_expr(clip, "x 0.5 +", [1, 2], func=self.__class__), **kwargs)

@BaseScalerMeta.cached_property
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)