visualnovel/config.py

"""All knobs in one place. Everything is overridable via environment variables so you
can flip between mock / real / different models without touching code.

Defaults are the "beefy" config (Qwen3-14B + SDXL-Turbo + Whisper-small), and MOCK is ON
by default so a fresh checkout runs end-to-end with zero model downloads.
"""

from __future__ import annotations

import os
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()  # loads .env from CWD or any parent; no-op if file absent

# --------------------------------------------------------------------------- #
#  Hugging Face authentication (required for gated models)
# --------------------------------------------------------------------------- #
HF_TOKEN: str | None = os.getenv("HF_TOKEN")

# True when launched with --mode debug (set by app.py before this import).
DEBUG: bool = os.getenv("VN_DEBUG", "0") == "1"

# --------------------------------------------------------------------------- #
#  Paths
# --------------------------------------------------------------------------- #
ROOT = Path(__file__).resolve().parent.parent
MODELS_DIR = Path(os.getenv("VN_MODELS_DIR", ROOT / "models"))
MEMORY_DIR = Path(os.getenv("VN_MEMORY_DIR", ROOT / "memory"))
RUNS_DIR = Path(os.getenv("VN_RUNS_DIR", ROOT / "runs"))
TEMPLATES_DIR = ROOT / "templates"
CACHE_DIR = Path(os.getenv("VN_CACHE_DIR", ROOT / ".cache" / "images"))
for _d in (MEMORY_DIR, RUNS_DIR, CACHE_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# --------------------------------------------------------------------------- #
#  Feature flags
# --------------------------------------------------------------------------- #
# MOCK on  -> orchestrator/painter/stt use deterministic fakes; no models needed.
# MOCK off -> real backends are instantiated (install the relevant uv extras first).
USE_MOCK = os.getenv("VN_MOCK", "1") == "1"

# VN_REAL_STT=1 -> use real Whisper even in mock mode (test STT without full model stack).
REAL_STT = os.getenv("VN_REAL_STT", "0") == "1"

# Detect HuggingFace Space runtime — SPACE_ID is always injected by the HF platform.
_ON_SPACE = bool(os.getenv("SPACE_ID"))

# LLM runtime backend when not mocking.
# Auto-selects transformers on HF Spaces (ZeroGPU); llamacpp everywhere else.
# Override with VN_LLM_BACKEND=modal for the Modal cloud path.
LLM_BACKEND = os.getenv("VN_LLM_BACKEND", "transformers" if _ON_SPACE else "llamacpp")

# Image generation backend when not mocking.
# "lightning" = SDXL-Lightning (4-step, GPU-local)  |  "modal" = Modal cloud  |  "local" = SDXL-Turbo
IMAGE_BACKEND = os.getenv(
    "VN_IMAGE_BACKEND",
    "modal" if LLM_BACKEND == "modal" else "lightning",
)

# Set to a .jsonl path to record a full per-turn trace (Open-Trace bonus).
TRACE_PATH = os.getenv("VN_TRACE")  # None disables tracing

# Force the gr.Blocks MVP UI instead of the custom frontend (Phase 0/1 de-risking).
MVP_UI = os.getenv("GRADIO_MVP_UI", "0") == "1"

# --------------------------------------------------------------------------- #
#  Model registry  (the "beefy" config — comfortably < 32B total)
#    Qwen3-14B (~14B) + SDXL-Turbo (~3.5B) + Whisper-small (~0.25B) ≈ ~18B
# --------------------------------------------------------------------------- #
LLM_REPO = os.getenv("VN_LLM_REPO", "Qwen/Qwen3-14B")  # transformers id
LLM_GGUF_REPO = os.getenv("VN_LLM_GGUF_REPO", "Qwen/Qwen3-14B-GGUF")  # llama.cpp source
LLM_GGUF_FILE = os.getenv("VN_LLM_GGUF_FILE", "Qwen3-14B-Q4_K_M.gguf")

# Defaults are backend-aware so the app works without an explicit .env override.
_lightning_model = "stabilityai/stable-diffusion-xl-base-1.0"
_lightning_lora = "ByteDance/SDXL-Lightning::sdxl_lightning_4step_lora.safetensors"
IMAGE_MODEL = os.getenv(
    "VN_IMAGE_MODEL",
    _lightning_model if IMAGE_BACKEND == "lightning" else "stabilityai/sdxl-turbo",
)
# Path or HF repo of YOUR fine-tuned anime-style LoRA (the Well-Tuned bonus). Empty = none.
IMAGE_LORA = os.getenv(
    "VN_IMAGE_LORA",
    _lightning_lora if IMAGE_BACKEND == "lightning" else "",
)
# fp16-safe VAE: the stock SDXL VAE has force_upcast, so diffusers upcasts it to fp32 at
# EVERY render (slow + per-render FutureWarning). The fix VAE decodes natively in fp16.
# Empty = keep the pipeline's stock VAE.
IMAGE_VAE = os.getenv("VN_IMAGE_VAE", "madebyollin/sdxl-vae-fp16-fix")

WHISPER_SIZE = os.getenv(
    "VN_WHISPER_SIZE", "large-v3-turbo"
)  # tiny|base|small|medium|large-v3|large-v3-turbo

TTS_BACKEND = os.getenv("VN_TTS_BACKEND", "mock")  # mock | kokoro

# --------------------------------------------------------------------------- #
#  Generation / budget params
# --------------------------------------------------------------------------- #
CONTEXT_TOKEN_BUDGET = int(os.getenv("VN_CTX_BUDGET", "3500"))  # the "Thousand Token Wood"
RECENT_TURNS_K = int(os.getenv("VN_RECENT_K", "6"))
# Compaction triggers when recent_turns exceeds 2*RECENT_TURNS_K and keeps RECENT_TURNS_K
# verbatim (see memory.should_compact / orchestrator.compact_memory).
IMAGE_SIZE = int(os.getenv("VN_IMAGE_SIZE", "768"))
_default_steps = "4" if IMAGE_BACKEND == "lightning" else "2"
IMAGE_STEPS = int(os.getenv("VN_IMAGE_STEPS", _default_steps))
# Backdrops use a non-zero guidance so the negative prompt (no people) is actually applied.
# Sprites keep 0.0 (distillation works best there). 1.5 is safe for both Turbo and Lightning.
BACKDROP_GUIDANCE = float(os.getenv("VN_BACKDROP_GUIDANCE", "1.5"))
MAX_PRESENT = int(os.getenv("VN_MAX_PRESENT", "3"))  # spirits on stage at once

# --------------------------------------------------------------------------- #
#  Fixed art-style base. The chosen theme only adds *setting* flavor on top of this,
#  so every playthrough keeps one coherent look. (See README: style is locked, theme varies.)
# --------------------------------------------------------------------------- #
STYLE_BASE = os.getenv(
    "VN_STYLE_BASE",
    "Japanese anime illustration, soft cel shading, clean linework, painterly background, "
    "expressive eyes, gentle film grain, storybook lighting",
)
NEGATIVE_PROMPT = os.getenv(
    "VN_NEGATIVE",
    "text, watermark, signature, lowres, blurry, deformed hands, extra limbs, "
    "multiple characters, jpeg artifacts",
)

# --------------------------------------------------------------------------- #
#  Pre-set setup-form themes (the player picks one from a dropdown).
#  Keep STYLE locked; `setting` is injected into prompts + the opening.
# --------------------------------------------------------------------------- #
THEMES: dict[str, dict[str, str]] = {
    "school": {
        "label": "School - slice of life",
        "setting": "a Japanese high school bathed in golden afternoon light",
        "palette": "warm, nostalgic",
    },
    "space": {
        "label": "Space - derelict starship",
        "setting": "a quiet derelict starship adrift among slow-turning stars",
        "palette": "cool teal and amber",
    },
    "cyberpunk_city": {
        "label": "Cyberpunk city",
        "setting": "a rain-slick neon megacity at night, signs humming in the drizzle",
        "palette": "magenta and cyan neon",
    },
    "seaside_town": {
        "label": "Seaside town",
        "setting": "a sleepy coastal town at first light, gulls and salt on the air",
        "palette": "soft pastel blues",
    },
}
TONES = ["romantic", "comedic", "dramatic", "bittersweet", "flirty"]


def detect_device() -> str:
    """Return the torch device family.

    NOTE: torch reports **both** NVIDIA CUDA and **AMD ROCm** as ``cuda`` — so an
    RX 7900 XTX shows up here as "cuda". Apple Silicon (Metal) shows up as "mps".
    """
    try:
        import torch  # noqa: PLC0415  (lazy — torch is an optional extra)

        if torch.backends.mps.is_available():
            return "mps"
        if torch.cuda.is_available():
            return "cuda"
    except Exception:
        pass
    return "cpu"