app.py

"""ACE-Step Inspire — creative text-to-song Space for ACE-Step 1.5."""

from __future__ import annotations

import logging
import os
import random
import sys
import time
import traceback
from typing import Optional

# ZeroGPU: import spaces BEFORE torch
try:
    import spaces

    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

for _proxy in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"):
    os.environ.pop(_proxy, None)
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

# ── Logging (stdout so Hugging Face Space logs capture everything) ───────────
LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO").upper()
logging.basicConfig(
    level=getattr(logging, LOG_LEVEL, logging.INFO),
    format="%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
    datefmt="%H:%M:%S",
    stream=sys.stdout,
    force=True,
)
log = logging.getLogger("ace-inspire")
# Keep third-party noise down unless debugging
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)

import gradio as gr
import torch
from diffusers import AceStepPipeline

from audio_export import AUDIO_FORMATS, DEFAULT_AUDIO_FORMAT, write_audio
from lyrics_gen import build_caption, generate_lyrics
from mood_engine import default_dims, resolve_mood_bundle, resolve_style_bundle
from presets import (
    ALL_STYLES,
    ALL_VOCALS,
    DEFAULT_DURATION,
    DEFAULT_GENRE,
    DEFAULT_MODEL,
    GENRES,
    INSTRUMENTS_ALL,
    MODELS,
    STRUCTURES,
    USE_CASES,
    VOCAL_LANGUAGES,
    suggest_use_case_bpm,
)
from console_ui import (
    AUTO_OPTS_HTML,
    BRAND_HTML,
    CONSOLE_CSS,
    CONSOLE_JS,
    DURATION_BAY_HTML,
    FORMAT_PADS_HTML,
    GENRE_PADS_HTML,
    INSTRUMENTAL_PAD_HTML,
    INSTRUMENTS_BROWSER_HTML,
    KNOB_HTML,
    LANGUAGE_PADS_HTML,
    MOOD_BOARD_HTML,
    PRESET_HTML,
    STRUCTURE_PADS_HTML,
    STYLE_PADS_HTML,
    USECASE_PADS_HTML,
    VOCAL_PADS_HTML,
)


def _gpu_mem_str() -> str:
    if not torch.cuda.is_available():
        return "cuda=unavailable"
    try:
        free, total = torch.cuda.mem_get_info()
        alloc = torch.cuda.memory_allocated()
        reserved = torch.cuda.memory_reserved()
        return (
            f"gpu_free={free/1e9:.2f}G/{total/1e9:.2f}G "
            f"alloc={alloc/1e9:.2f}G reserved={reserved/1e9:.2f}G"
        )
    except Exception as e:
        return f"gpu_mem_err={e}"


log.info(
    "Boot | HAS_SPACES=%s SPACE_ID=%s cuda_available=%s torch=%s",
    HAS_SPACES,
    os.environ.get("SPACE_ID"),
    torch.cuda.is_available(),
    torch.__version__,
)

# ── Pipeline cache (CPU-resident; moved to CUDA inside @spaces.GPU) ──────────
_pipes: dict[str, AceStepPipeline] = {}
_current_repo: Optional[str] = None


def _load_pipe(repo_id: str) -> AceStepPipeline:
    global _current_repo
    if repo_id in _pipes:
        log.info("Pipeline cache hit: %s", repo_id)
        return _pipes[repo_id]
    # Keep only one heavy pipeline in memory on ZeroGPU
    if _pipes:
        log.info("Clearing cached pipelines: %s", list(_pipes))
    _pipes.clear()
    t0 = time.perf_counter()
    log.info("Loading pipeline from_pretrained(%s) dtype=bfloat16 …", repo_id)
    try:
        pipe = AceStepPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
    except Exception as e:
        log.exception("from_pretrained failed for %s", repo_id)
        raise gr.Error(
            f"Failed to load model `{repo_id}`.\n"
            f"This Space only supports Diffusers AceStepPipeline checkpoints.\n\n{e}"
        ) from e
    if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_tiling"):
        pipe.vae.enable_tiling()
        log.debug("VAE tiling enabled")
    _pipes[repo_id] = pipe
    _current_repo = repo_id
    log.info("Pipeline ready: %s (%.1fs)", repo_id, time.perf_counter() - t0)
    return pipe


def _model_cfg(model_name: str) -> dict:
    return MODELS.get(model_name, MODELS[DEFAULT_MODEL])


# Preload default checkpoint on CPU during startup so ZeroGPU time isn't spent downloading.
try:
    _DEFAULT_REPO = MODELS[DEFAULT_MODEL]["repo_id"]
    log.info("Preloading default model on CPU: %s", _DEFAULT_REPO)
    _load_pipe(_DEFAULT_REPO)
    log.info("Default model ready | %s", _gpu_mem_str())
