helpers ¶

@cache
def get_gpu(device_id: int = 0) -> Device | None:
    """
    Return the GPU available for the requested device id.
    """
    from pyopencl import device_type, get_platforms

    if not (platforms := get_platforms()):
        logger.debug("No OpenCL platforms found.")
        return None

    try:
        # platform also resolves any software/cpu targets
        # but it doesn't seem like any one platform can have multiple devices
        gpus = [device for device in [pl.get_devices()[0] for pl in platforms] if device.type == device_type.GPU]
    except Exception as e:
        logger.debug(e, exc_info=e)
        return None

    try:
        return gpus[device_id]
    except KeyError:
        raise CustomValueError(f"No GPU found for device_id {device_id}!") from None

def is_gpu_available(device_id: int = 0) -> bool:
    """
    Check if a GPU is available for the requested device id.
    """
    return get_gpu(device_id) is not None

def pre_ss(
    clip: vs.VideoNode,
    function: VSFunctionNoArgs | None = None,
    rfactor: float = 2.0,
    sp: type[MixedScalerProcess[Any, Any]] | MixedScalerProcess[Any, Any] = ComplexSuperSamplerProcess,
    supersampler: ScalerLike | Callable[[vs.VideoNode, int, int], vs.VideoNode] | None = None,
    downscaler: ScalerLike | Callable[[vs.VideoNode, int, int], vs.VideoNode] | None = None,
    mod: int = 4,
    pscale: float = 1.0,
    planes: Planes = None,
    func: FuncExcept | None = None,
) -> vs.VideoNode:
    """
    Supersamples the input clip, applies a given function to the higher-resolution version, and then downscales it back to the original resolution.

    This function generalizes the behavior of
    [SuperSamplerProcess][vsaa.deinterlacers.SuperSamplerProcess] and
    [ComplexSuperSamplerProcess][vsscale.various.ComplexSuperSamplerProcess].

    - Examples:
        ```py
        out = pre_ss(clip, lambda clip: cool_function(clip, ...), planes=0)
        ```

        - Passing NNEDI3 as a supersampler:
        ```py
        from vsaa import SuperSamplerProcess

        out = pre_ss(clip, lambda clip: cool_function(clip, ...), SuperSamplerProcess)
        ```
        - This works too:
        ```py
        from vsaa import SuperSamplerProcess

        out = pre_ss(clip, sp=SuperSamplerProcess(function=lambda clip: cool_function(clip, ...)))
        ```

        - Specifying `supersampler` and `downscaler`:
        ```py
        from vskernels import Point

        out = pre_ss(clip, lambda clip: cool_function(clip, ...), supersampler=Point, downscaler=Point, planes=0)
        ```

    Args:
        clip: Source clip.
        function: A function to apply on the supersampled clip.
        rfactor: Scaling factor for supersampling. Defaults to 2.
        sp: A [MixedScalerProcess][vskernels.MixedScalerProcess] instance or class.
            Default is [ComplexSuperSamplerProcess[Lanczos]][vsscale.ComplexSuperSamplerProcess].
            It upscales with [Lanczos][vskernels.Lanczos] and downscales with [Point][vskernels.Point].
        supersampler: Scaler used to upscale the input clip if `sp` is not specified.
        downscaler: Downscaler used for undoing the upscaling done by the supersampler if `sp` is not specified.
        mod: Ensures the supersampled resolution is a multiple of this value. Defaults to 4.
        pscale: Scaling correction strength.
        planes: Which planes to process.
        func: An optional function to use for error handling.

    Returns:
        A clip with the given function applied at higher resolution, then downscaled back.
    """  # noqa: E501
    func_util = FunctionUtil(clip, func or pre_ss, planes)

    args = (
        func_util.work_clip,
        mod_x(func_util.work_clip.width * rfactor, mod),
        mod_x(func_util.work_clip.height * rfactor, mod),
    )

    if isinstance(sp, MixedScalerProcess):
        if pscale != 1.0:
            raise CustomValueError("pscale is not supported when using a MixedScalerProcess!", func_util.func)
        return func_util.return_clip(sp.scale(*args))

    if supersampler and downscaler and function:
        ss = _get_scaled(supersampler, *args, func_util.func)

        processed = function(ss)

        down = _get_scaled(downscaler, processed, clip.width, clip.height, func_util.func)
        down = pscale_blend(
            func_util.work_clip,
            down,
            lambda: _get_scaled(downscaler, ss, clip.width, clip.height, func_util.func),
            pscale,
            func=func_util.func,
        )

        return func_util.return_clip(down)

    if function:
        return pre_ss(clip, sp=sp(function=function), rfactor=rfactor, mod=mod, pscale=pscale, planes=planes, func=func)

    raise CustomTypeError

