jspace/gpu.py

"""ZeroGPU decorator shim.

On the Hugging Face Space, GPU-bound functions must be decorated ``@spaces.GPU``
so a GPU attaches for the call and releases after (§5). Locally (Mac / CI, no
``spaces`` package, no CUDA) we want the same modules to import and run on CPU
against a tiny model. ``GPU`` here is ``spaces.GPU`` when available and an
identity decorator otherwise.
"""

from __future__ import annotations

from typing import Callable

try:  # pragma: no cover - exercised only on the Space
    import spaces

    def GPU(fn: Callable | None = None, *, duration: int = 120):
        """``@spaces.GPU`` with an optional duration budget (seconds)."""
        if fn is not None:
            return spaces.GPU(fn)
        return spaces.GPU(duration=duration)

    HAS_ZEROGPU = True
except Exception:  # noqa: BLE001 - any import failure means "no ZeroGPU here"

    def GPU(fn: Callable | None = None, *, duration: int = 120):
        """Identity decorator used off-Space so modules import on CPU."""
        if fn is not None:
            return fn

        def _wrap(f: Callable) -> Callable:
            return f

        return _wrap

    HAS_ZEROGPU = False