except Exception as e:
    log.exception("Default model preload skipped: %s", e)


def apply_style_defaults(genre: str, style: str):
    """BPM + mood matrix + instruments for the selected style."""
    gname = genre or DEFAULT_GENRE
    g = GENRES.get(gname) or GENRES[DEFAULT_GENRE]
    sname = style or g["styles"][0]
    # Pad genre/style can race — resolve to a genre that owns this style.
    if sname not in g["styles"]:
        for name, gg in GENRES.items():
            if sname in gg["styles"]:
                gname, g = name, gg
                break
        else:
            sname = g["styles"][0]
    dims_s, label, key, meter, instruments, bpm = resolve_style_bundle(gname, sname)
    return (
        gr.update(value=bpm),
        gr.update(choices=INSTRUMENTS_ALL, value=instruments),
        dims_s,
        label,
        key,
        meter,
    )


def apply_use_case_bpm(use_case: str):
    """Set tempo from the use-case default when one is chosen."""
    bpm = suggest_use_case_bpm(use_case or "(none)")
    if bpm is None:
        return gr.update()
    return gr.update(value=int(bpm))


def apply_genre(genre: str):
    """Fill creative defaults from genre + first style (all still user-overridable)."""
    g = GENRES[genre]
    style0 = g["styles"][0]
    dims_s, label, key, meter, instruments, bpm = resolve_style_bundle(genre, style0)
    # Keep full choice unions so pad-driven values never fail Gradio validation.
    return (
        gr.update(choices=ALL_STYLES, value=style0),
        gr.update(choices=ALL_VOCALS, value=g["vocal"][0]),
        gr.update(choices=INSTRUMENTS_ALL, value=instruments),
        gr.update(value=bpm),
        dims_s,
        label,
        key,
        meter,
    )


def surprise_me(genre: str, theme: str, language: str = "English", instructions: str = ""):
    """Randomize creative controls within the selected genre."""
    g = GENRES.get(genre, GENRES[DEFAULT_GENRE])
    style = random.choice(g["styles"])
    vocal = random.choice(g["vocal"])
    n_inst = min(3, len(g["instruments"]))
    instruments = random.sample(g["instruments"], k=n_inst)
    bpm = random.randint(*g["bpm"])
    # jitter mood dims around genre defaults
    base = default_dims(genre)
    jittered = {k: max(0, min(100, v + random.randint(-18, 18))) for k, v in base.items()}
    dims_s, label, key, meter = resolve_mood_bundle(jittered, genre, g.get("keys"))
    structure = random.choice([s for s in STRUCTURES if s != "Instrumental (no lyrics)"])
    use_case = random.choice(USE_CASES[1:])
    lyrics = generate_lyrics(
        genre,
        label,
        theme or label,
        structure,
        instrumental=False,
        language=language,
        instructions=instructions or "",
    )
    caption = build_caption(
        genre,
        style,
        label,
        instruments,
        vocal,
        bpm,
        use_case,
        instrumental=False,
        auto_extra=True,
        notes=instructions or "",
    )
    log.info("Surprise | style=%s mood=%s bpm=%s key=%s lang=%s", style, label, bpm, key, language)
    return (
        style,
        vocal,
        instruments,
        bpm,
        dims_s,
        label,
        key,
        meter,
        structure,
        use_case,
        lyrics,
        caption,
    )


def on_generate_lyrics(genre, mood, theme, structure, instrumental, seed, language, instructions=""):
    seed_i = int(seed) if seed is not None and int(seed) >= 0 else None
    log.info(
        "Lyrics gen | genre=%s mood=%s structure=%s instrumental=%s seed=%s lang=%s theme=%r notes=%r",
        genre,
        mood,
        structure,
        instrumental,
        seed_i,
        language,
        (theme or "")[:80],
        (instructions or "")[:80],
    )
    text = generate_lyrics(
        genre=genre,
        mood=mood or "emotional",
        theme=theme or mood or genre,
        structure_name=structure,
        instrumental=bool(instrumental),
        seed=seed_i,
        language=language,
        instructions=instructions or "",
    )
    log.info("Lyrics gen done | chars=%d lang=%s", len(text), language)
    return text