def pscale_blend(
    clip: vs.VideoNode,
    processed: vs.VideoNode,
    unprocessed: vs.VideoNode | Callable[[], vs.VideoNode],
    pscale: float = 0.0,
    planes: Planes = None,
    func: FuncExcept | None = None,
) -> vs.VideoNode:
    """
    Apply pscale-based error correction/scaling expression.

    This function calculates a blend of the processed clip, the original clip,
    and an unprocessed, re-scaled and downscaled clip to correct downscaling artifacts/ringing.

    The formula evaluated is:
        `y + (1 - pscale) * (x - z)`

    Where:
        - `x` is `clip` (original input clip)
        - `y` is `processed` (processed and downscaled clip)
        - `z` is `unprocessed` (unprocessed and downscaled clip)

    - Behavior of `pscale`:
        - `1.0`: No correction is applied (`y` is returned).
        - `0.5`: The scaling change is blended back at half strength.
        - `0.0`: The scaling change is blended back at full strength.
        - `> 1.0`: Subtracts scaling changes (acts as a low-pass/blur).
        - `< 0.0`: Amplifies scaling changes (acts as a low-quality sharpen).

    Args:
        clip: Original input clip (x).
        processed: Processed and resampled/downscaled clip (y).
        unprocessed: Unprocessed resampled/downscaled clip or a callable that returns it (z).
        pscale: Scale factor/correction strength.
        planes: Planes to process.
        func: Function returned for custom error handling.
    Returns:
        The blended/corrected clip.
    """
    func = func or pscale_blend

    if pscale == 1.0:
        return processed

    if callable(unprocessed):
        unprocessed = unprocessed()

    return norm_expr([clip, processed, unprocessed], f"x z - {1.0 - pscale} * y +", planes, func=func)

def scale_var_clip(
    clip: vs.VideoNode,
    scaler: Scaler | Callable[[tuple[int, int]], Scaler],
    width: int | Callable[[tuple[int, int]], int] | None,
    height: int | Callable[[tuple[int, int]], int],
    shift: tuple[float, float] | Callable[[tuple[int, int]], tuple[float, float]] = (0, 0),
    debug: bool = False,
) -> vs.VideoNode:
    """
    Scale a variable clip to constant or variable resolution.

    Args:
        clip: Source clip.
        scaler: A scaler instance or a callable that returns a scaler instance.
        width: A width integer or a callable that returns the width. If None, it will be calculated from the height and
            the aspect ratio of the clip.
        height: A height integer or a callable that returns the height.
        shift: Optional top shift, left shift tuple or a callable that returns the shifts. Defaults to no shift.
        debug: If True, the `var_width` and `var_height` props will be added to the clip.

    Returns:
        Scaled clip.
    """
    cached_clips = dict[str, vs.VideoNode]()

    no_accepts_var = list[Scaler]()

    def _eval_scale(f: vs.VideoFrame, n: int) -> vs.VideoNode:
        key = f"{f.width}_{f.height}"

        if key not in cached_clips:
            res = (f.width, f.height)

            norm_scaler = scaler(res) if callable(scaler) else scaler
            norm_shift = shift(res) if callable(shift) else shift
            norm_height = height(res) if callable(height) else height

            if width is None:
                norm_width = get_w(norm_height, f.width / f.height)
            else:
                norm_width = width(res) if callable(width) else width

            part_scaler = partial(norm_scaler.scale, width=norm_width, height=norm_height, shift=norm_shift)

            scaled = clip
            if (scaled.width, scaled.height) != (norm_width, norm_height):
                if norm_scaler not in no_accepts_var:
                    try:
                        scaled = part_scaler(clip)
                    except Exception:  # noqa: BLE001
                        no_accepts_var.append(norm_scaler)

                if norm_scaler in no_accepts_var:
                    const_clip = clip.resize.Point(f.width, f.height)

                    scaled = part_scaler(const_clip)

            if debug:
                scaled = scaled.std.SetFrameProps(var_width=f.width, var_height=f.height)

            cached_clips[key] = scaled

        return cached_clips[key]

    out_clip = clip if callable(width) or callable(height) else clip.std.BlankClip(width, height)

    return out_clip.std.FrameEval(_eval_scale, clip, clip)