src/utils/zero_gpu.py

"""Optional Hugging Face ZeroGPU decorator helpers."""

from __future__ import annotations

from collections.abc import Callable
from typing import TypeVar


F = TypeVar("F", bound=Callable)


def zero_gpu(duration: int = 180) -> Callable[[F], F]:
    """Return a ZeroGPU decorator when available, otherwise a no-op decorator."""
    try:
        import spaces  # type: ignore[import-not-found]
    except Exception:
        return _identity_decorator

    return spaces.GPU(duration=duration)


def _identity_decorator(func: F) -> F:
    return func