def peek_prompt(genre, style, mood, instruments, vocal, bpm, use_case, instrumental, lyrics, instructions=""):
    """Build current caption + show lyrics separately (lyrics are NOT inside the caption)."""
    caption = build_caption(
        genre=genre,
        style=style,
        mood=mood,
        instruments=instruments if isinstance(instruments, list) else [],
        vocal=vocal,
        bpm=int(bpm) if bpm else 120,
        use_case=use_case,
        instrumental=bool(instrumental),
        auto_extra=True,
        notes=instructions or "",
    )
    lyric_text = "[Instrumental]" if instrumental else ((lyrics or "").strip() or "(no lyrics yet)")
    view = (
        "=== CAPTION / STYLE PROMPT ===\n"
        f"{caption}\n\n"
        "=== LYRICS (separate ACE-Step input, not part of caption) ===\n"
        f"{lyric_text}"
    )
    log.info("Prompt peek | caption=%r", caption[:200])
    return view, gr.update(visible=True)


def _meter_label(timesignature: str) -> str:
    ts = str(timesignature or "4")
    return "6/8" if ts == "6" else f"{ts}/4"


def _generate_impl(
    model_name,
    lyrics,
    duration,
    bpm,
    keyscale,
    timesignature,
    language_label,
    instrumental,
    steps,
    guidance,
    shift,
    seed,
    random_seed,
    genre=None,
    style=None,
    mood=None,
    instruments=None,
    vocal=None,
    use_case=None,
    audio_format=None,
    auto_play=False,  # client-side only; kept in signature for Gradio wiring
    instructions="",
):
    t_run = time.perf_counter()
    log.info("=" * 60)
    log.info(
        "Generate start | model=%r duration=%s bpm=%s key=%s meter=%s lang=%s instrumental=%s format=%s",
        model_name,
        duration,
        bpm,
        keyscale,
        timesignature,
        language_label,
        instrumental,
        audio_format,
    )
    log.info("CUDA before | available=%s | %s", torch.cuda.is_available(), _gpu_mem_str())

    try:
        cfg = _model_cfg(model_name)
        repo_id = cfg["repo_id"]
        log.info("Resolved model | name=%r repo=%s turbo=%s defaults=%s", model_name, repo_id, cfg["turbo"], cfg)

        t0 = time.perf_counter()
        pipe = _load_pipe(repo_id)
        log.info("Moving pipeline to CUDA … | %s", _gpu_mem_str())
        pipe.to("cuda")
        log.info("Pipeline on CUDA (%.1fs) | %s", time.perf_counter() - t0, _gpu_mem_str())

        duration = int(duration)
        duration = max(10, min(duration, 3600))
        fmt = (audio_format or DEFAULT_AUDIO_FORMAT).strip()
        if fmt not in AUDIO_FORMATS:
            fmt = DEFAULT_AUDIO_FORMAT

        if instrumental:
            lyrics_text = "[Instrumental]"
        else:
            lyrics_text = (lyrics or "").strip() or "[Instrumental]"

        seed_val = -1 if seed is None else int(seed)
        use_seed = random.randint(0, 2**31 - 1) if random_seed or seed_val < 0 else seed_val
        generator = torch.Generator(device="cuda").manual_seed(use_seed)

        steps = int(steps) if steps else cfg["steps"]
        guidance = float(guidance) if guidance is not None else cfg["guidance"]
        shift = float(shift) if shift is not None else cfg["shift"]
        if cfg["turbo"]:
            guidance = 1.0

        lang = VOCAL_LANGUAGES.get(language_label, "en")

        try:
            bpm_i = int(float(bpm)) if bpm is not None and str(bpm).strip() != "" else None
        except (TypeError, ValueError):
            bpm_i = None
        if bpm_i is not None and bpm_i < 1:
            bpm_i = None

        # ACE docs: BPM soft-control range is ~30–300. Outside that, metas are OOD.
        bpm_meta = None
        bpm_note = ""
        if bpm_i is not None:
            bpm_meta = max(30, min(300, bpm_i))
            if bpm_meta != bpm_i:
                bpm_note = f" (metas clamped {bpm_i}→{bpm_meta}; ACE range 30–300)"
                log.warning("BPM %s outside ACE metas range; clamping to %s", bpm_i, bpm_meta)

        log.info("Auto-building prompt | live_bpm=%s metas_bpm=%s", bpm_i, bpm_meta)
        notes = (instructions or "").strip()
        prompt = build_caption(
            genre=genre or DEFAULT_GENRE,
            style=style or "",
            mood=mood or "",
            instruments=instruments if isinstance(instruments, list) else [],
            vocal=vocal or "",
            bpm=int(bpm_i or bpm_meta or 120),
            use_case=use_case or "(none)",
            instrumental=bool(instrumental),
            auto_extra=True,
            notes=notes,
        )
        if not prompt or not str(prompt).strip():
            raise gr.Error("Could not build a prompt from the current controls.")

        # Tempo is soft text conditioning; turbo has no CFG — stress BPM in the instruction.
        bpm_for_inst = bpm_meta if bpm_meta is not None else int(bpm_i or 120)
        instruction = (
            "Fill the audio semantic mask based on the given conditions. "
            f"Strictly match tempo: exactly {bpm_for_inst} BPM with a clear steady pulse at that speed. "
            "Do not use half-time or double-time feels that hide the target BPM:"
        )
        if notes:
            instruction = f"{instruction} Additional creative notes (follow these; do not sing them): {notes}"

        kwargs = dict(
            prompt=str(prompt).strip(),
            lyrics=lyrics_text,
            audio_duration=float(duration),
            vocal_language=lang,
            num_inference_steps=steps,
            guidance_scale=guidance,
            shift=shift,
            generator=generator,
            instruction=instruction,
            bpm=bpm_meta,
            keyscale=keyscale or None,
            timesignature=str(timesignature) if timesignature else None,
            task_type="text2music",
        )

        log.info(
            "Inference params | bpm_meta=%s bpm_req=%s key=%s meter=%s steps=%s guidance=%s shift=%s seed=%s duration=%ss lang=%s turbo=%s",
            bpm_meta,
            bpm_i,
            keyscale,
            timesignature,
            steps,
            guidance,
            shift,
            use_seed,
            duration,
            lang,
            cfg["turbo"],
        )
        log.info("Instruction: %r", instruction)
        log.info("Prompt (%d chars): %r", len(kwargs["prompt"]), kwargs["prompt"][:240])
        log.info("Lyrics (%d chars): %r", len(lyrics_text), lyrics_text[:240].replace("\n", " | "))

        if cfg["turbo"] and bpm_meta is not None:
            log.info(
                "Tempo note: Turbo is guidance-distilled (no CFG). BPM is soft metadata only — "
                "XL SFT follows tempo more reliably."
            )

        t_inf = time.perf_counter()
        log.info("pipe(...) starting …")
        try:
            output = pipe(**kwargs)
            audio = output.audios[0]
        except Exception as e:
            log.error("pipe(...) failed after %.1fs | %s", time.perf_counter() - t_inf, _gpu_mem_str())
            log.error("Traceback:\n%s", traceback.format_exc())
            raise gr.Error(f"Generation failed:\n{type(e).__name__}: {e}") from e

        log.info(
            "pipe(...) done in %.1fs | audio_type=%s shape=%s | %s",
            time.perf_counter() - t_inf,
            type(audio).__name__,
            getattr(audio, "shape", None),
            _gpu_mem_str(),
        )

        if isinstance(audio, torch.Tensor):
            audio = audio.detach().float().cpu().numpy()
        if audio.ndim == 2:
            if audio.shape[0] <= 8 and audio.shape[0] < audio.shape[1]:
                audio = audio.T

        sr = getattr(pipe, "sample_rate", 48000)
        out_path = write_audio(audio, sr, fmt)
        size_mb = os.path.getsize(out_path) / 1e6
        ext = os.path.splitext(out_path)[1].lstrip(".").upper() or fmt

        meta = (
            f"Model: {repo_id}\n"
            f"Seed: {use_seed} | Steps: {steps} | Guidance: {guidance} | Shift: {shift}\n"
            f"Duration: {duration}s | BPM: {bpm_meta}{bpm_note} | Key: {keyscale} | Meter: {_meter_label(str(timesignature))}\n"
            f"Language: {lang} | Format: {ext}\n"
            f"Wall time: {time.perf_counter() - t_run:.1f}s | File: {size_mb:.1f} MB @ {sr} Hz"
        )
        if cfg["turbo"]:
            meta += (
                "\nTempo: Turbo follows BPM softly (no CFG). "
                "For stronger tempo lock, switch to XL SFT and raise Guidance."
            )
        log.info(
            "Generate OK | %.1fs | file=%s format=%s (%.1f MB) | %s",
            time.perf_counter() - t_run,
            out_path,
            ext,
            size_mb,
            _gpu_mem_str(),
        )
        log.info("=" * 60)
        prompt_bundle = (
            "=== CAPTION / STYLE PROMPT ===\n"
            f"{prompt}\n\n"
            "=== LYRICS (separate ACE-Step input, not part of caption) ===\n"
            f"{lyrics_text}"
        )
        return (
            gr.update(value=out_path, autoplay=bool(auto_play)),
            meta,
            prompt_bundle,
            "<div class='empty-hint' style='color:#3dffb0'>Track loaded on deck.</div>",
            out_path,
        )

    except gr.Error:
        log.error("Generate aborted (Gradio Error) after %.1fs", time.perf_counter() - t_run)
        raise
    except Exception as e:
        log.error("Generate crashed after %.1fs | %s", time.perf_counter() - t_run, _gpu_mem_str())
        log.error("Traceback:\n%s", traceback.format_exc())
        raise gr.Error(f"Unexpected error:\n{type(e).__name__}: {e}") from e
    finally:
        sys.stdout.flush()
        sys.stderr.flush()


