audio-demo/webui/app.py

"""
Pipeline studio: Demucs (strip) -> AuK (edit voice, with preview/undo) ->
MOSS-TTS (generate speech in the edited voice). Plus a Characters tab for
browsing the-service-station's cast and creating new characters/emotion
candidates without SSH + manual YAML editing.

Orchestrates across the three existing isolated venvs (demucs/, auk/, tts/)
via subprocess, rather than merging their conflicting dependencies
(MOSS-TTS needs transformers 5.0.0, AuK needs 4.57.6) into one environment.

Run with: webui/.venv/bin/python webui/app.py

2026-09-15 refresh, after building the-service-station's cast surfaced real
gaps in this UI:
- Speed editing was a silent no-op (AuK needs --gen_seconds or it just
  preserves the input's duration - see vendor/AuK/docs/COOKBOOK.md 3.2).
- AuK's Emotion task has a fixed 8-word vocabulary; anything else needs the
  Timbre task's free-form template instead.
- A chained AuK edit can leak fragments of its own instruction text into
  the spoken content. This went undetected through a full listen-and-
  approve cycle on a real character voice because nothing had ever
  transcribed the *reference* audio, only generated lines. Every AuK edit
  and MOSS generation here now gets transcribed automatically so that's
  visible immediately instead of discovered later by ear.
"""

import glob
import json
import os
import re
import subprocess
import tempfile
import time
from datetime import date
from pathlib import Path

import gradio as gr
import yaml

REPO_ROOT = Path(__file__).resolve().parent.parent
DEMUCS_DIR = REPO_ROOT / "demucs"
AUK_DIR = REPO_ROOT / "auk"
TTS_DIR = REPO_ROOT / "tts"

DEMUCS_PYTHON = DEMUCS_DIR / ".venv" / "bin" / "python3"
AUK_INFER = AUK_DIR / ".venv" / "bin" / "auk-infer"
AUK_CKPT = AUK_DIR / "ckpts" / "AuK" / "auk_base.safetensors"
AUK_FLASH_CKPT = AUK_DIR / "ckpts" / "AuK-Flash" / "auk_flash.safetensors"
AUK_QWEN_PATH = AUK_DIR / "ckpts" / "Qwen2.5-Omni-3B"
TTS_PYTHON = TTS_DIR / ".venv" / "bin" / "python3"
TTS_GENERATE = TTS_DIR / "generate.py"
TTS_GENERATE_SFX = TTS_DIR / "generate_sfx.py"

MODEL_BASE = "Base (full quality)"
MODEL_FLASH = "Flash (fast draft, 4-step)"

# The-service-station checkout on this same box - where character cards
# live. Override with the env var if running this UI somewhere else.
SERVICE_STATION_DIR = Path(os.environ.get("SERVICE_STATION_REPO", "/root/projects/the-service-station"))
SS_VOICES_DIR = SERVICE_STATION_DIR / "voices"

WORKDIR = Path(tempfile.mkdtemp(prefix="audio-demo-webui-"))

# AuK's Emotion task has exactly this vocabulary (vendor/AuK/docs/COOKBOOK.md
# 4.1) - anything else needs the Timbre task's free-form template instead
# (see auk_emotion_instruction below). Kept as the dropdown's options since
# these are the only words guaranteed to hit the dedicated task.
EMOTIONS = ["happy", "angry", "sad", "fearful", "surprised", "disgusted", "calm", "excited"]
PITCH_OPTIONS = ["-3", "-2", "-1", "+1", "+2", "+3"]
SPEED_OPTIONS = ["0.5", "0.75", "1.25", "1.5", "2.0"]
VOLUME_OPTIONS = ["+5", "+10", "+15", "-5", "-10", "-15"]
LANGUAGES = [
    "English", "Spanish", "French", "German", "Portuguese", "Italian",
    "Chinese", "Japanese", "Korean", "Russian",
]

COMFYUI_BASE = "http://localhost:8188"


def auk_emotion_instruction(word: str) -> str:
    """Same logic as the-service-station/workflows/generate_episode_lines.py's
    auk_emotion_instruction() - duplicated rather than shared because this
    repo and the-service-station use separate venvs/dependency sets, but
    kept in sync deliberately. AuK's Emotion task only understands the 8
    words in EMOTIONS; anything else goes through Timbre's free-form
    description instead, which has no fixed vocabulary."""
    if word in EMOTIONS:
        return f"Change the emotion to {word}."
    return f'Keep the spoken content unchanged and change the timbre to: "the same voice, sounding {word}".'


def _comfyui_queue_empty() -> bool:
    try:
        import urllib.error
        import urllib.request
        with urllib.request.urlopen(f"{COMFYUI_BASE}/queue", timeout=3) as resp:
            data = json.loads(resp.read())
            return not data.get("queue_running") and not data.get("queue_pending")
    except Exception:
        return True  # ComfyUI not reachable -> nothing to worry about


