File size: 863 Bytes
4df4c39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
"""Runtime selection without importing CUDA or Spaces in local mode."""

import os


def runtime_mode():
    mode = os.getenv("IDOL_RUNTIME", "auto").lower()
    if mode not in {"auto", "local", "zerogpu"}:
        raise ValueError("IDOL_RUNTIME must be auto, local, or zerogpu")
    hosted = os.getenv("SPACES_ZERO_GPU", "").lower() not in {"", "0", "false"}
    if mode == "local" and hosted:
        raise ValueError("Local mode cannot run on ZeroGPU hardware")
    return ("zerogpu" if hosted else "local") if mode == "auto" else mode


MODE = runtime_mode()


def gpu_task(function=None, *, duration=None):
    if function is None:
        return lambda fn: gpu_task(fn, duration=duration)
    if MODE == "zerogpu":
        import spaces

        return spaces.GPU(duration=duration or int(os.getenv("IDOL_GPU_SECONDS", "90")))(function)
    return function