core/gpu.py

"""ZeroGPU helpers.

`gpu` is a decorator: on HF Spaces (ZeroGPU) it allocates a GPU per call;
locally (without the `spaces` package) it becomes a transparent no-op, so the
same code runs on the local GPUs and on the Space.
"""
import functools

try:
    import spaces  # available only on ZeroGPU Spaces
    _HAS_SPACES = True
except Exception:
    _HAS_SPACES = False


def gpu(duration: int = 60):
    """Allocate a GPU for `duration`s on the call (ZeroGPU). No-op locally."""
    def decorate(fn):
        if _HAS_SPACES:
            return spaces.GPU(duration=duration)(fn)

        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            return fn(*args, **kwargs)

        return wrapper

    return decorate


def device() -> str:
    import torch
    return "cuda" if torch.cuda.is_available() else "cpu"