runtime_env.py

import os
from pathlib import Path


ROOT_DIR = Path(__file__).resolve().parent


def configure_runtime_env() -> None:
    os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1")
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
    os.environ.setdefault("ATTN_BACKEND", "flash_attn_3")
    os.environ.setdefault(
        "FLEX_GEMM_AUTOTUNE_CACHE_PATH", str(ROOT_DIR / "autotune_cache.json")
    )
    os.environ.setdefault("FLEX_GEMM_AUTOTUNER_VERBOSE", "1")


configure_runtime_env()