shared/gpu.py

"""ZeroGPU helpers.

`gpu` é um decorator: no HF Spaces (ZeroGPU) ele aloca uma GPU por chamada;
localmente (sem o pacote `spaces`) vira um no-op transparente. Assim o mesmo
código roda nas 3060 locais e no Space.
"""
import functools

try:
    import spaces  # disponível só nos Spaces ZeroGPU
    _HAS_SPACES = True
except Exception:
    _HAS_SPACES = False


def gpu(duration: int = 60):
    """Aloca GPU por `duration`s na chamada (ZeroGPU). No-op local."""
    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"