File size: 551 Bytes
e20e3d9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
"""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