zerogpu.py

"""Helpers for running Tabras workloads on Hugging Face ZeroGPU.

The real `spaces.GPU` decorator is available inside Hugging Face Spaces.  Local
and test environments may not have the `spaces` package installed, so this
module falls back to an identity decorator while keeping call sites identical.
"""

from collections.abc import Callable
from typing import Any, TypeVar, overload

F = TypeVar("F", bound=Callable[..., Any])

try:
    import spaces
except ImportError:  # pragma: no cover - exercised when HF Spaces SDK is absent
    spaces = None


@overload
def gpu(func: F, /) -> F: ...


@overload
def gpu(func: None = None, /, **kwargs: Any) -> Callable[[F], F]: ...


# Decorate a callable with `spaces.GPU` when available, otherwise leave it as-is.
def gpu(func: F | None = None, /, **kwargs: Any) -> F | Callable[[F], F]:
    if spaces is None:
        if func is None:
            return lambda wrapped: wrapped
        return func
    decorator = spaces.GPU(**kwargs)
    if func is None:
        return decorator
    return decorator(func)