def _free_comfyui_memory() -> None:
    import urllib.error
    import urllib.request
    if not _comfyui_queue_empty():
        return
    try:
        req = urllib.request.Request(
            f"{COMFYUI_BASE}/free",
            data=json.dumps({"unload_models": True, "free_memory": True}).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        urllib.request.urlopen(req, timeout=5)
    except (urllib.error.URLError, OSError):
        pass


def _run(cmd: list[str], cwd: Path, env: dict | None = None) -> None:
    result = subprocess.run(
        cmd, cwd=cwd, env=env, capture_output=True, text=True, timeout=1800
    )
    if result.returncode != 0:
        tail = "\n".join(result.stderr.strip().splitlines()[-25:])
        raise gr.Error(f"Command failed ({cmd[0]}):\n{tail}")


def _ensure_wav(path: str) -> str:
    """AuK's loader (librosa/soundfile) can't read compressed containers like
    .m4a/.mp3 - only demucs's loader handles those. Normalize once, here, at
    the point a reference is set, so everything downstream is always wav."""
    if Path(path).suffix.lower() == ".wav":
        return path
    out_path = WORKDIR / f"converted_{int(time.time() * 1000)}.wav"
    result = subprocess.run(
        ["ffmpeg", "-nostdin", "-y", "-i", path, "-ar", "44100", "-ac", "1", str(out_path)],
        capture_output=True, text=True, timeout=120,
    )
    if result.returncode != 0:
        tail = "\n".join(result.stderr.strip().splitlines()[-15:])
        raise gr.Error(f"Could not convert '{path}' to wav:\n{tail}")
    return str(out_path)


def _duration(path: str) -> float:
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", path],
        capture_output=True, text=True,
    )
    return float(r.stdout.strip())


def _ld_library_path_for_venv(venv_dir: Path) -> str:
    matches = list(venv_dir.glob(".venv/lib/python*/site-packages"))
    if not matches:
        return ""
    site_packages = matches[0]
    nvidia_dir = site_packages / "nvidia"
    if not nvidia_dir.is_dir():
        return ""
    lib_dirs = sorted(p for p in nvidia_dir.glob("*/lib") if p.is_dir())
    return ":".join(str(p) for p in lib_dirs)


TTS_LD_LIBRARY_PATH = _ld_library_path_for_venv(TTS_DIR)


# ---------------------------------------------------------------------------
# Whisper content verification - added 2026-09-15 after a chained AuK edit
# leaked fragments of its own instruction text into a character's spoken
# reference, undetected through a full listen-and-approve round because
# nobody had transcribed the reference itself. Loaded once, lazily, since
# most sessions of this UI won't need it immediately.
# ---------------------------------------------------------------------------
_whisper_model = None


def _whisper_transcribe(path: str) -> str:
    global _whisper_model
    try:
        if _whisper_model is None:
            from faster_whisper import WhisperModel
            # "small" mis-heard casual contractions as formal phrasing on
            # short/quiet clips (e.g. "it's gonna" -> "it is going to be") -
            # "medium" is noticeably more faithful to actual speech patterns
            # at the cost of a few more seconds per transcription on CPU,
            # worth it since this check exists specifically to catch content
            # problems - a transcript that itself distorts casual speech
            # defeats that purpose.
            _whisper_model = WhisperModel("medium", device="cpu", compute_type="int8")
        segments, _ = _whisper_model.transcribe(path, language="en", beam_size=5)
        return " ".join(s.text.strip() for s in segments)
    except Exception as e:
        return f"(transcription failed: {e})"


# ---------------------------------------------------------------------------
# Stage 1: input / demucs stripping
# ---------------------------------------------------------------------------


def strip_vocals(input_audio: str):
    if not input_audio:
        raise gr.Error("Upload or record some audio first.")
    _free_comfyui_memory()
    out_dir = WORKDIR / f"demucs_{int(time.time())}"
    out_dir.mkdir(parents=True, exist_ok=True)
    _run(
        [
            str(DEMUCS_PYTHON), "-m", "demucs",
            "--two-stems=vocals", "-n", "htdemucs",
            "-o", str(out_dir), input_audio,
        ],
        cwd=DEMUCS_DIR,
    )
    stem_name = Path(input_audio).stem
    vocals = out_dir / "htdemucs" / stem_name / "vocals.wav"
    no_vocals = out_dir / "htdemucs" / stem_name / "no_vocals.wav"
    return str(vocals), str(no_vocals)


# ---------------------------------------------------------------------------
# Stage 2: AuK voice editing, with history for undo/reset
#
# history_state is a list of {"path": str, "instruction": str|None} dicts,
# one entry per actual AuK call (not per button click - a click can queue
# several edits, e.g. enhance+emotion, and each becomes its own history
# entry so "Save as character" can derive a sensibly-scoped
# generation_recipe and Undo is granular).
# ---------------------------------------------------------------------------


def use_as_reference(audio_path: str, history: list):
    if not audio_path:
        raise gr.Error("No audio to use yet.")
    audio_path = _ensure_wav(audio_path)
    history = [{"path": audio_path, "instruction": None}]
    transcript = _whisper_transcribe(audio_path)
    return audio_path, history, _history_label(history), transcript


def _history_label(history: list) -> str:
    n = len(history) - 1
    if n <= 0:
        return "No edits applied yet." if history else "No reference set yet."
    steps = "\n".join(f"{i}. {h['instruction']}" for i, h in enumerate(history[1:], 1))
    return f"{n} edit(s) applied on top of the original reference:\n{steps}"


