utils/gpu_runtime.py

"""Taichi lifecycle helpers for HF ZeroGPU (init on demand, release before render)."""

from __future__ import annotations

import os

import taichi as ti

_TAICHI_READY = False


def hf_space_mode() -> bool:
    return os.environ.get("ENDOGSIM_HF_SPACE") == "1"


def ensure_taichi_runtime() -> None:
    global _TAICHI_READY
    if _TAICHI_READY:
        return
    ti_mem = 0.05 if hf_space_mode() else 0.1
    ti.init(arch=ti.cuda, device_memory_fraction=ti_mem)
    _TAICHI_READY = True


def release_taichi_runtime() -> None:
    global _TAICHI_READY
    if not _TAICHI_READY:
        return
    if hf_space_mode():
        ti.reset()
    _TAICHI_READY = False