src/config.py

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path


def _bool_env(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _path_env(name: str, default: str) -> Path:
    return Path(os.getenv(name, default)).expanduser().resolve()


@dataclass(frozen=True)
class Settings:
    # Model identity. This Space intentionally targets a single model family.
    gguf_repo_id: str = os.getenv("GGUF_REPO_ID", "unsloth/diffusiongemma-26B-A4B-it-GGUF")
    gguf_filename: str = os.getenv("GGUF_FILENAME", "diffusiongemma-26B-A4B-it-Q4_K_M.gguf")
    model_name: str = os.getenv("MODEL_NAME", "unsloth/diffusiongemma-26B-A4B-it-GGUF:Q4_K_M")

    # Runtime directories. Prefer /data when a persistent Storage Bucket is attached.
    data_dir: Path = _path_env("DATA_DIR", "/data" if Path("/data").exists() else "/tmp/diffusiongemma")
    model_cache_dir: Path = _path_env("MODEL_CACHE_DIR", "/data/models" if Path("/data").exists() else "/tmp/diffusiongemma/models")
    bin_dir: Path = _path_env("BIN_DIR", "/data/bin" if Path("/data").exists() else "/tmp/diffusiongemma/bin")
    llama_src_dir: Path = _path_env("LLAMA_SRC_DIR", "/data/llama.cpp" if Path("/data").exists() else "/tmp/diffusiongemma/llama.cpp")

    llama_diffusion_bin: Path = _path_env(
        "LLAMA_DIFFUSION_BIN",
        "/data/bin/llama-diffusion-cli" if Path("/data").exists() else "/tmp/diffusiongemma/bin/llama-diffusion-cli",
    )
    llama_diffusion_bin_url: str | None = os.getenv("LLAMA_DIFFUSION_BIN_URL")

    # Build/download behavior.
    build_llama_diffusion: bool = _bool_env("BUILD_LLAMA_DIFFUSION", True)
    download_model_on_startup: bool = _bool_env("DOWNLOAD_MODEL_ON_STARTUP", False)
    prepare_runtime_on_startup: bool = _bool_env("PREPARE_RUNTIME_ON_STARTUP", False)
    llama_build_cuda: bool = _bool_env("LLAMA_BUILD_CUDA", True)
    llama_cmake_extra_args: str = os.getenv("LLAMA_CMAKE_EXTRA_ARGS", "")

    # ZeroGPU and API limits.
    zero_gpu_size: str = os.getenv("ZEROGPU_SIZE", "large")
    api_time_limit_seconds: int = int(os.getenv("API_TIME_LIMIT_SECONDS", "1200"))
    cli_timeout_seconds: int = int(os.getenv("CLI_TIMEOUT_SECONDS", "1100"))

    default_max_tokens: int = int(os.getenv("DEFAULT_MAX_TOKENS", "512"))
    max_max_tokens: int = int(os.getenv("MAX_MAX_TOKENS", "2048"))
    min_gpu_duration_seconds: int = int(os.getenv("MIN_GPU_DURATION_SECONDS", "120"))
    max_gpu_duration_seconds: int = int(os.getenv("MAX_GPU_DURATION_SECONDS", "180"))
    gpu_base_seconds: int = int(os.getenv("GPU_BASE_SECONDS", "45"))
    seconds_per_block_step: float = float(os.getenv("SECONDS_PER_BLOCK_STEP", "0.5"))

    # Diffusion sampler knobs. Keep these conservative; expose fewer knobs than a normal GGUF runner.
    n_gpu_layers: int = int(os.getenv("N_GPU_LAYERS", "99"))
    diffusion_max_steps: int = int(os.getenv("DIFFUSION_MAX_STEPS", "48"))
    diffusion_visual_default: bool = _bool_env("DIFFUSION_VISUAL_DEFAULT", False)
    diffusion_kv_cache: str = os.getenv("DIFFUSION_KV_CACHE", "auto")

    # Prompt behavior.
    default_system_prompt: str = os.getenv(
        "DEFAULT_SYSTEM_PROMPT",
        "You are DiffusionGemma, a concise and useful assistant. Respond only with the final answer. Do not include hidden reasoning, scratchpad notes, channel tags, or performance metrics.",
    )
    thinking_enabled_default: bool = _bool_env("THINKING_ENABLED_DEFAULT", False)

    # Output cleanup.
    strip_prompt_echo: bool = _bool_env("STRIP_PROMPT_ECHO", True)
    prompt_mode: str = os.getenv("PROMPT_MODE", "auto")  # auto | arg | stdin


settings = Settings()

if settings.zero_gpu_size not in {"large", "xlarge"}:
    raise RuntimeError("ZEROGPU_SIZE must be 'large' or 'xlarge'.")