ENHANCE_INSTRUCTIONS = {
    "none": None,
    "General (denoise + dereverb)": (
        "Preserve all speakers, remove noise and reverberation, and output "
        "clean speech of the same length."
    ),
    "Telephone/call recording": (
        "Repair the telephone effect and restore natural, clear speech."
    ),
}


def _auk_model_args(model_choice: str) -> list[str]:
    """AuK-Flash is the same architecture distilled to 4 fixed diffusion
    steps with CFG=0 (vendor README) - same --qwen_path, different --ckpt
    plus --nfe/--cfg. VRAM footprint is essentially identical to Base
    (measured in the README), so no extra offload concerns switching."""
    if model_choice == MODEL_FLASH:
        if not AUK_FLASH_CKPT.exists():
            raise gr.Error(
                f"AuK-Flash isn't downloaded yet ({AUK_FLASH_CKPT} not found) - "
                f"run `hf download tencent/AuK-Flash --local-dir ./ckpts/AuK-Flash` "
                f"from the auk/ directory first, or switch back to Base."
            )
        return ["--ckpt", str(AUK_FLASH_CKPT), "--qwen_path", str(AUK_QWEN_PATH), "--nfe", "4", "--cfg", "0"]
    return ["--ckpt", str(AUK_CKPT), "--qwen_path", str(AUK_QWEN_PATH)]


def apply_edit(
    current_ref: str,
    history: list,
    model_choice: str,
    enhance: str,
    timbre_description: str,
    emotion: str,
    other_emotion: str,
    pitch: str,
    speed: str,
    deaccent: bool,
    volume: str,
    whisper_direction: str,
    nonverbal_instruction: str,
    nonverbal_gen_seconds: float,
    custom_instruction: str,
):
    if not current_ref:
        raise gr.Error("Set a voice reference first (upload, strip with Demucs, or use a previous edit).")

    instructions = []
    # Enhancement goes first when picked - clean the recording up before any
    # character/accent editing, not after (garbage in, garbage timbre-edited).
    if enhance and ENHANCE_INSTRUCTIONS.get(enhance):
        instructions.append(ENHANCE_INSTRUCTIONS[enhance])
    if timbre_description and timbre_description.strip():
        # This is AuK's Timbre task - the one actually used to build every
        # character's base voice in this project (accent, age, register,
        # personality all described in one free-form sentence). It has no
        # fixed vocabulary, unlike Emotion below.
        instructions.append(
            f'Keep the spoken content unchanged and change the timbre to: "{timbre_description.strip()}".'
        )
    if emotion and emotion != "none":
        instructions.append(auk_emotion_instruction(emotion))
    if other_emotion and other_emotion.strip():
        # Free-text emotion word, not one of AuK's 8 - always goes through
        # the Timbre task's free-form template (see auk_emotion_instruction).
        instructions.append(auk_emotion_instruction(other_emotion.strip()))
    if pitch and pitch != "none":
        direction = "Raise" if pitch.startswith("+") else "Lower"
        instructions.append(f"{direction} the pitch by {pitch.lstrip('+-')} semitone(s).")
    if speed and speed != "none":
        speed_gen_seconds = _duration(current_ref) / float(speed)
        instructions.append((f"Adjust the speech speed to {speed}x.", speed_gen_seconds))
    if deaccent:
        instructions.append("Remove the regional accent while preserving the speaker's voice and content.")
    if volume and volume != "none":
        direction = "Increase" if volume.startswith("+") else "Decrease"
        instructions.append(f"{direction} the volume by {volume.lstrip('+-')} dB.")
    if whisper_direction and whisper_direction != "none":
        if whisper_direction == "To whisper":
            instructions.append("Convert this speech into a soft whisper while preserving the speaker and content.")
        else:
            instructions.append("Convert this whispered speech into a normal speaking voice while preserving the speaker and content.")
    if nonverbal_instruction and nonverbal_instruction.strip():
        # Add/remove changes duration unpredictably (vendor cookbook 4.4 sets
        # --gen_seconds for both directions, not just "add") - no way to
        # auto-estimate this one like Speed, the user has to know roughly
        # what they want. Defaults to the current reference's own duration
        # if left blank.
        target_seconds = float(nonverbal_gen_seconds) if nonverbal_gen_seconds else _duration(current_ref)
        instructions.append((nonverbal_instruction.strip(), target_seconds))
    if custom_instruction and custom_instruction.strip():
        instructions.append(custom_instruction.strip())

    if not instructions:
        raise gr.Error("Pick at least one edit.")

    _free_comfyui_memory()
    ref = current_ref
    out_dir = WORKDIR / f"auk_{int(time.time())}"
    out_dir.mkdir(parents=True, exist_ok=True)
    last_transcript = ""
    for i, instruction in enumerate(instructions):
        gen_seconds = None
        if isinstance(instruction, tuple):
            instruction, gen_seconds = instruction
        out_path = out_dir / f"edit_{i}.wav"
        cmd = (
            [str(AUK_INFER)] + _auk_model_args(model_choice) +
            ["--audio", ref, "--instruction", instruction, "--output", str(out_path), "--cpu_offload"]
        )
        if gen_seconds is not None:
            cmd += ["--gen_seconds", str(round(gen_seconds, 2))]
        _run(cmd, cwd=AUK_DIR)
        ref = str(out_path)
        last_transcript = _whisper_transcribe(ref)
        history = history + [{"path": ref, "instruction": instruction}]

    return ref, history, _history_label(history), last_transcript


