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
|