if HAS_SPACES:
    generate_music = spaces.GPU(duration=300)(_generate_impl)
else:
    generate_music = _generate_impl


# ── UI ───────────────────────────────────────────────────────────────────────

dark_theme = gr.themes.Base(
    primary_hue="amber",
    secondary_hue="slate",
    neutral_hue="zinc",
    font=[gr.themes.GoogleFont("Rajdhani"), "ui-sans-serif", "system-ui"],
    font_mono=[gr.themes.GoogleFont("Share Tech Mono"), "monospace"],
).set(
    body_background_fill="#07080b",
    body_text_color="#e8ecf4",
    block_background_fill="#141821",
    block_border_color="#2a3142",
    block_label_text_color="#6b7385",
    button_primary_background_fill="#ffb020",
    button_primary_text_color="#1a1200",
    border_color_primary="#2a3142",
    input_background_fill="#0a0c11",
)

with gr.Blocks(
    title="INSPIRE · ACE-Step",
    theme=dark_theme,
    css=CONSOLE_CSS,
    js=CONSOLE_JS,
    elem_classes=["console-shell"],
) as demo:
    gr.HTML(BRAND_HTML)

    # ── Generate (1/8) + playback deck (7/8), equal height ────────────────
    # Plain Row (not Group) so Gradio does not paint a full-bleed card wider than the racks.
    with gr.Row(elem_classes=["deck-shell", "deck-shell-row"], equal_height=True):
        with gr.Column(scale=1, min_width=0, elem_classes=["deck-gen-col"]):
            generate_btn = gr.Button(
                "Generate\nsong",
                variant="primary",
                elem_id="generate-song-btn",
            )
        with gr.Column(scale=7, min_width=0, elem_classes=["deck-panel"], elem_id="deck-panel"):
            gr.HTML("<div class='deck-label'>PLAYBACK DECK</div>")
            deck_hint = gr.HTML(
                "<div class='empty-hint'>No track yet — set the bay, then hit GENERATE SONG.</div>",
                elem_id="deck-hint",
            )
            audio_out = gr.Audio(
                label=None,
                show_label=False,
                type="filepath",
                elem_id="deck-audio",
                interactive=False,
                min_width=0,
                show_download_button=True,
                autoplay=False,
            )
            # Hidden download target — clicked by JS when "Download automatically" is on
            autodl_btn = gr.DownloadButton(
                label="Download track",
                value=None,
                elem_id="inspire-autodl",
                visible=True,
            )

    _dims0, _mood0, _key0, _meter0 = resolve_mood_bundle(None, DEFAULT_GENRE, GENRES[DEFAULT_GENRE]["keys"])

    # ── Three racks ─────────────────────────────────────────────────────────
    with gr.Row(elem_classes=["rack-row"]):
        # LEFT — machine
        with gr.Column(scale=2, min_width=200, elem_classes=["rack-panel"]):
            gr.HTML("<div class='rack-title'>MACHINE</div>")
            gr.HTML(DURATION_BAY_HTML)
            duration = gr.Number(
                value=DEFAULT_DURATION,
                precision=0,
                minimum=10,
                maximum=3600,
                label="Duration (sec)",
                elem_id="duration",
                elem_classes=["pad-hidden"],
            )
            model = gr.Dropdown(
                choices=list(MODELS.keys()),
                value=DEFAULT_MODEL,
                label="Model",
                elem_id="model",
            )
            gr.HTML(FORMAT_PADS_HTML)
            audio_format = gr.Dropdown(
                choices=list(AUDIO_FORMATS),
                value=DEFAULT_AUDIO_FORMAT,
                label="Export format",
                elem_id="audio_format",
                elem_classes=["pad-hidden"],
            )
            gr.HTML(AUTO_OPTS_HTML)
            auto_download = gr.Checkbox(
                label="Download automatically",
                value=False,
                elem_id="auto_download",
                elem_classes=["pad-hidden"],
            )
            auto_play = gr.Checkbox(
                label="Play automatically",
                value=False,
                elem_id="auto_play",
                elem_classes=["pad-hidden"],
            )
            gr.HTML(PRESET_HTML)

        # CENTER — creative
        with gr.Column(scale=4, min_width=360, elem_classes=["rack-panel"]):
            gr.HTML("<div class='rack-title'>MIX BAY · CREATIVE</div>")
            gr.HTML(USECASE_PADS_HTML)
            use_case = gr.Dropdown(
                choices=USE_CASES,
                value="(none)",
                label="Use case",
                elem_id="use_case",
                elem_classes=["pad-hidden"],
            )
            gr.HTML(GENRE_PADS_HTML)
            genre = gr.Dropdown(
                choices=list(GENRES.keys()),
                value=DEFAULT_GENRE,
                label="Genre",
                elem_id="genre",
                elem_classes=["pad-hidden"],
            )
            gr.HTML(STYLE_PADS_HTML)
            style = gr.Dropdown(
                choices=ALL_STYLES,
                value=GENRES[DEFAULT_GENRE]["styles"][0],
                label="Style",
                elem_id="style",
                elem_classes=["pad-hidden"],
            )
            gr.HTML(VOCAL_PADS_HTML)
            vocal = gr.Dropdown(
                choices=ALL_VOCALS,
                value=GENRES[DEFAULT_GENRE]["vocal"][0],
                label="Vocal character",
                elem_id="vocal",
                elem_classes=["pad-hidden"],
            )
            theme = gr.Textbox(
                label="Theme",
                placeholder="neon heartbreak, victory after failure…",
                lines=1,
                max_lines=1,
                elem_id="inspire_theme",
            )
            instructions = gr.Textbox(
                label="Additional instructions",
                placeholder="Writer / production notes — used as context, not sung "
                "(e.g. female POV, no rain metaphors, keep verses short…)",
                lines=3,
                elem_id="inspire_instructions",
            )

            gr.HTML("<div class='rack-title' style='margin-top:0.55rem'>MOOD MATRIX</div>")
            gr.HTML(MOOD_BOARD_HTML)
            # Hidden fields driven by the custom mood board / used by generation
            mood_dims = gr.Textbox(value=_dims0, elem_id="mood_dims", label="dims")
            mood = gr.Textbox(value=_mood0, elem_id="mood_label_box", label="mood")
            keyscale = gr.Textbox(value=_key0, elem_id="key_box", label="key")
            timesignature = gr.Textbox(value=_meter0, elem_id="meter_box", label="meter")

            gr.HTML(INSTRUMENTS_BROWSER_HTML)
            instruments = gr.CheckboxGroup(
                choices=INSTRUMENTS_ALL,
                value=[i for i in GENRES[DEFAULT_GENRE]["instruments"] if i in INSTRUMENTS_ALL][:3]
                or INSTRUMENTS_ALL[:3],
                label="Instruments / textures",
                elem_id="instrument-pads",
                elem_classes=["pad-hidden"],
            )

        # RIGHT — tempo + lyrics
        with gr.Column(scale=2, min_width=200, elem_classes=["rack-panel"]):
            gr.HTML("<div class='rack-title'>PERFORMANCE</div>")
            gr.HTML(KNOB_HTML)
            bpm = gr.Number(
                value=GENRES[DEFAULT_GENRE]["default_bpm"],
                precision=0,
                minimum=1,
                label="BPM",
                elem_id="bpm_number",
                elem_classes=["pad-hidden"],
            )

            gr.HTML("<div class='rack-title' style='margin-top:0.85rem'>LYRICS</div>")
            gr.HTML(LANGUAGE_PADS_HTML)
            language = gr.Dropdown(
                choices=list(VOCAL_LANGUAGES.keys()),
                value="English",
                label="Language",
                elem_id="language",
                elem_classes=["pad-hidden"],
            )
            gr.HTML(STRUCTURE_PADS_HTML)
            structure = gr.Dropdown(
                choices=list(STRUCTURES.keys()),
                value=list(STRUCTURES.keys())[0],
                label="Structure",
                elem_id="structure",
                elem_classes=["pad-hidden"],
            )
            gr.HTML(INSTRUMENTAL_PAD_HTML)
            instrumental = gr.Checkbox(
                label="Instrumental",
                value=False,
                elem_id="instrumental",
                elem_classes=["pad-hidden"],
            )
            gen_lyrics_btn = gr.Button("Generate lyrics", elem_id="lyrics-btn")
            lyrics = gr.Textbox(
                label="Pad",
                lines=10,
                placeholder="[verse] / [chorus] …",
                elem_id="lyrics",
            )

    # Hidden state for prompt (filled on generate / peek)
    prompt_state = gr.State("")
    meta_state = gr.State("No run yet.")

    # ── Action bar ──────────────────────────────────────────────────────────
    with gr.Row(elem_classes=["action-bar"]):
        surprise_btn = gr.Button("Surprise", elem_id="surprise-btn")
        prompt_btn = gr.Button("Prompt", elem_id="prompt-btn")
        info_btn = gr.Button("Info", elem_id="info-btn")
        expert_btn = gr.Button("Expert", elem_id="expert-btn")

    # ── Modals ──────────────────────────────────────────────────────────────
    with gr.Column(visible=False, elem_classes=["modal-shell"]) as prompt_modal:
        with gr.Group(elem_classes=["modal-card"]):
            gr.HTML("<div class='rack-title'>CAPTION + LYRICS</div>")
            prompt_view = gr.Textbox(label=None, show_label=False, lines=14, interactive=True)
            close_prompt = gr.Button("Close")

    with gr.Column(visible=False, elem_classes=["modal-shell"]) as info_modal:
        with gr.Group(elem_classes=["modal-card"]):
            gr.HTML("<div class='rack-title'>RUN INFO</div>")
            meta_view = gr.Textbox(label=None, show_label=False, lines=8, interactive=False)
            close_info = gr.Button("Close")

    with gr.Column(visible=False, elem_classes=["modal-shell"]) as expert_modal:
        with gr.Group(elem_classes=["modal-card", "expert-modal-card"]):
            gr.HTML("<div class='rack-title'>EXPERT</div>")
            steps = gr.Slider(4, 60, value=_model_cfg(DEFAULT_MODEL)["steps"], step=1, label="Steps", elem_id="steps")
            guidance = gr.Slider(1.0, 15.0, value=_model_cfg(DEFAULT_MODEL)["guidance"], step=0.5, label="Guidance", elem_id="guidance")
            shift = gr.Slider(1.0, 5.0, value=3.0, step=0.5, label="Shift", elem_id="shift")
            seed = gr.Number(value=-1, precision=0, label="Audio seed", elem_id="seed")
            random_seed = gr.Checkbox(value=True, label="Random seed", elem_id="random_seed")
            lyric_seed = gr.Number(value=-1, precision=0, label="Lyric seed", elem_id="lyric_seed")
            close_expert = gr.Button("Close")

    # BPM / pad-driven fields are hidden via .pad-hidden in CONSOLE_CSS

    # ── Wiring ───────────────────────────────────────────────────────────────
    genre.change(
        fn=apply_genre,
        inputs=[genre],
        outputs=[style, vocal, instruments, bpm, mood_dims, mood, keyscale, timesignature],
    )
    style.change(
        fn=apply_style_defaults,
        inputs=[genre, style],
        outputs=[bpm, instruments, mood_dims, mood, keyscale, timesignature],
    )
    use_case.change(
        fn=apply_use_case_bpm,
        inputs=[use_case],
        outputs=[bpm],
    )

    def _sync_expert(model_name):
        cfg = _model_cfg(model_name)
        return cfg["steps"], cfg["guidance"], cfg["shift"]

    model.change(fn=_sync_expert, inputs=[model], outputs=[steps, guidance, shift])

    gen_lyrics_btn.click(
        fn=on_generate_lyrics,
        inputs=[genre, mood, theme, structure, instrumental, lyric_seed, language, instructions],
        outputs=[lyrics],
        js="""
(genre, mood, theme, structure, instrumental, seed, language, instructions) => {
  const L = (window.readLiveConsole && window.readLiveConsole()) || {};
  return [
    L.genre != null ? L.genre : genre,
    mood,
    (L.theme != null && L.theme !== '') ? L.theme : theme,
    L.structure != null ? L.structure : structure,
    L.instrumental != null ? L.instrumental : instrumental,
    seed,
    L.language != null ? L.language : language,
    (L.instructions != null && L.instructions !== '') ? L.instructions : instructions,
  ];
}
""",
    )

    surprise_btn.click(
        fn=surprise_me,
        inputs=[genre, theme, language, instructions],
        outputs=[style, vocal, instruments, bpm, mood_dims, mood, keyscale, timesignature, structure, use_case, lyrics, prompt_view],
        js="""
(genre, theme, language, instructions) => {
  const L = (window.readLiveConsole && window.readLiveConsole()) || {};
  return [
    L.genre != null ? L.genre : genre,
    (L.theme != null && L.theme !== '') ? L.theme : theme,
    L.language != null ? L.language : language,
    (L.instructions != null && L.instructions !== '') ? L.instructions : instructions,
  ];
}
""",
    ).then(
        fn=lambda p: p,
        inputs=[prompt_view],
        outputs=[prompt_state],
    ).then(
        fn=lambda: None,
        js="""
() => {
  if (window.releaseConsoleOwned) window.releaseConsoleOwned();
  const pull = () => { if (window.syncBpmFromGradio) window.syncBpmFromGradio(); };
  setTimeout(pull, 80);
  setTimeout(pull, 250);
  return [];
}
""",
    )

    # Visible console UI is source of truth — Gradio hidden fields often lag.
    _SYNC_LIVE_JS = """
(model, lyrics, duration, bpm, keyscale, timesignature, language, instrumental, steps, guidance, shift, seed, random_seed, genre, style, mood, instruments, vocal, use_case, audio_format, auto_play, instructions) => {
  const L = (window.readLiveConsole && window.readLiveConsole()) || {};
  return [
    model,
    lyrics,
    L.duration != null ? L.duration : duration,
    L.bpm != null ? L.bpm : bpm,
    L.keyscale != null && L.keyscale !== '' ? L.keyscale : keyscale,
    L.timesignature != null && L.timesignature !== '' ? L.timesignature : timesignature,
    L.language != null ? L.language : language,
    L.instrumental != null ? L.instrumental : instrumental,
    steps,
    guidance,
    shift,
    seed,
    random_seed,
    L.genre != null ? L.genre : genre,
    L.style != null ? L.style : style,
    L.mood != null && L.mood !== '' ? L.mood : mood,
    (L.instruments && L.instruments.length) ? L.instruments : instruments,
    L.vocal != null ? L.vocal : vocal,
    L.use_case != null ? L.use_case : use_case,
    L.audio_format != null ? L.audio_format : audio_format,
    L.auto_play != null ? L.auto_play : auto_play,
    (L.instructions != null && L.instructions !== '') ? L.instructions : instructions,
  ];
}
"""

    generate_btn.click(
        fn=generate_music,
        inputs=[
            model,
            lyrics,
            duration,
            bpm,
            keyscale,
            timesignature,
            language,
            instrumental,
            steps,
            guidance,
            shift,
            seed,
            random_seed,
            genre,
            style,
            mood,
            instruments,
            vocal,
            use_case,
            audio_format,
            auto_play,
            instructions,
        ],
        outputs=[audio_out, meta_view, prompt_view, deck_hint, autodl_btn],
        js=_SYNC_LIVE_JS,
    ).then(
        fn=lambda p, m: (p, m),
        inputs=[prompt_view, meta_view],
        outputs=[prompt_state, meta_state],
    ).then(
        fn=lambda: None,
        js="""
() => {
  const L = (window.readLiveConsole && window.readLiveConsole()) || {};
  const wantDl = !!L.auto_download;
  if (wantDl) {
    let dlDone = false;
    const clickDl = () => {
      if (dlDone) return true;
      const r = document.getElementById('inspire-autodl');
      if (!r) return false;
      const el = r.matches('button, a') ? r : r.querySelector('button, a[download], a[href]');
      if (!el) return false;
      dlDone = true;
      el.click();
      return true;
    };
    if (!clickDl()) {
      setTimeout(clickDl, 300);
      setTimeout(clickDl, 900);
    }
  }

  // Gradio autoplay handles Play automatically (JS play() is blocked after the GPU wait).
  // If two <audio> nodes race, keep only the one that just started.
  const root = document.getElementById('deck-audio');
  if (root && !root.dataset.playDedupe) {
    root.dataset.playDedupe = '1';
    root.addEventListener('play', (e) => {
      const active = e.target;
      if (!(active instanceof HTMLMediaElement)) return;
      root.querySelectorAll('audio').forEach((el) => {
        if (el !== active && !el.paused) {
          try { el.pause(); } catch (_) {}
        }
      });
    }, true);
  }

  return [];
}
""",
    )

    prompt_btn.click(
        fn=peek_prompt,
        inputs=[genre, style, mood, instruments, vocal, bpm, use_case, instrumental, lyrics, instructions],
        outputs=[prompt_view, prompt_modal],
        js="""
(genre, style, mood, instruments, vocal, bpm, use_case, instrumental, lyrics, instructions) => {
  const L = (window.readLiveConsole && window.readLiveConsole()) || {};
  return [
    L.genre != null ? L.genre : genre,
    L.style != null ? L.style : style,
    L.mood != null && L.mood !== '' ? L.mood : mood,
    (L.instruments && L.instruments.length) ? L.instruments : instruments,
    L.vocal != null ? L.vocal : vocal,
    L.bpm != null ? L.bpm : bpm,
    L.use_case != null ? L.use_case : use_case,
    L.instrumental != null ? L.instrumental : instrumental,
    lyrics,
    (L.instructions != null && L.instructions !== '') ? L.instructions : instructions,
  ];
}
""",
    )
    info_btn.click(
        fn=lambda state, view: (view or state or "No run yet.", gr.update(visible=True)),
        inputs=[meta_state, meta_view],
        outputs=[meta_view, info_modal],
    )
    expert_btn.click(fn=lambda: gr.update(visible=True), outputs=[expert_modal])
    close_prompt.click(fn=lambda: gr.update(visible=False), outputs=[prompt_modal])
    close_info.click(fn=lambda: gr.update(visible=False), outputs=[info_modal])
    close_expert.click(fn=lambda: gr.update(visible=False), outputs=[expert_modal])


if __name__ == "__main__":
    demo.queue(max_size=10).launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
    )