def undo_edit(history: list):
    if len(history) <= 1:
        raise gr.Error("Nothing to undo.")
    history = history[:-1]
    transcript = _whisper_transcribe(history[-1]["path"])
    return history[-1]["path"], history, _history_label(history), transcript


def reset_to_original(history: list):
    if not history:
        raise gr.Error("No reference set yet.")
    history = [history[0]]
    transcript = _whisper_transcribe(history[0]["path"])
    return history[0]["path"], history, _history_label(history), transcript


# ---------------------------------------------------------------------------
# Stage 3: MOSS-TTS generation
# ---------------------------------------------------------------------------


def generate_speech(reference: str, text: str, language: str):
    if not reference:
        raise gr.Error("Set a voice reference first (upload, strip, and/or edit it above).")
    if not text or not text.strip():
        raise gr.Error("Enter some text to synthesize.")

    _free_comfyui_memory()
    out_path = WORKDIR / f"tts_{int(time.time())}.wav"
    env = None
    if TTS_LD_LIBRARY_PATH:
        env = os.environ.copy()
        env["LD_LIBRARY_PATH"] = TTS_LD_LIBRARY_PATH

    _run(
        [
            str(TTS_PYTHON), str(TTS_GENERATE),
            "--text", text,
            "--reference", reference,
            "--language", language,
            "--output", str(out_path),
        ],
        cwd=TTS_DIR,
        env=env,
    )
    transcript = _whisper_transcribe(str(out_path))
    return str(out_path), transcript


def generate_sound_effect(description: str):
    """MOSS-SoundEffect - text-only ambience/SFX generation, no reference
    audio or character voice involved (different generation primitive from
    the AuK/MOSS-TTS voice-cloning flow above, so it's a standalone section
    rather than wired into that pipeline's history/reference state).
    No whisper verification here - there's no speech content to check."""
    if not description or not description.strip():
        raise gr.Error("Enter a description of the sound effect/ambience.")

    _free_comfyui_memory()
    out_path = WORKDIR / f"sfx_{int(time.time())}.wav"
    env = None
    if TTS_LD_LIBRARY_PATH:
        env = os.environ.copy()
        env["LD_LIBRARY_PATH"] = TTS_LD_LIBRARY_PATH

    _run(
        [str(TTS_PYTHON), str(TTS_GENERATE_SFX), "--text", description.strip(), "--output", str(out_path)],
        cwd=TTS_DIR,
        env=env,
    )
    return str(out_path)


# ---------------------------------------------------------------------------
# Characters (the-service-station cast)
# ---------------------------------------------------------------------------


def _slugify(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")


def read_character_cards() -> dict:
    """name -> parsed card dict (frontmatter + prose), read-only. Never
    round-trips these through yaml.dump - only used for display and for
    resolving a base reference_audio, never for rewriting an existing file
    (that would silently strip hand-authored inline comments)."""
    cards = {}
    if not SS_VOICES_DIR.is_dir():
        return cards
    for card_path in sorted(SS_VOICES_DIR.glob("*/CHARACTER.md")):
        try:
            text = card_path.read_text(encoding="utf-8")
            parts = text.split("---", 2)
            if len(parts) < 3:
                continue
            frontmatter = yaml.safe_load(parts[1]) or {}
            frontmatter["_prose"] = parts[2].strip()
            frontmatter["_dir"] = card_path.parent
            cards[frontmatter.get("name", card_path.parent.name)] = frontmatter
        except Exception as e:
            cards[card_path.parent.name] = {"_error": str(e), "_dir": card_path.parent}
    return cards


def character_names() -> list[str]:
    return sorted(read_character_cards().keys())


def load_character(name: str):
    if not name:
        return "Pick a character.", None
    cards = read_character_cards()
    card = cards.get(name)
    if not card:
        return f"No card found for '{name}'.", None
    if "_error" in card:
        return f"Could not parse {name}'s card: {card['_error']}", None

    ref_audio = card.get("reference_audio")
    ref_path = str(SERVICE_STATION_DIR / ref_audio) if ref_audio else None

    lib = card.get("emotion_library") or {}
    lib_lines = [f"- **{emo}**: {entry.get('status', '?')}" for emo, entry in lib.items()] or ["(none approved yet)"]

    cache_dir = card["_dir"] / "emotion-cache"
    cached = sorted(p.stem for p in cache_dir.glob("*.wav")) if cache_dir.is_dir() else []
    cache_lines = [f"- {e} (cached, not yet approved)" for e in cached] or ["(none pending)"]

    info = (
        f"**{name}** — {card.get('role', '?')}\n\n"
        f"Status: `{card.get('status', '?')}` | Engine: `{card.get('engine', '?')}` | "
        f"Default emotion (documentation only): `{card.get('default_emotion', '?')}`\n\n"
        f"### Approved emotion_library\n" + "\n".join(lib_lines) + "\n\n"
        f"### Pending (emotion-cache, not approved)\n" + "\n".join(cache_lines) + "\n\n"
        f"---\n{card.get('_prose', '')}"
    )
    return info, ref_path


def list_character_voices(name: str) -> dict[str, str]:
    """label -> reference audio path, for every voice available for a
    character: the base voice, every approved emotion_library entry, and
    every pending (not yet approved) emotion-cache candidate. Used both by
    the clone-as-character tool below and could back a similar picker
    anywhere else a character's variants need listing."""
    cards = read_character_cards()
    card = cards.get(name)
    if not card or "_error" in card:
        return {}
    voices = {}
    ref_audio = card.get("reference_audio")
    if ref_audio:
        voices["Base (natural voice)"] = str(SERVICE_STATION_DIR / ref_audio)
    for emo, entry in (card.get("emotion_library") or {}).items():
        ref = entry.get("reference_audio")
        if ref:
            voices[f"{emo} (approved)"] = str(SERVICE_STATION_DIR / ref)
    cache_dir = card["_dir"] / "emotion-cache"
    if cache_dir.is_dir():
        for p in sorted(cache_dir.glob("*.wav")):
            voices[f"{p.stem} (cached, not approved)"] = str(p)
    return voices


def refresh_voice_choices(name: str):
    voices = list_character_voices(name)
    labels = list(voices.keys())
    return gr.update(choices=labels, value=labels[0] if labels else None)


def clone_as_character(character_name: str, voice_label: str, text: str, language: str):
    if not character_name:
        raise gr.Error("Pick a character first.")
    if not voice_label:
        raise gr.Error("Pick a voice (base or one of its emotions).")
    if not text or not text.strip():
        raise gr.Error("Enter some text for them to say.")

    voices = list_character_voices(character_name)
    reference = voices.get(voice_label)
    if not reference:
        raise gr.Error(f"Could not resolve '{voice_label}' for {character_name} - try Refresh.")

    _free_comfyui_memory()
    env = None
    if TTS_LD_LIBRARY_PATH:
        env = os.environ.copy()
        env["LD_LIBRARY_PATH"] = TTS_LD_LIBRARY_PATH
    out_path = WORKDIR / f"clone_{int(time.time())}.wav"
    _run(
        [str(TTS_PYTHON), str(TTS_GENERATE), "--text", text.strip(),
         "--reference", reference, "--language", language, "--output", str(out_path)],
        cwd=TTS_DIR, env=env,
    )
    transcript = _whisper_transcribe(str(out_path))
    return str(out_path), transcript


def create_character(
    name: str, role: str, prose: str, language: str,
    current_ref: str, history: list,
):
    if not name or not name.strip():
        raise gr.Error("Give the character a name.")
    if not current_ref or not history:
        raise gr.Error("Set up a reference and edit it above first - there's nothing to save yet.")

    slug = _slugify(name)
    char_dir = SS_VOICES_DIR / slug / "webui-created"
    if (SS_VOICES_DIR / slug / "CHARACTER.md").exists():
        raise gr.Error(f"voices/{slug}/CHARACTER.md already exists - this UI won't overwrite an existing character.")
    char_dir.mkdir(parents=True, exist_ok=True)

    source_dest = char_dir / f"{slug}_source.wav"
    ref_dest = char_dir / f"{slug}_ref_v1.wav"
    subprocess.run(["cp", history[0]["path"], str(source_dest)], check=True)
    subprocess.run(["cp", current_ref, str(ref_dest)], check=True)

    recipe = [
        {"tool": "auk", "instruction": h["instruction"]}
        for h in history[1:] if h["instruction"]
    ]

    frontmatter = {
        "name": name.strip(),
        "role": role.strip() or "Standing guest",
        "engine": "moss-tts",
        "model": "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5",
        "reference_audio": str(ref_dest.relative_to(SERVICE_STATION_DIR)),
        "source_clip": str(source_dest.relative_to(SERVICE_STATION_DIR)),
        "source_clip_note": "Created via the webui pipeline studio - see generation_recipe for the exact edit chain.",
        "generation_recipe": recipe,
        "seed": None,
        "emotion_library": {},
        "language": language,
        "status": "draft",
        "last_updated": date.today().isoformat(),
        "provenance": "webui",
        "appearances": [],
    }
    frontmatter_yaml = yaml.dump(frontmatter, sort_keys=False, allow_unicode=True)
    body = f"---\n{frontmatter_yaml}---\n\n# {name.strip()}\n\n{prose.strip()}\n"
    (SS_VOICES_DIR / slug / "CHARACTER.md").write_text(body, encoding="utf-8")

    return (
        f"Created voices/{slug}/CHARACTER.md with status: draft. "
        f"Not committed - review, listen, and commit it yourself (or ask an agent to) "
        f"once you're happy, same as every other voice in this project."
    )


def generate_emotion_candidate(character_name: str, emotion_word: str, model_choice: str):
    if not character_name:
        raise gr.Error("Pick a character first.")
    if not emotion_word or not emotion_word.strip():
        raise gr.Error("Enter an emotion word.")
    emotion_word = emotion_word.strip()

    cards = read_character_cards()
    card = cards.get(character_name)
    if not card or "_error" in card:
        raise gr.Error(f"Could not load {character_name}'s card.")

    base_reference = str(SERVICE_STATION_DIR / card["reference_audio"])
    cache_dir = card["_dir"] / "emotion-cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = cache_dir / f"{_slugify(emotion_word)}.wav"

    _free_comfyui_memory()
    instruction = auk_emotion_instruction(emotion_word)
    cmd = (
        [str(AUK_INFER)] + _auk_model_args(model_choice) +
        ["--audio", base_reference, "--instruction", instruction,
         "--output", str(cache_path), "--cpu_offload"]
    )
    _run(cmd, cwd=AUK_DIR)

    ref_transcript = _whisper_transcribe(str(cache_path))

    # Preview clone so the candidate can actually be judged, not just the
    # edited reference in isolation.
    env = None
    if TTS_LD_LIBRARY_PATH:
        env = os.environ.copy()
        env["LD_LIBRARY_PATH"] = TTS_LD_LIBRARY_PATH
    preview_text = f"Testing the {emotion_word} emotion for {character_name}."
    preview_path = WORKDIR / f"emotion_preview_{int(time.time())}.wav"
    _run(
        [str(TTS_PYTHON), str(TTS_GENERATE), "--text", preview_text,
         "--reference", str(cache_path), "--language", "English", "--output", str(preview_path)],
        cwd=TTS_DIR, env=env,
    )
    preview_transcript = _whisper_transcribe(str(preview_path))

    model_note = (
        " Generated with **Flash** - if this direction is right, regenerate the same "
        "word with **Base** before promoting it, since Flash trades fidelity for speed "
        "and this is a candidate you might lock in permanently."
        if model_choice == MODEL_FLASH else ""
    )
    status = (
        f"Saved to `{cache_path.relative_to(SERVICE_STATION_DIR)}` (cached, **not approved**).{model_note}\n\n"
        f"Reference transcript: {ref_transcript}\n\n"
        f"Preview clone transcript: {preview_transcript}\n\n"
        f"If this sounds right, promote it manually per `.claude/skills/produce-episode/"
        f"SKILL.md`'s \"Adding an emotion to a character's library\" - this UI won't do "
        f"that step for you."
    )
    return str(cache_path), str(preview_path), status


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="Audio-demo pipeline studio") as demo:
    with gr.Tabs():
        with gr.Tab("Pipeline Studio"):
            gr.Markdown(
                "# Pipeline studio\n"
                "Strip a voice with **Demucs**, edit it with **AuK**, then generate new "
                "speech in that voice with **MOSS-TTS**. Each stage shells out to its own "
                "isolated venv, so this box's GPU is the real constraint — heavy edits "
                "(long clips) can OOM AuK; keep reference clips short (under ~25s)."
            )

            reference_state = gr.State(None)
            history_state = gr.State([])

            with gr.Row():
                with gr.Column():
                    gr.Markdown("## 1. Input")
                    input_audio = gr.Audio(label="Upload or record", type="filepath")
                    strip_btn = gr.Button("Strip vocals with Demucs")
                    with gr.Row():
                        vocals_out = gr.Audio(label="Vocals (isolated)", type="filepath")
                        no_vocals_out = gr.Audio(label="Background (no vocals)", type="filepath")
                    with gr.Row():
                        use_vocals_btn = gr.Button("Use vocals as reference")
                        use_no_vocals_btn = gr.Button("Use background as reference")
                        use_original_btn = gr.Button("Use uploaded audio as-is")

                with gr.Column():
                    gr.Markdown("## 2. Voice editor (AuK)")
                    current_ref_audio = gr.Audio(label="Current reference", type="filepath", interactive=False)
                    ref_transcript_box = gr.Textbox(
                        label="Whisper transcript of current reference (editable - fix any mis-hearing before using it)",
                        info="Whisper can mishear casual speech as more formal phrasing on short/quiet clips (e.g. \"it's gonna\" -> \"it is going to be\") - correct it here before sending it anywhere.",
                    )
                    use_ref_transcript_btn = gr.Button("Use this transcript as 'Text to speak' →")
                    history_label = gr.Markdown("No reference set yet.")
                    model_dd = gr.Radio(
                        [MODEL_BASE, MODEL_FLASH], value=MODEL_BASE, label="AuK model",
                        info="Draft with Flash (fast, 4-step, lower fidelity), confirm with Base before committing a voice or emotion you're actually going to keep.",
                    )
                    enhance_dd = gr.Dropdown(
                        list(ENHANCE_INSTRUCTIONS.keys()), value="none",
                        label="Enhance (denoise/dereverb/restore) - runs first, before character edits",
                    )
                    timbre_tb = gr.Textbox(
                        label="Timbre / voice description (accent, age, register, personality)",
                        placeholder=(
                            'e.g. "British male, mid-30s, dry and understated, slight Essex '
                            'accent, deadpan delivery" - this is what actually built every '
                            "character's base voice in this project. Free-form, no fixed "
                            "vocabulary (unlike Emotion below)."
                        ),
                        lines=2,
                    )
                    emotion_dd = gr.Dropdown(
                        ["none"] + EMOTIONS, value="none",
                        label="Emotion (AuK's fixed vocabulary - anything else, use the field below)",
                    )
                    other_emotion_tb = gr.Textbox(
                        label="Other emotion word (not in AuK's list above)",
                        placeholder='e.g. "bemused" - a single word for a momentary reaction, wrapped '
                                     "as a Timbre instruction automatically (different from the "
                                     "voice-description field above, which is for the character's "
                                     "lasting accent/age/personality)",
                    )
                    pitch_dd = gr.Dropdown(["none"] + PITCH_OPTIONS, value="none", label="Pitch (semitones)")
                    speed_dd = gr.Dropdown(["none"] + SPEED_OPTIONS, value="none", label="Speed")
                    gr.Markdown(
                        "*Pitch editing has been unreliable in testing — even \"1 semitone\" "
                        "has overshot to ~4+ semitones on some voices. Prefer a free-form "
                        "custom instruction like \"a little lower in register\" over the Pitch "
                        "dropdown for character work.*\n\n"
                        "*Speed now sets AuK's required `--gen_seconds` automatically. On "
                        "references longer than ~15s this can stall (a 27s target hung this "
                        "box for 20+ minutes swapping to disk with 0% GPU use, while a 12s "
                        "source completed in seconds) - trim your reference shorter first if "
                        "using Speed.*"
                    )
                    with gr.Accordion("More AuK tasks (de-accent, volume, whisper, nonverbal)", open=False):
                        deaccent_cb = gr.Checkbox(label="De-accent (remove regional accent, keep voice)")
                        volume_dd = gr.Dropdown(
                            ["none"] + VOLUME_OPTIONS, value="none", label="Volume (dB)",
                        )
                        whisper_dd = gr.Dropdown(
                            ["none", "To whisper", "From whisper"], value="none",
                            label="Whisper conversion",
                        )
                        nonverbal_tb = gr.Textbox(
                            label="Nonverbal editing (advanced)",
                            placeholder='e.g. "Add a cough before \'We tested\'" or "Remove all laughs from the audio."',
                        )
                        nonverbal_seconds_nb = gr.Number(
                            label="Target duration for the nonverbal edit above (seconds)",
                            value=None,
                            info="AuK needs an explicit target length for add/remove nonverbal edits - it can't be auto-computed like Speed. Defaults to the current reference's own duration; increase it a little for \"add\", decrease it a little for \"remove\".",
                        )
                    custom_instruction = gr.Textbox(
                        label="Custom instruction (advanced, optional)",
                        placeholder='e.g. "Lower the pitch very slightly."',
                    )
                    apply_btn = gr.Button("Apply edit(s)", variant="primary")
                    with gr.Row():
                        undo_btn = gr.Button("Undo last edit")
                        reset_btn = gr.Button("Reset to original")

                with gr.Column():
                    gr.Markdown("## 3. Generate (MOSS-TTS)")
                    gen_text = gr.Textbox(label="Text to speak", lines=6)
                    gen_language = gr.Dropdown(LANGUAGES, value="English", label="Language")
                    generate_btn = gr.Button("Generate speech", variant="primary")
                    generated_audio = gr.Audio(label="Generated speech", type="filepath")
                    gen_transcript_box = gr.Textbox(label="Whisper transcript of generated speech")

                    gr.Markdown("## 4. Save as a the-service-station character")
                    new_char_name = gr.Textbox(label="Name")
                    new_char_role = gr.Textbox(label="Role", placeholder="e.g. Standing guest - weather")
                    new_char_prose = gr.Textbox(label="Character brief (prose)", lines=4)
                    save_char_btn = gr.Button("Save as new character (status: draft)")
                    save_char_status = gr.Markdown()

            with gr.Row():
                with gr.Column():
                    gr.Markdown(
                        "## 5. Sound effects (MOSS-SoundEffect)\n"
                        "Text-only ambience/SFX generation - no reference audio or character "
                        "voice involved, a separate model from the voice pipeline above."
                    )
                    sfx_description = gr.Textbox(
                        label="Description",
                        placeholder=(
                            'e.g. "A door creaking open slowly, then closing with a soft thud." '
                            'or "A busy outdoor market with distant chatter and footsteps."'
                        ),
                        lines=2,
                    )
                    sfx_btn = gr.Button("Generate sound effect", variant="primary")
                    sfx_audio = gr.Audio(label="Generated sound effect", type="filepath")

            strip_btn.click(strip_vocals, inputs=[input_audio], outputs=[vocals_out, no_vocals_out])
            use_vocals_btn.click(
                use_as_reference, inputs=[vocals_out, history_state],
                outputs=[current_ref_audio, history_state, history_label, ref_transcript_box],
            )
            use_no_vocals_btn.click(
                use_as_reference, inputs=[no_vocals_out, history_state],
                outputs=[current_ref_audio, history_state, history_label, ref_transcript_box],
            )
            use_original_btn.click(
                use_as_reference, inputs=[input_audio, history_state],
                outputs=[current_ref_audio, history_state, history_label, ref_transcript_box],
            )
            apply_btn.click(
                apply_edit,
                inputs=[current_ref_audio, history_state, model_dd, enhance_dd, timbre_tb, emotion_dd,
                        other_emotion_tb, pitch_dd, speed_dd, deaccent_cb, volume_dd, whisper_dd,
                        nonverbal_tb, nonverbal_seconds_nb, custom_instruction],
                outputs=[current_ref_audio, history_state, history_label, ref_transcript_box],
            )
            undo_btn.click(
                undo_edit, inputs=[history_state],
                outputs=[current_ref_audio, history_state, history_label, ref_transcript_box],
            )
            reset_btn.click(
                reset_to_original, inputs=[history_state],
                outputs=[current_ref_audio, history_state, history_label, ref_transcript_box],
            )
            generate_btn.click(
                generate_speech, inputs=[current_ref_audio, gen_text, gen_language],
                outputs=[generated_audio, gen_transcript_box],
            )
            use_ref_transcript_btn.click(lambda t: t, inputs=[ref_transcript_box], outputs=[gen_text])
            save_char_btn.click(
                create_character,
                inputs=[new_char_name, new_char_role, new_char_prose, gen_language, current_ref_audio, history_state],
                outputs=[save_char_status],
            )
            sfx_btn.click(generate_sound_effect, inputs=[sfx_description], outputs=[sfx_audio])

        with gr.Tab("Characters"):
            gr.Markdown(
                f"# The Service Station cast\n"
                f"Reading from `{SERVICE_STATION_DIR}`. Approving a character or an emotion "
                f"is a manual step (editing the card) - this tab creates drafts and "
                f"candidates for you to listen to and promote yourself, it never marks "
                f"anything `approved` on its own."
            )
            with gr.Row():
                char_dd = gr.Dropdown(label="Character", choices=character_names())
                refresh_btn = gr.Button("Refresh")
            char_info = gr.Markdown()
            char_audio = gr.Audio(label="Reference audio", type="filepath", interactive=False)

            gr.Markdown(
                "## Clone text as this character\n"
                "Say anything in the selected character's voice - the base voice, an "
                "approved emotion, or a pending (not yet approved) candidate from the "
                "cache. Useful both for \"how would they say this?\" and for testing an "
                "emotion candidate with real dialogue instead of the fixed test line below."
            )
            clone_voice_dd = gr.Dropdown(label="Voice", choices=list(list_character_voices(None).keys()))
            clone_text_tb = gr.Textbox(label="Text to speak", lines=3, placeholder="Type anything.")
            clone_language_dd = gr.Dropdown(LANGUAGES, value="English", label="Language")
            clone_btn = gr.Button("Clone", variant="primary")
            clone_audio = gr.Audio(label="Cloned speech", type="filepath", interactive=False)
            clone_transcript_box = gr.Textbox(label="Whisper transcript")

            gr.Markdown("## Generate a new emotion candidate")
            with gr.Row():
                emotion_char_dd = gr.Dropdown(label="Character", choices=character_names())
                emotion_word_tb = gr.Textbox(label="Emotion word", placeholder="e.g. surprised, bemused, wistful")
            emotion_model_dd = gr.Radio(
                [MODEL_BASE, MODEL_FLASH], value=MODEL_FLASH, label="AuK model",
                info="Flash defaults here since this is exploring candidates - regenerate the same word with Base before promoting one.",
            )
            gen_emotion_btn = gr.Button("Generate candidate", variant="primary")
            emotion_ref_audio = gr.Audio(label="Edited reference (candidate)", type="filepath", interactive=False)
            emotion_preview_audio = gr.Audio(label="Preview clone", type="filepath", interactive=False)
            emotion_status = gr.Markdown()

            refresh_btn.click(
                lambda: (gr.update(choices=character_names()), gr.update(choices=character_names())),
                outputs=[char_dd, emotion_char_dd],
            )
            char_dd.change(load_character, inputs=[char_dd], outputs=[char_info, char_audio])
            char_dd.change(refresh_voice_choices, inputs=[char_dd], outputs=[clone_voice_dd])
            clone_btn.click(
                clone_as_character,
                inputs=[char_dd, clone_voice_dd, clone_text_tb, clone_language_dd],
                outputs=[clone_audio, clone_transcript_box],
            )
            gen_emotion_btn.click(
                generate_emotion_candidate,
                inputs=[emotion_char_dd, emotion_word_tb, emotion_model_dd],
                outputs=[emotion_ref_audio, emotion_preview_audio, emotion_status],
            ).then(
                # So a freshly-generated candidate is immediately pickable in
                # "Clone text as this character" without reselecting the
                # character to force a refresh. Keyed off char_dd (the
                # browser above), not emotion_char_dd - the clone tool
                # always uses whichever character char_dd is showing, and
                # those two dropdowns can legitimately point at different
                # characters at once.
                refresh_voice_choices, inputs=[char_dd], outputs=[clone_voice_dd],
            )
            # Dropdown choices are otherwise fixed at server-start time, so a
            # character created earlier in the same running session (or by
            # someone else) wouldn't show up without this - refresh
            # automatically on every page load, in addition to the manual
            # Refresh button above for mid-session use without reloading.
            demo.load(
                lambda: (gr.update(choices=character_names()), gr.update(choices=character_names())),
                outputs=[char_dd, emotion_char_dd],
            )

if __name__ == "__main__":
    # Gradio refuses to serve a file back to the browser unless it's under
    # the app's own directory, the system temp dir, or explicitly allowed -
    # every character's reference audio lives in the-service-station repo,
    # a different path entirely, so every single Characters-tab selection
    # errored without this (found via a real user report + the server log,
    # not by inspection - the error only surfaces server-side, not in the
    # browser's generic error toast).
    demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=[str(SERVICE_STATION_DIR)])