app.py

from __future__ import annotations

import os

# ZeroGPU and third-party caches must be configured before importing libraries.
os.environ.setdefault("HF_HOME", "/tmp/.cache/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

try:  # ZeroGPU requires spaces to be imported before torch.
    import spaces
except ImportError:  # Local UI and tests work without the ZeroGPU package.
    class _SpacesFallback:
        @staticmethod
        def GPU(duration=None, size=None):
            def decorate(function):
                return function

            return decorate

    spaces = _SpacesFallback()  # type: ignore[assignment]

import base64
import hashlib
import html
import inspect
import math
import re
import shutil
import sys
import tempfile
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterator
from urllib.parse import parse_qs, urlparse

import gradio as gr

from notation import estimate_score_settings, generate_notation, safe_filename_stem
from studio_component import StudioViewer, initial_studio_value, normalize_viewer_value


ROOT = Path(__file__).resolve().parent
MODEL_ID = "MuScriptor/muscriptor-large"
MODEL_VARIANT = "large"
# The official API explicitly recommends keeping CFG at 1 for released,
# post-RL checkpoints. Instrument conditioning is handled separately below.
MODEL_CFG_COEF = 1.0
DECODE_PRESETS: dict[str, tuple[bool, int]] = {
    "fast": (False, 1),
    "balanced": (False, 2),
    "quality": (False, 4),
    "creative": (True, 1),
}
# Five minutes by default. Space owners can tune both limits without editing code.
MAX_AUDIO_SECONDS = float(os.environ.get("MAX_AUDIO_SECONDS", "300"))

# spaces 0.51 delegates the legal-duration policy to the ZeroGPU scheduler and
# does not impose the old 120-second client limit. Use a conservative 300-second
# application ceiling by default; Space owners can still lower it to favour
# queue priority without editing the planner.
ZEROGPU_MAX_FUNCTION_SECONDS = 300
ZEROGPU_MIN_CONFIGURED_CAP_SECONDS = 30
ZEROGPU_GPU_SIZES = frozenset({"large", "xlarge"})


def _configured_gpu_duration_cap(name: str, default: int = 300) -> int:
    """Read an owner cap, bounded by this Space's deliberate safety ceiling."""

    try:
        requested = int(os.environ.get(name, str(default)))
    except (TypeError, ValueError):
        requested = default
    return max(
        ZEROGPU_MIN_CONFIGURED_CAP_SECONDS,
        min(ZEROGPU_MAX_FUNCTION_SECONDS, requested),
    )


def _configured_runtime_scale(name: str) -> float:
    """Expose one simple calibration knob without allowing unsafe extremes."""

    try:
        value = float(os.environ.get(name, "1.0"))
    except (TypeError, ValueError):
        value = 1.0
    if not math.isfinite(value):
        value = 1.0
    return max(0.5, min(2.0, value))


def _configured_gpu_size(name: str, default: str) -> str:
    """Resolve a supported ZeroGPU slice without accepting invalid secrets."""

    requested = str(os.environ.get(name, default)).strip().lower()
    return requested if requested in ZEROGPU_GPU_SIZES else default


GPU_DURATION_CAP = _configured_gpu_duration_cap("MAX_GPU_SECONDS")
TRANSCRIPTION_GPU_SIZE = _configured_gpu_size("MUSCRIPTOR_GPU_SIZE", "xlarge")
# Runtime profiles below were measured on the default half-GPU `large` slice.
# A full xlarge slice exposes more memory and batching capacity, but a real
# 3:34 Fast run still took about 127 seconds at batch 1. Do not discount
# single-stream wall time merely because the full slice is selected; use its
# additional capacity only when the planner chooses a larger batch.
TRANSCRIPTION_GPU_COMPUTE_FACTORS = {
    "large": 1.0,
    "xlarge": 1.0,
}
TRANSCRIPTION_GPU_COMPUTE_FACTOR = TRANSCRIPTION_GPU_COMPUTE_FACTORS[
    TRANSCRIPTION_GPU_SIZE
]
TRANSCRIPTION_GPU_LABEL = "XL" if TRANSCRIPTION_GPU_SIZE == "xlarge" else "Large"
# The runtime model uses real-time factors (GPU seconds per second of audio),
# calibrated from the deployed large slice. A five-minute Fast transcription
# is modelled at ~181s sequentially before the safety margin, matching observed
# runs that finish just below 200s. The planner chooses the smallest batch
# needed to fit one valid ZeroGPU call; xlarge makes that batch less likely to
# be memory-bound, rather than pretending batch-1 decoding is twice as fast.
TRANSCRIPTION_RUNTIME_SCALE = _configured_runtime_scale("MUSCRIPTOR_RUNTIME_SCALE")
TRANSCRIPTION_RESERVATION_SAFETY = 1.10
TRANSCRIPTION_BATCH_EFFICIENCY = 0.82
TRANSCRIPTION_MIN_RESERVATION_SECONDS = 35
TRANSCRIPTION_BATCHED_MIN_RESERVATION_SECONDS = 110
TRANSCRIPTION_GPU_PROFILES: dict[str, tuple[float, float]] = {
    # (fixed setup seconds, sequential real-time factor)
    "fast": (16.0, 0.55),
    "balanced": (18.0, 1.15),
    "quality": (20.0, 2.00),
    "creative": (17.0, 0.72),
}
# MuScriptor officially supports batched GPU decoding when prelude forcing is
# disabled. Use the smallest batch that fits the ZeroGPU window: short clips
# retain strict sequential chunk continuity, while long clips trade a limited
# amount of boundary quality for enough throughput to remain deployable.
TRANSCRIPTION_MAX_BATCH_SIZE_BY_GPU: dict[str, dict[str, int]] = {
    "large": {"fast": 8, "balanced": 8, "quality": 4, "creative": 8},
    # The full 96 GB slice can safely leave more beam hypotheses resident.
    "xlarge": {"fast": 8, "balanced": 8, "quality": 5, "creative": 8},
}
TRANSCRIPTION_MAX_BATCH_SIZE = TRANSCRIPTION_MAX_BATCH_SIZE_BY_GPU[
    TRANSCRIPTION_GPU_SIZE
]

# Demucs is an optional source-preparation pass. The standard four-source
# htdemucs model is intentionally used instead of the experimental six-source
# model whose piano stem is documented as unreliable by the project.
DEMUCS_MODEL_NAME = os.environ.get("DEMUCS_MODEL_NAME", "htdemucs")
DEMUCS_SEGMENT_SECONDS = 7
DEMUCS_DEFAULT_PRESET = "balanced"
DEMUCS_PRESETS: dict[str, dict[str, float | int | str]] = {
    # The segment length stays fixed: htdemucs only supports up to 7.8 seconds,
    # and a stable chunk size is friendlier to ZeroGPU memory allocation.
    "fast": {"label": "Fast", "shifts": 0, "overlap": 0.10, "runtime_rtf": 0.12},
    "balanced": {"label": "Balanced", "shifts": 1, "overlap": 0.25, "runtime_rtf": 0.16},
    "quality": {"label": "Quality", "shifts": 2, "overlap": 0.50, "runtime_rtf": 0.30},
}
DEMUCS_SHIFTS = int(DEMUCS_PRESETS[DEMUCS_DEFAULT_PRESET]["shifts"])
DEMUCS_OVERLAP = float(DEMUCS_PRESETS[DEMUCS_DEFAULT_PRESET]["overlap"])
DEMUCS_GPU_SIZE = "large"
DEMUCS_GPU_LABEL = "Large"
DEMUCS_GPU_DURATION_CAP = _configured_gpu_duration_cap("DEMUCS_MAX_GPU_SECONDS")
DEMUCS_RUNTIME_SCALE = _configured_runtime_scale("DEMUCS_RUNTIME_SCALE")
DEMUCS_RESERVATION_SAFETY = 1.12
DEMUCS_SETUP_SECONDS = 12.0
DEMUCS_MIN_RESERVATION_SECONDS = 35
DEMUCS_STEM_ORDER = ("vocals", "drums", "bass", "other")
DEMUCS_STEM_LABELS = {
    "full": "Full mix",
    "vocals": "Vocals",
    "drums": "Drums",
    "bass": "Bass",
    "other": "Other",
}

# The remote importer deliberately accepts only public YouTube video URLs.
# This keeps yt-dlp from becoming a generic server-side URL fetcher and makes
# the same duration policy apply to uploads, recordings and remote sources.
YOUTUBE_ALLOWED_HOSTS = frozenset(
    {
        "youtube.com",
        "www.youtube.com",
        "m.youtube.com",
        "music.youtube.com",
        "youtu.be",
    }
)
YOUTUBE_MAX_URL_LENGTH = 2048
YOUTUBE_MAX_DOWNLOAD_BYTES = 96 * 1024 * 1024
YOUTUBE_SOCKET_TIMEOUT = 20
# The normal yt-dlp clients still begin with a YouTube webpage request. Shared
# datacenter IPs can be challenged at that layer even for a completely public
# video. This fallback asks the same official extractor to use public clients
# without that initial webpage request; it is deliberately not an auth bypass.
YOUTUBE_WEBPAGE_FREE_EXTRACTOR_ARGS = {
    "youtube": {
        "player_client": ["android_vr", "web_safari", "web_embedded"],
        "player_skip": ["webpage"],
    }
}

COLORS = (
    "#35c8ff",
    "#f6a623",
    "#fb7185",
    "#9b87f5",
    "#67e8a5",
    "#60a5fa",
    "#f472b6",
    "#a3e635",
    "#2dd4bf",
    "#f6c667",
)

FALLBACK_INSTRUMENTS = (
    "acoustic_piano",
    "electric_piano",
    "clean_electric_guitar",
    "distorted_electric_guitar",
    "acoustic_guitar",
    "electric_bass",
    "acoustic_bass",
    "violin",
    "viola",
    "cello",
    "strings",
    "trumpet",
    "trombone",
    "flutes",
    "clarinet",
    "soprano_and_alto_sax",
    "tenor_sax",
    "voice",
    "drums",
)

INSTRUMENT_CATEGORY_MEMBERS: tuple[tuple[str, tuple[str, ...]], ...] = (
    (
        "Keys & mallets",
        ("acoustic_piano", "electric_piano", "organ", "chromatic_percussion"),
    ),
    (
        "Guitars & basses",
        (
            "acoustic_guitar",
            "clean_electric_guitar",
            "distorted_electric_guitar",
            "acoustic_bass",
            "electric_bass",
        ),
    ),
    (
        "Strings & ensemble",
        (
            "violin",
            "viola",
            "cello",
            "contrabass",
            "orchestral_harp",
            "strings",
            "string_ensemble",
            "synth_strings",
        ),
    ),
    (
        "Brass",
        ("trumpet", "trombone", "tuba", "french_horn", "brass_section"),
    ),
    (
        "Woodwinds & saxophones",
        (
            "soprano_and_alto_sax",
            "tenor_sax",
            "baritone_sax",
            "oboe",
            "english_horn",
            "bassoon",
            "clarinet",
            "flutes",
        ),
    ),
    ("Percussion", ("timpani", "drums")),
    ("Voice, synth & other", ("voice", "orchestra_hit", "synth_lead", "synth_pad")),
)


MODEL: Any | None = None
MODEL_ERROR = ""
DEMUCS_SEPARATOR: Any | None = None
DEMUCS_ERROR = ""
NoteStartEvent: Any = None
NoteEndEvent: Any = None
ProgressEvent: Any = None
INSTRUMENT_NAMES = list(FALLBACK_INSTRUMENTS)


def _load_model() -> Any | None:
    global MODEL_ERROR, NoteStartEvent, NoteEndEvent, ProgressEvent, INSTRUMENT_NAMES
    try:
        import torch
        from huggingface_hub import hf_hub_download
        from muscriptor.events import NoteEndEvent as _NoteEndEvent
        from muscriptor.events import NoteStartEvent as _NoteStartEvent
        from muscriptor.events import ProgressEvent as _ProgressEvent
        from muscriptor.tokenizer.mt3 import MT3Tokenizer, MT3_FULL_PLUS_GROUP_NAMES
        from muscriptor.transcription_model import (
            TranscriptionModel,
            _build_model,
            _remap_single_codebook_keys,
            _resolve_config,
            _resolve_source,
        )
        from safetensors.torch import load_file
    except Exception as exc:
        MODEL_ERROR = f"MuScriptor dependencies are unavailable: {type(exc).__name__}: {exc}"
        return None

    NoteStartEvent = _NoteStartEvent
    NoteEndEvent = _NoteEndEvent
    ProgressEvent = _ProgressEvent
    INSTRUMENT_NAMES = list(MT3_FULL_PLUS_GROUP_NAMES)

    token = os.environ.get("HF_TOKEN")
    if not token:
        MODEL_ERROR = (
            "HF_TOKEN is missing. Accept the MuScriptor large model license, then "
            "add a read-only HF_TOKEN secret in the Space settings."
        )
        return None

    try:
        source = _resolve_source(MODEL_VARIANT)
        weights_path = Path(
            hf_hub_download(
                repo_id=MODEL_ID,
                filename="model.safetensors",
                token=token,
            )
        )
        config = _resolve_config(source, weights_path)
        device = torch.device("cuda")
        model = _build_model(device, config)
        model.eval()
        state_dict = _remap_single_codebook_keys(load_file(str(weights_path), device="cpu"))
        model.load_state_dict(state_dict)
        model.to("cuda")
        tokenizer = MT3Tokenizer(instrument_vocabulary="MT3_FULL_PLUS", max_shift_steps=1001)
        return TranscriptionModel(model=model, tokenizer=tokenizer, device=device)
    except Exception as exc:
        MODEL_ERROR = f"The model could not be loaded: {type(exc).__name__}: {exc}"
        return None


if os.environ.get("MUSCRIPTOR_SKIP_MODEL_LOAD") != "1":
    started = time.perf_counter()
    print(f"[MuScriptor Studio] Loading the {MODEL_VARIANT} model…", flush=True)
    MODEL = _load_model()
    if MODEL is None:
        print(f"[MuScriptor Studio] Interface-only mode: {MODEL_ERROR}", file=sys.stderr, flush=True)
    else:
        print(f"[MuScriptor Studio] Model ready in {time.perf_counter() - started:.2f}s.", flush=True)


def _load_demucs_separator() -> Any | None:
    """Load the public htdemucs model once and register it with ZeroGPU."""

    global DEMUCS_ERROR
    try:
        from demucs.api import Separator

        separator = Separator(
            model=DEMUCS_MODEL_NAME,
            device="cuda",
            shifts=DEMUCS_SHIFTS,
            overlap=DEMUCS_OVERLAP,
            split=True,
            segment=DEMUCS_SEGMENT_SECONDS,
            jobs=0,
            progress=False,
        )
        # ZeroGPU emulates this placement at module scope and restores the
        # packed model inside the decorated separation call.
        separator.model.to("cuda")
        return separator
    except Exception as exc:
        DEMUCS_ERROR = f"Demucs could not be loaded: {type(exc).__name__}: {exc}"
        return None


if os.environ.get("MUSCRIPTOR_SKIP_DEMUCS_LOAD") != "1" and os.environ.get("MUSCRIPTOR_SKIP_MODEL_LOAD") != "1":
    started = time.perf_counter()
    print(f"[MuScriptor Studio] Loading Demucs {DEMUCS_MODEL_NAME}…", flush=True)
    DEMUCS_SEPARATOR = _load_demucs_separator()
    if DEMUCS_SEPARATOR is None:
        print(f"[MuScriptor Studio] Demucs unavailable: {DEMUCS_ERROR}", file=sys.stderr, flush=True)
    else:
        print(f"[MuScriptor Studio] Demucs ready in {time.perf_counter() - started:.2f}s.", flush=True)


def _audio_path(value: Any) -> str | None:
    if value is None:
        return None
    if isinstance(value, (str, Path)):
        return str(value)
    if isinstance(value, dict):
        path = value.get("path") or value.get("name")
        return str(path) if path else None
    path = getattr(value, "path", None) or getattr(value, "name", None)
    return str(path) if path else None


def _audio_duration(path: str | None) -> float:
    if not path:
        return 0.0
    try:
        import soundfile as sf

        return float(sf.info(path).duration)
    except Exception:
        return 0.0


def _short_duration(seconds: float) -> str:
    total = max(0, int(round(float(seconds or 0))))
    minutes, remainder = divmod(total, 60)
    return f"{minutes}:{remainder:02d}"


def _validate_youtube_url(value: str | None) -> str:
    """Accept known video URL shapes and return a canonical YouTube watch URL."""

    candidate = str(value or "").strip()
    if not candidate:
        raise ValueError("Paste a public YouTube video URL first.")
    if len(candidate) > YOUTUBE_MAX_URL_LENGTH:
        raise ValueError("This YouTube URL is too long.")
    if "://" not in candidate:
        candidate = f"https://{candidate}"
    parsed = urlparse(candidate)
    hostname = (parsed.hostname or "").lower().rstrip(".")
    if parsed.scheme.lower() != "https" or hostname not in YOUTUBE_ALLOWED_HOSTS:
        raise ValueError("Use a valid HTTPS URL from youtube.com or youtu.be.")
    try:
        custom_port = parsed.port
    except ValueError as exc:
        raise ValueError("This YouTube URL contains an invalid port.") from exc
    if parsed.username or parsed.password or custom_port:
        raise ValueError("Authenticated or custom-port YouTube URLs are not supported.")

    path_parts = [part for part in parsed.path.split("/") if part]
    video_id = ""
    if hostname == "youtu.be" and path_parts:
        video_id = path_parts[0]
    elif path_parts == ["watch"]:
        video_id = (parse_qs(parsed.query).get("v") or [""])[0]
    elif len(path_parts) == 2 and path_parts[0] in {"shorts", "embed", "live"}:
        video_id = path_parts[1]
    elif path_parts and path_parts[0].lower() == "playlist":
        raise ValueError("Playlists are not supported. Paste one video URL instead.")
    else:
        raise ValueError("Paste a direct YouTube video, Short or youtu.be link.")

    if not re.fullmatch(r"[A-Za-z0-9_-]{6,32}", video_id):
        raise ValueError("This URL does not contain a valid YouTube video identifier.")
    return f"https://www.youtube.com/watch?v={video_id}"


def _validate_youtube_info(info: Any) -> tuple[str, float, str]:
    """Validate metadata before any media bytes are downloaded."""

    if not info or not hasattr(info, "get"):
        raise ValueError("YouTube did not return readable video metadata.")
    if info.get("_type") in {"playlist", "multi_video"} or info.get("entries") is not None:
        raise ValueError("Playlists and channel pages are not supported. Choose one video.")
    live_status = str(info.get("live_status") or "").lower()
    if bool(info.get("is_live")) or live_status in {"is_live", "is_upcoming"}:
        raise ValueError("Live and upcoming streams are not supported.")
    try:
        duration = float(info.get("duration") or 0)
    except (TypeError, ValueError):
        duration = 0.0
    if not math.isfinite(duration) or duration <= 0:
        raise ValueError("The video duration could not be verified, so it was not imported.")
    if duration > MAX_AUDIO_SECONDS:
        raise ValueError(
            f"This video is {_short_duration(duration)}. The Space accepts up to "
            f"{_short_duration(MAX_AUDIO_SECONDS)}."
        )
    title = str(info.get("title") or "YouTube audio").strip() or "YouTube audio"
    video_id = str(info.get("id") or "video").strip() or "video"
    return title, duration, video_id


def _youtube_import_status_html(
    state: str = "idle",
    *,
    title: str = "",
    duration: float = 0,
    message: str = "",
) -> str:
    safe_title = html.escape(title or "YouTube video")
    safe_message = html.escape(message or "The video could not be imported.")
    copy: dict[str, tuple[str, str]] = {
        "idle": (
            "Paste a public YouTube video",
            "Use content you own or are authorised to process · one video · no playlists or live streams.",
        ),
        "working": (
            "Checking the YouTube source…",
            "Reading metadata first, then loading only the audio into this session.",
        ),
        "complete": (
            f"{safe_title} loaded",
            f"{_short_duration(duration)} · ready in the audio player · the remote source is not retained.",
        ),
        "error": ("YouTube import stopped", safe_message),
    }
    heading, detail = copy.get(state, copy["idle"])
    return (
        f'<div class="youtube-import-status" data-state="{html.escape(state)}">'
        f'<i aria-hidden="true"></i><div><strong>{heading}</strong><span>{detail}</span></div></div>'
    )


def reset_youtube_import_status() -> str:
    return _youtube_import_status_html()


def begin_youtube_import(url: str | None) -> tuple[str, dict[str, Any]]:
    return (
        _youtube_import_status_html("working"),
        gr.update(interactive=False, value="Loading audio…"),
    )


def _youtube_raw_error(exc: Exception, messages: list[str] | None = None) -> str:
    candidates = [str(exc or ""), *(messages or [])]
    cleaned = []
    for candidate in candidates:
        message = re.sub(r"\x1b\[[0-9;]*m", "", str(candidate or "")).strip()
        message = re.sub(r"^ERROR:\s*", "", message, flags=re.IGNORECASE)
        if message and message not in cleaned:
            cleaned.append(message)
    return " · ".join(cleaned)[-1600:]


def _youtube_is_access_challenge(message: str) -> bool:
    lowered = str(message or "").lower().replace("’", "'")
    return any(
        marker in lowered
        for marker in (
            "sign in to confirm",
            "not a bot",
            "po token",
            "proof of origin",
            "http error 403",
            "request is blocked",
        )
    )


def _clean_youtube_error(exc: Exception) -> str:
    message = _youtube_raw_error(exc)
    lowered = message.lower()
    if _youtube_is_access_challenge(message):
        return (
            "YouTube blocked this Space's shared server address with an anti-bot check, "
            "even though the video may be public. This is not a privacy-setting error. "
            "Upload the audio file instead; reliable server imports require a PO-token provider."
        )
    if "429" in lowered or "too many requests" in lowered:
        return (
            "YouTube temporarily rate-limited the shared server address used by this Space. "
            "Try again later or upload the audio file directly."
        )
    if any(marker in lowered for marker in ("private video", "members-only", "login required")):
        return "This video requires authentication. Only public, unauthenticated sources are supported."
    return (message or "The video could not be imported.")[:420]


class _YouTubeCaptureLogger:
    """Keep actionable yt-dlp diagnostics without writing its noisy output to UI logs."""

    def __init__(self) -> None:
        self.messages: list[str] = []

    def debug(self, message: str) -> None:
        pass

    def info(self, message: str) -> None:
        pass

    def warning(self, message: str) -> None:
        if message:
            self.messages.append(str(message))

    def error(self, message: str) -> None:
        if message:
            self.messages.append(str(message))


def _extract_youtube_preview(
    yt_dlp_module: Any,
    url: str,
    base_options: dict[str, Any],
) -> tuple[Any, dict[str, Any]]:
    """Read metadata, retrying once without the challenge-prone webpage layer."""

    attempts = ({}, {"extractor_args": YOUTUBE_WEBPAGE_FREE_EXTRACTOR_ARGS})
    challenge_errors: list[str] = []
    last_error: Exception | None = None
    for attempt_index, attempt_options in enumerate(attempts):
        logger = _YouTubeCaptureLogger()
        options = {**base_options, **attempt_options, "logger": logger}
        try:
            with yt_dlp_module.YoutubeDL(options) as downloader:
                return downloader.extract_info(url, download=False), attempt_options
        except Exception as exc:
            last_error = exc
            detail = _youtube_raw_error(exc, logger.messages)
            if not _youtube_is_access_challenge(detail) or attempt_index == len(attempts) - 1:
                if challenge_errors:
                    detail = " · ".join([*challenge_errors, detail])
                raise RuntimeError(detail or "YouTube metadata extraction failed.") from exc
            challenge_errors.append(detail)
    raise RuntimeError("YouTube metadata extraction failed.") from last_error


def _download_youtube_audio(
    url: str,
    progress: gr.Progress,
) -> tuple[Path, str, float, str]:
    """Download one public YouTube video's best audio and convert it to WAV."""

    try:
        import yt_dlp
    except ImportError as exc:
        raise RuntimeError(
            "yt-dlp is unavailable. Rebuild the Space with the provided requirements.txt."
        ) from exc

    base_options: dict[str, Any] = {
        "quiet": True,
        "no_warnings": True,
        "noplaylist": True,
        "socket_timeout": YOUTUBE_SOCKET_TIMEOUT,
        "retries": 2,
        "extractor_retries": 2,
        "fragment_retries": 2,
        "concurrent_fragment_downloads": 1,
        "cachedir": False,
    }
    output_dir: Path | None = None
    try:
        progress(0.03, desc="Checking the public YouTube video")
        preview, extraction_options = _extract_youtube_preview(yt_dlp, url, base_options)
        title, duration, video_id = _validate_youtube_info(preview)

        output_dir = Path(tempfile.mkdtemp(prefix="muscriptor-youtube-"))

        def progress_hook(status: dict[str, Any]) -> None:
            if status.get("status") == "finished":
                progress(0.9, desc="Converting the audio to WAV")
                return
            if status.get("status") != "downloading":
                return
            downloaded = float(status.get("downloaded_bytes") or 0)
            total = float(status.get("total_bytes") or status.get("total_bytes_estimate") or 0)
            fraction = downloaded / total if total > 0 else 0.08
            progress(0.12 + 0.72 * max(0.0, min(1.0, fraction)), desc="Downloading YouTube audio")

        download_options = {
            **base_options,
            **extraction_options,
            "format": "bestaudio/best",
            "outtmpl": str(output_dir / "%(id)s.%(ext)s"),
            "restrictfilenames": True,
            "max_filesize": YOUTUBE_MAX_DOWNLOAD_BYTES,
            "overwrites": True,
            "continuedl": False,
            "progress_hooks": [progress_hook],
            "postprocessors": [
                {
                    "key": "FFmpegExtractAudio",
                    "preferredcodec": "wav",
                }
            ],
        }
        progress(0.1, desc=f"Loading {_short_duration(duration)} of audio")
        download_logger = _YouTubeCaptureLogger()
        try:
            with yt_dlp.YoutubeDL({**download_options, "logger": download_logger}) as downloader:
                downloaded_info = downloader.extract_info(url, download=True)
        except Exception as exc:
            raise RuntimeError(_youtube_raw_error(exc, download_logger.messages)) from exc
        downloaded_title, downloaded_duration, downloaded_id = _validate_youtube_info(downloaded_info)

        wav_candidates = list(output_dir.glob("*.wav"))
        if not wav_candidates:
            raise RuntimeError("FFmpeg did not produce the expected WAV audio file.")
        wav_path = max(wav_candidates, key=lambda candidate: candidate.stat().st_size)
        title_slug = safe_filename_stem(downloaded_title, "youtube-audio")[:80]
        id_slug = safe_filename_stem(downloaded_id, "video")[:24]
        final_path = output_dir / f"{title_slug}-{id_slug}.wav"
        if wav_path != final_path:
            wav_path.replace(final_path)

        decoded_duration = _audio_duration(str(final_path))
        if decoded_duration <= 0:
            raise RuntimeError("The converted YouTube audio could not be read.")
        if decoded_duration > MAX_AUDIO_SECONDS + 1:
            raise RuntimeError("The converted audio exceeds this Space's duration limit.")
        progress(1, desc="YouTube audio loaded")
        return final_path, downloaded_title, downloaded_duration, downloaded_id
    except Exception:
        if output_dir is not None:
            shutil.rmtree(output_dir, ignore_errors=True)
        raise


def import_youtube_audio(
    url: str | None,
    progress: gr.Progress = gr.Progress(),
) -> tuple[Any, Any, Any, Any, Any, Any, str, dict[str, Any]]:
    """Load a validated public YouTube source into the existing Gradio player."""

    started = time.perf_counter()
    try:
        normalized_url = _validate_youtube_url(url)
        audio_path, title, duration, video_id = _download_youtube_audio(normalized_url, progress)
        source_session, selector, stem_status, stem_button, stem_files = register_source_audio(
            str(audio_path)
        )
        source_session.update(
            {
                "source_kind": "youtube",
                "source_title": title,
                "source_video_id": video_id,
                "duration": duration,
            }
        )
        print(
            f"[MuScriptor Studio] Imported public YouTube video {video_id} "
            f"({_short_duration(duration)}) in {time.perf_counter() - started:.1f}s.",
            flush=True,
        )
        return (
            str(audio_path),
            source_session,
            selector,
            stem_status,
            stem_button,
            stem_files,
            _youtube_import_status_html("complete", title=title, duration=duration),
            gr.update(interactive=True, value="Load another video"),
        )
    except Exception as exc:
        message = _clean_youtube_error(exc)
        print(f"[MuScriptor Studio] YouTube import stopped: {message}", file=sys.stderr, flush=True)
        return (
            gr.skip(),
            gr.skip(),
            gr.skip(),
            gr.skip(),
            gr.skip(),
            gr.skip(),
            _youtube_import_status_html("error", message=message),
            gr.update(interactive=True, value="Try again"),
        )


def _stem_source_path(session: dict[str, Any] | None, selection: str = "full") -> str | None:
    session = session or {}
    key = str(selection or "full")
    if key == "full":
        return _audio_path(session.get("original_path"))
    return _audio_path((session.get("stems") or {}).get(key))


def _stem_choices(session: dict[str, Any] | None = None) -> list[tuple[str, str]]:
    session = session or {}
    stems = session.get("stems") or {}
    choices = [(DEMUCS_STEM_LABELS["full"], "full")]
    choices.extend(
        (DEMUCS_STEM_LABELS[name], name)
        for name in DEMUCS_STEM_ORDER
        if stems.get(name)
    )
    return choices


def _stem_files(session: dict[str, Any] | None = None) -> list[str] | None:
    session = session or {}
    values = [
        str(path)
        for name in DEMUCS_STEM_ORDER
        if (path := (session.get("stems") or {}).get(name))
    ]
    return values or None


def _stem_selector_update(
    session: dict[str, Any] | None = None,
    *,
    interactive: bool | None = None,
    value: str | None = None,
) -> dict[str, Any]:
    session = session or {}
    choices = _stem_choices(session)
    selected = value or str(session.get("selected") or "full")
    if selected not in {choice_value for _, choice_value in choices}:
        selected = "full"
    enabled = len(choices) > 1 if interactive is None else bool(interactive)
    return gr.update(choices=choices, value=selected, interactive=enabled, visible=True)


def _stem_status_html(
    state: str = "waiting",
    *,
    audio_name: str = "",
    selected: str = "full",
    message: str = "",
    preset: str = DEMUCS_DEFAULT_PRESET,
    reservation_seconds: int = 0,
) -> str:
    safe_name = html.escape(audio_name or "the source audio")
    label = html.escape(DEMUCS_STEM_LABELS.get(selected, _display_name(selected)))
    _, settings = _demucs_settings(preset)
    preset_label = html.escape(str(settings["label"]))
    copy: dict[str, tuple[str, str]] = {
        "waiting": (
            "Optional stem separation",
            "Upload or record audio, then isolate vocals, drums, bass and accompaniment.",
        ),
        "ready": (
            "Ready for source separation",
            f"{safe_name} stays available as the full mix. Demucs runs only when requested.",
        ),
        "working": (
            "Separating four audio stems…",
            f"ZeroGPU is running the {preset_label} preset on vocals, drums, bass and accompaniment.",
        ),
        "complete": (
            "Four stems ready",
            f"{preset_label} separation complete. Choose a source below before transcription.",
        ),
        "selected": (
            f"{label} loaded in the audio player",
            "This is now the source MuScriptor will transcribe.",
        ),
        "unavailable": (
            "Demucs unavailable",
            "The audio can still be transcribed as a full mix on this runtime.",
        ),
        "error": (
            "Stem separation stopped",
            html.escape(message or "Demucs could not process this audio source."),
        ),
    }
    title, detail = copy.get(state, copy["waiting"])
    if state == "working" and reservation_seconds:
        detail += (
            f" Up to {int(reservation_seconds)}s of ZeroGPU {DEMUCS_GPU_LABEL} "
            "runtime is requested."
        )
    return (
        f'<div class="stem-status" data-state="{html.escape(state)}"><i aria-hidden="true"></i><div>'
        f"<strong>{title}</strong><span>{detail}</span>"
        f'</div><b>{html.escape(DEMUCS_MODEL_NAME)} · {preset_label}</b></div>'
    )


def register_source_audio(audio: Any) -> tuple[dict[str, Any], dict[str, Any], str, dict[str, Any], None]:
    """Reset stale stems only when the user uploads or records a new source."""

    path = _audio_path(audio)
    if not path:
        return (
            {},
            _stem_selector_update({}, interactive=False),
            _stem_status_html(),
            gr.update(interactive=False, value="Separate into stems"),
            None,
        )
    session = {
        "original_path": path,
        "original_name": Path(path).name,
        "selected": "full",
        "stems": {},
        "model": DEMUCS_MODEL_NAME,
    }
    available = DEMUCS_SEPARATOR is not None
    status = _stem_status_html(
        "ready" if available else "unavailable",
        audio_name=Path(path).name,
    )
    return (
        session,
        _stem_selector_update(session, interactive=False),
        status,
        gr.update(
            interactive=available,
            value="Separate into stems" if available else "Demucs unavailable",
        ),
        None,
    )


def clear_source_audio() -> tuple[dict[str, Any], dict[str, Any], str, dict[str, Any], None]:
    return register_source_audio(None)


def begin_stem_separation(
    audio: Any,
    session: dict[str, Any] | None,
    demucs_preset: str = DEMUCS_DEFAULT_PRESET,
) -> tuple[str, dict[str, Any], dict[str, Any]]:
    path = _stem_source_path(session, "full") or _audio_path(audio)
    if not path:
        raise gr.Error("Add or record audio before running Demucs.")
    duration = _audio_duration(path)
    if duration <= 0:
        raise gr.Error("The audio duration could not be read.")
    if duration > MAX_AUDIO_SECONDS:
        raise gr.Error(
            f"This Space accepts audio up to {MAX_AUDIO_SECONDS / 60:.1f} minutes."
        )
    runtime_plan = _demucs_runtime_plan(duration, demucs_preset)
    if not runtime_plan["fits"]:
        supported = _short_duration(runtime_plan["max_supported_seconds"])
        raise gr.Error(
            f"{str(runtime_plan['preset']).title()} Demucs cannot fit this recording "
            f"inside a {DEMUCS_GPU_DURATION_CAP}s ZeroGPU call. Choose Fast or Balanced, "
            f"or trim this preset to about {supported}."
        )
    return (
        _stem_status_html(
            "working",
            audio_name=Path(path).name,
            preset=demucs_preset,
            reservation_seconds=int(runtime_plan["requested_seconds"]),
        ),
        gr.update(interactive=False, value="Separating stems…"),
        _stem_selector_update(session, interactive=False),
    )


def _demucs_settings(preset: str | None = None) -> tuple[str, dict[str, float | int | str]]:
    """Resolve a bounded UI preset without exposing unsafe free-form GPU settings."""

    key = str(preset or DEMUCS_DEFAULT_PRESET).strip().lower()
    if key not in DEMUCS_PRESETS:
        key = DEMUCS_DEFAULT_PRESET
    return key, DEMUCS_PRESETS[key]


def _estimate_demucs_gpu_duration(
    audio: Any,
    session: dict[str, Any] | None = None,
    demucs_preset: str = DEMUCS_DEFAULT_PRESET,
    *args: Any,
    **kwargs: Any,
) -> int:
    path = _stem_source_path(session, "full") or _audio_path(audio)
    duration = _audio_duration(path) or 15.0
    plan = _demucs_runtime_plan(duration, demucs_preset)
    return int(plan["requested_seconds"])


def _demucs_runtime_plan(
    duration: float,
    demucs_preset: str | None,
) -> dict[str, Any]:
    """Estimate one Demucs pass and expose the safe duration for its preset."""

    preset, settings = _demucs_settings(demucs_preset)
    runtime_rtf = float(settings["runtime_rtf"])
    safe_duration = max(0.0, float(duration or 0.0))
    # Demucs complexity is mostly governed by duration, overlap and shifts.
    # Balanced is calibrated around one minute for a five-minute source; the
    # safety factor covers model restoration, encoding and slightly slower runs.
    raw_seconds = DEMUCS_SETUP_SECONDS + safe_duration * runtime_rtf
    estimate = math.ceil(raw_seconds * DEMUCS_RUNTIME_SCALE * DEMUCS_RESERVATION_SAFETY)
    fits = estimate <= DEMUCS_GPU_DURATION_CAP
    requested_seconds = (
        min(DEMUCS_GPU_DURATION_CAP, max(DEMUCS_MIN_RESERVATION_SECONDS, estimate))
        if fits
        else min(DEMUCS_GPU_DURATION_CAP, ZEROGPU_MIN_CONFIGURED_CAP_SECONDS)
    )
    available_compute_seconds = max(
        0.0,
        DEMUCS_GPU_DURATION_CAP / (DEMUCS_RUNTIME_SCALE * DEMUCS_RESERVATION_SAFETY)
        - DEMUCS_SETUP_SECONDS,
    )
    max_supported_seconds = max(
        0,
        math.floor(available_compute_seconds / runtime_rtf),
    )
    return {
        "preset": preset,
        "raw_seconds": raw_seconds,
        "estimated_seconds": estimate,
        "requested_seconds": requested_seconds,
        "fits": fits,
        "max_supported_seconds": max_supported_seconds,
    }


@spaces.GPU(size=DEMUCS_GPU_SIZE, duration=_estimate_demucs_gpu_duration)
def separate_audio_stems(
    audio: Any,
    session: dict[str, Any] | None,
    demucs_preset: str = DEMUCS_DEFAULT_PRESET,
    progress: gr.Progress = gr.Progress(),
) -> tuple[dict[str, Any], dict[str, Any], str, dict[str, Any], list[str] | None]:
    """Run htdemucs once and retain the full mix plus four selectable WAV stems."""

    previous = dict(session or {})
    path = _stem_source_path(previous, "full") or _audio_path(audio)
    audio_name = str(previous.get("original_name") or (Path(path).name if path else ""))
    retry_update = gr.update(interactive=DEMUCS_SEPARATOR is not None, value="Retry separation")
    if not path:
        return (
            previous,
            _stem_selector_update(previous),
            _stem_status_html("error", message="Add or record audio before running Demucs."),
            gr.update(interactive=False, value="Separate into stems"),
            _stem_files(previous),
        )
    duration = _audio_duration(path)
    if duration <= 0:
        return (
            previous,
            _stem_selector_update(previous),
            _stem_status_html("error", audio_name=audio_name, message="The audio duration could not be read."),
            retry_update,
            _stem_files(previous),
        )
    if duration > MAX_AUDIO_SECONDS:
        return (
            previous,
            _stem_selector_update(previous),
            _stem_status_html(
                "error",
                audio_name=audio_name,
                message=f"This Space accepts audio up to {MAX_AUDIO_SECONDS / 60:.1f} minutes.",
            ),
            retry_update,
            _stem_files(previous),
        )
    if DEMUCS_SEPARATOR is None:
        return (
            previous,
            _stem_selector_update(previous),
            _stem_status_html("unavailable", audio_name=audio_name),
            gr.update(interactive=False, value="Demucs unavailable"),
            _stem_files(previous),
        )

    runtime_plan = _demucs_runtime_plan(duration, demucs_preset)
    if not runtime_plan["fits"]:
        supported = _short_duration(runtime_plan["max_supported_seconds"])
        message = (
            f"{str(runtime_plan['preset']).title()} Demucs cannot fit this recording "
            f"inside a {DEMUCS_GPU_DURATION_CAP}s ZeroGPU call. Choose Fast or Balanced, "
            f"or trim this preset to about {supported}."
        )
        return (
            previous,
            _stem_selector_update(previous),
            _stem_status_html("error", audio_name=audio_name, message=message),
            retry_update,
            _stem_files(previous),
        )

    print(
        "[MuScriptor Studio] Demucs plan: "
        f"{duration:.1f}s audio · {runtime_plan['preset']} · "
        f"request {runtime_plan['requested_seconds']}s ZeroGPU {DEMUCS_GPU_SIZE}.",
        flush=True,
    )
    started = time.perf_counter()
    separator = DEMUCS_SEPARATOR
    preset_key, settings = _demucs_settings(demucs_preset)
    shifts = int(settings["shifts"])
    overlap = float(settings["overlap"])
    segment_frames = int(separator.samplerate * DEMUCS_SEGMENT_SECONDS)

    def on_progress(info: dict[str, Any]) -> None:
        if info.get("state") != "end":
            return
        audio_length = max(1, int(info.get("audio_length") or 1))
        segment_end = min(audio_length, int(info.get("segment_offset") or 0) + segment_frames)
        model_count = max(1, int(info.get("models") or 1))
        model_index = max(0, int(info.get("model_idx_in_bag") or 0))
        shift_index = max(0, int(info.get("shift_idx") or 0))
        shift_count = max(1, shifts)
        unit = (model_index * shift_count + shift_index + segment_end / audio_length) / (
            model_count * shift_count
        )
        progress(min(0.92, 0.04 + 0.86 * unit), desc="Separating audio with Demucs")

    try:
        from demucs.api import save_audio

        progress(0.02, desc="Preparing the full mix")
        separator.update_parameter(
            shifts=shifts,
            overlap=overlap,
            segment=DEMUCS_SEGMENT_SECONDS,
            callback=on_progress,
            callback_arg={"audio_name": audio_name, "preset": preset_key},
        )
        _, separated = separator.separate_audio_file(Path(path))
        missing = [name for name in DEMUCS_STEM_ORDER if name not in separated]
        if missing:
            raise RuntimeError(f"Demucs did not return: {', '.join(missing)}")

        progress(0.94, desc="Writing four WAV stems")
        export_stem = safe_filename_stem(Path(audio_name).stem, "demucs-source")
        output_dir = Path(tempfile.mkdtemp(prefix="muscriptor-demucs-"))
        stem_paths: dict[str, str] = {}
        for name in DEMUCS_STEM_ORDER:
            stem_path = output_dir / f"{export_stem}-{name}-stem.wav"
            save_audio(
                separated[name].detach().cpu(),
                stem_path,
                samplerate=separator.samplerate,
                clip="rescale",
                bits_per_sample=16,
                as_float=False,
            )
            stem_paths[name] = str(stem_path)

        result = {
            "original_path": path,
            "original_name": audio_name,
            "selected": "full",
            "stems": stem_paths,
            "model": DEMUCS_MODEL_NAME,
            "demucs_preset": preset_key,
            "demucs_shifts": shifts,
            "demucs_overlap": overlap,
            "duration": duration,
            "elapsed": time.perf_counter() - started,
            "reserved_gpu_seconds": int(runtime_plan["requested_seconds"]),
        }
        progress(1, desc="Four stems ready")
        print(
            f"[MuScriptor Studio] Demucs separated {audio_name} "
            f"({duration:.1f}s) in {result['elapsed']:.1f}s.",
            flush=True,
        )
        return (
            result,
            _stem_selector_update(result, interactive=True, value="full"),
            _stem_status_html("complete", audio_name=audio_name, preset=preset_key),
            gr.update(interactive=True, value="Separate original again"),
            _stem_files(result),
        )
    except Exception as exc:
        message = f"{type(exc).__name__}: {exc}"
        print(f"[MuScriptor Studio] Demucs stopped: {message}", file=sys.stderr, flush=True)
        return (
            previous,
            _stem_selector_update(previous),
            _stem_status_html("error", audio_name=audio_name, message=message),
            retry_update,
            _stem_files(previous),
        )
    finally:
        try:
            separator.update_parameter(
                shifts=DEMUCS_SHIFTS,
                overlap=DEMUCS_OVERLAP,
                segment=DEMUCS_SEGMENT_SECONDS,
                callback=None,
                callback_arg=None,
            )
        except Exception:
            pass


def select_stem_source(
    selection: str | None,
    session: dict[str, Any] | None,
) -> tuple[Any, dict[str, Any], str]:
    updated = dict(session or {})
    key = str(selection or "full")
    path = _stem_source_path(updated, key)
    if not path or not Path(path).exists():
        return (
            gr.update(),
            updated,
            _stem_status_html("error", message=f"The {DEMUCS_STEM_LABELS.get(key, key)} source is unavailable."),
        )
    updated["selected"] = key
    return (
        path,
        updated,
        _stem_status_html(
            "selected",
            audio_name=str(updated.get("original_name") or Path(path).name),
            selected=key,
        ),
    )


def _estimate_gpu_duration(
    audio: Any,
    instruments: list[str] | None = None,
    decode_preset: str = "fast",
    temperature: float = 1.0,
    *args: Any,
    **kwargs: Any,
) -> int:
    duration = _audio_duration(_audio_path(audio)) or 15.0
    plan = _transcription_runtime_plan(duration, decode_preset)
    return int(plan["requested_seconds"])


def _transcription_runtime_plan(
    duration: float,
    decode_preset: str | None,
) -> dict[str, Any]:
    """Choose the smallest supported batch that fits one ZeroGPU call."""

    preset = str(decode_preset or "fast")
    if preset not in TRANSCRIPTION_GPU_PROFILES:
        preset = "fast"
    safe_duration = max(0.0, float(duration or 0.0))
    chunks = max(1, math.ceil(safe_duration / 5.0))
    setup_seconds, sequential_rtf = TRANSCRIPTION_GPU_PROFILES.get(
        preset,
        TRANSCRIPTION_GPU_PROFILES["fast"],
    )
    max_batch_size = TRANSCRIPTION_MAX_BATCH_SIZE[preset]
    batch_size = 1
    speedup = 1.0
    sequential_compute_seconds = (
        safe_duration * sequential_rtf * TRANSCRIPTION_GPU_COMPUTE_FACTOR
    )
    raw_seconds = setup_seconds + sequential_compute_seconds
    estimate = math.ceil(
        raw_seconds * TRANSCRIPTION_RUNTIME_SCALE * TRANSCRIPTION_RESERVATION_SAFETY
    )
    if estimate > GPU_DURATION_CAP:
        for candidate in range(2, max_batch_size + 1):
            candidate_speedup = 1.0 + TRANSCRIPTION_BATCH_EFFICIENCY * (candidate - 1)
            candidate_raw_seconds = (
                setup_seconds + sequential_compute_seconds / candidate_speedup
            )
            candidate_estimate = math.ceil(
                candidate_raw_seconds
                * TRANSCRIPTION_RUNTIME_SCALE
                * TRANSCRIPTION_RESERVATION_SAFETY
            )
            batch_size = candidate
            speedup = candidate_speedup
            raw_seconds = candidate_raw_seconds
            estimate = candidate_estimate
            if candidate_estimate <= GPU_DURATION_CAP:
                break

    fits = estimate <= GPU_DURATION_CAP
    reservation_floor = (
        TRANSCRIPTION_BATCHED_MIN_RESERVATION_SECONDS
        if batch_size > 1
        else TRANSCRIPTION_MIN_RESERVATION_SECONDS
    )
    requested_seconds = (
        min(GPU_DURATION_CAP, max(reservation_floor, estimate))
        if fits
        else min(GPU_DURATION_CAP, ZEROGPU_MIN_CONFIGURED_CAP_SECONDS)
    )
    maximum_speedup = 1.0 + TRANSCRIPTION_BATCH_EFFICIENCY * (max_batch_size - 1)
    available_compute_seconds = max(
        0.0,
        GPU_DURATION_CAP / (TRANSCRIPTION_RUNTIME_SCALE * TRANSCRIPTION_RESERVATION_SAFETY)
        - setup_seconds,
    )
    max_supported_seconds = math.floor(
        available_compute_seconds
        * maximum_speedup
        / (sequential_rtf * TRANSCRIPTION_GPU_COMPUTE_FACTOR)
    )
    return {
        "preset": preset,
        "gpu_size": TRANSCRIPTION_GPU_SIZE,
        "compute_factor": TRANSCRIPTION_GPU_COMPUTE_FACTOR,
        "chunks": chunks,
        "batch_size": batch_size,
        "prelude_forcing": batch_size == 1,
        "speedup": speedup,
        "raw_seconds": raw_seconds,
        "estimated_seconds": estimate,
        "requested_seconds": requested_seconds,
        "fits": fits,
        "max_supported_seconds": max(0, max_supported_seconds),
    }


def _decode_settings(preset: str | None) -> tuple[bool, int]:
    return DECODE_PRESETS.get(str(preset or "fast"), DECODE_PRESETS["fast"])


def _accepts_keyword(function: Any, keyword: str) -> bool:
    """Check an installed API before passing a post-release keyword to it."""

    try:
        parameters = inspect.signature(function).parameters.values()
    except (TypeError, ValueError):
        return False
    return any(
        parameter.name == keyword
        or parameter.kind is inspect.Parameter.VAR_KEYWORD
        for parameter in parameters
    )


def _temperature_control(preset: str | None) -> dict[str, Any]:
    return gr.update(interactive=str(preset) == "creative")


def _resolve_instrument_lock(
    instruments: list[str] | None,
) -> list[str]:
    """Return the deduplicated hard constraint sent to MuScriptor."""

    return list(dict.fromkeys(instruments or []))


def _display_name(name: str) -> str:
    special = {
        "drums": "Drums",
        "voice": "Voice",
        "flutes": "Flutes",
        "electric_bass": "Electric bass",
        "acoustic_bass": "Acoustic bass",
        "acoustic_piano": "Acoustic piano",
        "electric_piano": "Electric piano",
        "soprano_and_alto_sax": "Soprano / alto saxophone",
        "tenor_sax": "Tenor saxophone",
    }
    return special.get(name, name.replace("_", " ").capitalize())


def _instrument_category_choices(
    names: list[str] | tuple[str, ...],
) -> list[tuple[str, list[tuple[str, str]]]]:
    """Return every available model instrument exactly once, grouped for the UI."""

    available = set(names)
    assigned: set[str] = set()
    categories: list[tuple[str, list[tuple[str, str]]]] = []
    for category, members in INSTRUMENT_CATEGORY_MEMBERS:
        choices = [(_display_name(name), name) for name in members if name in available]
        if choices:
            categories.append((category, choices))
            assigned.update(value for _, value in choices)
    other = [(_display_name(name), name) for name in names if name not in assigned]
    if other:
        categories.append(("Other model groups", other))
    return categories


def _color_for(name: str) -> str:
    digest = hashlib.sha1(name.encode("utf-8")).digest()
    return COLORS[int.from_bytes(digest[:2], "big") % len(COLORS)]


def _slug(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-") or "track"


def _data_uri(data: bytes, mime: str = "audio/midi") -> str:
    return f"data:{mime};base64,{base64.b64encode(data).decode('ascii')}"


def _program_for(name: str) -> int:
    if name == "drums" or MODEL is None:
        return 0
    try:
        return int(MODEL._program_for_instrument(name))
    except Exception:
        return 0


def _notes_from_events(events: list[Any]) -> dict[str, list[dict[str, float | int]]]:
    notes: dict[str, list[dict[str, float | int]]] = defaultdict(list)
    if NoteEndEvent is None:
        return {}
    for event in events:
        if not isinstance(event, NoteEndEvent):
            continue
        start = event.start_event
        notes[start.instrument].append(
            {
                "pitch": int(start.pitch),
                "start": round(float(start.start_time), 4),
                "end": round(max(float(event.end_time), float(start.start_time) + 0.03), 4),
                "velocity": 100,
            }
        )
    for values in notes.values():
        values.sort(key=lambda item: (item["start"], item["pitch"]))
    return dict(notes)


def _track_payloads(
    notes: dict[str, list[dict[str, float | int]]],
    requested: list[str] | None = None,
    midi_by_instrument: dict[str, bytes] | None = None,
) -> list[dict[str, Any]]:
    names = set(notes)
    names.update(requested or [])
    ordered = sorted(
        names,
        key=lambda name: (
            notes.get(name, [{}])[0].get("start", float("inf")) if notes.get(name) else float("inf"),
            name,
        ),
    )
    tracks: list[dict[str, Any]] = []
    for index, name in enumerate(ordered):
        track_notes = notes.get(name, [])
        midi = (midi_by_instrument or {}).get(name)
        tracks.append(
            {
                # A semantic id keeps mute/solo and synth channels stable while
                # new instruments appear during progressive transcription.
                "id": name,
                "key": name,
                "name": _display_name(name),
                "color": _color_for(name),
                "note_count": len(track_notes),
                "program": _program_for(name),
                "is_drum": name == "drums",
                "midi": _data_uri(midi) if midi else "",
                "notes": track_notes,
            }
        )
    return tracks


def _viewer_payload(
    *,
    state: str,
    status: str,
    progress: float,
    completed_windows: int | None = None,
    total_windows: int | None = None,
    audio_name: str,
    elapsed: float,
    duration: float,
    tracks: list[dict[str, Any]],
    full_midi: bytes | None = None,
    original_audio: str = "",
    available_until: float | None = None,
) -> dict[str, Any]:
    inferred_total_windows = math.ceil(duration / 5.0) if duration > 0 else 0
    window_total = max(
        0,
        int(total_windows if total_windows is not None else inferred_total_windows),
    )
    if completed_windows is None:
        if state in {"complete", "ready", "generating_score"}:
            window_completed = window_total
        elif available_until is not None:
            window_completed = math.floor(max(0.0, float(available_until)) / 5.0 + 1e-9)
        else:
            window_completed = math.floor(max(0.0, float(progress)) * window_total + 1e-9)
    else:
        window_completed = int(completed_windows)
    window_completed = max(0, min(window_total, window_completed))
    live_end = (
        duration
        if available_until is None and state in {"complete", "ready"}
        else float(available_until or 0)
    )
    value = initial_studio_value()
    value.update(
        {
            "state": state,
            "status": status,
            "progress": round(max(0.0, min(1.0, progress)), 4),
            "completed_windows": window_completed,
            "total_windows": window_total,
            "audio_name": audio_name,
            "export_stem": safe_filename_stem(Path(audio_name).stem, "muscriptor-transcription"),
            "elapsed": round(elapsed, 2),
            "duration": round(duration, 3),
            "available_until": round(max(0.0, min(duration, live_end)), 3),
            "original_audio": original_audio,
            "note_count": sum(track["note_count"] for track in tracks),
            "tracks": tracks,
            "full_midi": _data_uri(full_midi) if full_midi else "",
        }
    )
    return value


def _served_file_url(path: str | None, *, inline: bool = False) -> str:
    """Expose a temporary result through Gradio's authenticated file route.

    Gradio forces cache-created PDFs to download. Marking only our generated
    score files as static makes their response ``Content-Disposition: inline``,
    which lets the custom viewer embed the complete PDF safely.
    """
    if not path:
        return ""
    try:
        if inline:
            gr.set_static_paths(path)
        file_data = demo.serve_static_file(path)
    except (NameError, AttributeError, FileNotFoundError, ValueError):
        return ""
    if not isinstance(file_data, dict):
        return ""
    return str(file_data.get("url") or "")


def _served_audio_url(path: str | None) -> str:
    """Return the stable URL used by the custom original-audio player."""

    return _served_file_url(path)


def _event_for_instrument(event: Any, instrument: str) -> bool:
    if NoteStartEvent is not None and isinstance(event, NoteStartEvent):
        return event.instrument == instrument
    if NoteEndEvent is not None and isinstance(event, NoteEndEvent):
        return event.start_event.instrument == instrument
    return False


def _write_midi_outputs(
    events: list[Any],
    track_names: list[str],
    export_stem: str,
) -> tuple[Path, dict[str, Path], bytes, dict[str, bytes]]:
    if MODEL is None:
        raise RuntimeError(MODEL_ERROR or "Model not loaded.")
    output_dir = Path(tempfile.mkdtemp(prefix="muscriptor-midi-"))
    full_bytes = MODEL.events_to_midi_bytes(iter(events))
    full_path = output_dir / f"{export_stem}-full-transcription.mid"
    full_path.write_bytes(full_bytes)
    paths: dict[str, Path] = {}
    payloads: dict[str, bytes] = {}
    for name in track_names:
        track_events = [event for event in events if _event_for_instrument(event, name)]
        data = MODEL.events_to_midi_bytes(iter(track_events))
        path = output_dir / f"{export_stem}-{_slug(name)}.mid"
        path.write_bytes(data)
        paths[name] = path
        payloads[name] = data
    return full_path, paths, full_bytes, payloads


def _error_result(
    message: str,
    audio_name: str = "",
    duration: float = 0.0,
    original_audio: str | None = None,
):
    original_audio_url = _served_audio_url(original_audio)
    viewer = _viewer_payload(
        state="error",
        status=message,
        progress=0,
        audio_name=audio_name,
        elapsed=0,
        duration=duration,
        tracks=[],
        original_audio=original_audio_url,
    )
    return (
        original_audio,
        None,
        None,
        viewer,
        {"state": "error", "message": message, "viewer": viewer},
        gr.update(interactive=False),
    )


@spaces.GPU(size=TRANSCRIPTION_GPU_SIZE, duration=_estimate_gpu_duration)
def transcribe_audio(
    audio: Any,
    instruments: list[str] | None,
    decode_preset: str,
    temperature: float,
) -> Iterator[
    tuple[
        str | None,
        str | None,
        list[str] | None,
        dict[str, Any],
        dict[str, Any],
        dict[str, Any],
    ]
]:
    path = _audio_path(audio)
    if not path:
        yield _error_result("Upload an audio file before starting transcription.")
        return

    audio_name = Path(path).name
    duration = _audio_duration(path)
    if duration <= 0:
        yield _error_result("The audio duration could not be determined.", audio_name, original_audio=path)
        return
    if duration > MAX_AUDIO_SECONDS:
        yield _error_result(
            f"This Space accepts recordings up to {MAX_AUDIO_SECONDS / 60:.1f} minutes.",
            audio_name,
            duration,
            path,
        )
        return
    if MODEL is None:
        yield _error_result(
            MODEL_ERROR or "The MuScriptor model is unavailable.",
            audio_name,
            duration,
            path,
        )
        return

    original_audio_url = _served_audio_url(path)
    export_stem = safe_filename_stem(Path(audio_name).stem, "muscriptor-transcription")

    requested = _resolve_instrument_lock(instruments)
    use_sampling, beam = _decode_settings(decode_preset)
    runtime_plan = _transcription_runtime_plan(duration, decode_preset)
    if not runtime_plan["fits"]:
        supported = _short_duration(runtime_plan["max_supported_seconds"])
        yield _error_result(
            f"{str(runtime_plan['preset']).title()} decoding cannot fit this recording "
            f"inside a {GPU_DURATION_CAP}s ZeroGPU call. Choose Fast or Balanced, "
            f"or trim this preset to about {supported}.",
            audio_name,
            duration,
            path,
        )
        return
    batch_size = int(runtime_plan["batch_size"])
    planned_prelude_forcing = bool(runtime_plan["prelude_forcing"])
    supports_prelude_forcing = _accepts_keyword(MODEL.transcribe, "prelude_forcing")
    prelude_forcing = planned_prelude_forcing and supports_prelude_forcing
    if planned_prelude_forcing and not supports_prelude_forcing:
        print(
            "[MuScriptor Studio] Installed MuScriptor API has no prelude_forcing "
            "parameter; continuing in compatible sequential mode. Rebuild the Space "
            "to install the pinned upstream teacher-forcing commit.",
            file=sys.stderr,
            flush=True,
        )
    print(
        "[MuScriptor Studio] Transcription plan: "
        f"{duration:.1f}s audio · {runtime_plan['preset']} · "
        f"ZeroGPU {TRANSCRIPTION_GPU_SIZE} · batch {batch_size} · "
        f"request {runtime_plan['requested_seconds']}s.",
        flush=True,
    )
    events: list[Any] = []
    started = time.perf_counter()
    last_completed = -1
    last_total = math.ceil(duration / 5.0)

    try:
        transcribe_options: dict[str, Any] = {
            "instruments": requested or None,
            "use_sampling": use_sampling,
            "temperature": float(temperature),
            "cfg_coef": MODEL_CFG_COEF,
            "beam_size": beam,
            "batch_size": batch_size,
        }
        if supports_prelude_forcing:
            transcribe_options["prelude_forcing"] = planned_prelude_forcing
        stream = MODEL.transcribe(path, **transcribe_options)
        for event in stream:
            events.append(event)
            if ProgressEvent is None or not isinstance(event, ProgressEvent):
                continue
            if event.completed == last_completed:
                continue
            last_completed = event.completed
            last_total = max(1, int(event.total))
            notes = _notes_from_events(events)
            tracks = _track_payloads(notes, requested)
            fraction = event.completed / max(1, event.total)
            status = (
                (
                    f"ZeroGPU {TRANSCRIPTION_GPU_LABEL} ready · sequential continuity mode"
                    if batch_size == 1 and prelude_forcing
                    else (
                        f"ZeroGPU {TRANSCRIPTION_GPU_LABEL} ready · compatible sequential mode"
                        if batch_size == 1
                        else (
                            f"ZeroGPU {TRANSCRIPTION_GPU_LABEL} ready · "
                            f"quota-safe throughput mode (batch {batch_size})"
                        )
                    )
                )
                if event.completed == 0
                else f"Transcribed window {event.completed}/{event.total}"
            )
            viewer = _viewer_payload(
                state="transcribing",
                status=status,
                progress=fraction,
                completed_windows=event.completed,
                total_windows=event.total,
                audio_name=audio_name,
                elapsed=time.perf_counter() - started,
                duration=duration,
                tracks=tracks,
                original_audio=original_audio_url,
                available_until=min(duration, event.completed * 5.0),
            )
            yield (
                path,
                None,
                None,
                viewer,
                {"state": "running", "viewer": viewer},
                gr.update(interactive=False),
            )

        elapsed = time.perf_counter() - started
        notes = _notes_from_events(events)
        preliminary = _track_payloads(notes)
        track_names = [track["key"] for track in preliminary if track["note_count"]]
        full_path, track_paths, full_bytes, midi_by_instrument = _write_midi_outputs(
            events,
            track_names,
            export_stem,
        )
        tracks = _track_payloads(notes, midi_by_instrument=midi_by_instrument)
        for track in tracks:
            if track["key"] in track_paths:
                track["midi_name"] = track_paths[track["key"]].name
        viewer = _viewer_payload(
            state="complete",
            status=f"Transcription complete in {elapsed:.1f}s · score available on demand",
            progress=1,
            completed_windows=last_total,
            total_windows=last_total,
            audio_name=audio_name,
            elapsed=elapsed,
            duration=duration,
            tracks=tracks,
            full_midi=full_bytes,
            original_audio=original_audio_url,
            available_until=duration,
        )
        track_file_values = [str(track_paths[name]) for name in track_names]
        session = {
            "state": "complete",
            "title": Path(audio_name).stem,
            "audio_name": audio_name,
            "duration": duration,
            "elapsed": elapsed,
            "tracks": tracks,
            "viewer": viewer,
            "midi_files": [str(full_path), *track_file_values],
            "export_stem": export_stem,
            "decode_preset": decode_preset,
            "beam_size": beam,
            "use_sampling": use_sampling,
            "batch_size": batch_size,
            "prelude_forcing": prelude_forcing,
            "prelude_forcing_supported": supports_prelude_forcing,
            "gpu_size": TRANSCRIPTION_GPU_SIZE,
            "reserved_gpu_seconds": int(runtime_plan["requested_seconds"]),
            "instrument_lock_active": bool(requested),
            "locked_instruments": requested,
        }
        print(
            "[MuScriptor Studio] Transcription complete: "
            f"{duration:.1f}s audio · {elapsed:.1f}s inference · "
            f"{runtime_plan['requested_seconds']}s {TRANSCRIPTION_GPU_SIZE} requested · "
            f"batch {batch_size}.",
            flush=True,
        )
        yield (
            path,
            str(full_path),
            track_file_values,
            viewer,
            session,
            gr.update(interactive=False),
        )
    except Exception as exc:
        elapsed = time.perf_counter() - started
        message = f"Transcription stopped: {type(exc).__name__}: {exc}"
        print(f"[MuScriptor Studio] {message}", file=sys.stderr, flush=True)
        tracks = _track_payloads(_notes_from_events(events), requested)
        viewer = _viewer_payload(
            state="error",
            status=message,
            progress=max(0, last_completed) / max(1, last_total),
            completed_windows=max(0, last_completed),
            total_windows=last_total,
            audio_name=audio_name,
            elapsed=elapsed,
            duration=duration,
            tracks=tracks,
            original_audio=original_audio_url,
        )
        yield (
            path,
            None,
            None,
            viewer,
            {"state": "error", "message": message, "viewer": viewer},
            gr.update(interactive=False),
        )


def _score_analysis_html(analysis: dict[str, Any] | None = None, *, state: str = "waiting") -> str:
    if state == "waiting":
        return (
            '<div class="score-analysis" data-state="waiting"><i></i><div>'
            '<strong>Automatic notation presets</strong>'
            '<span>Tempo, meter, key, pickup and rhythm grid will be estimated from the decoded notes.</span>'
            '</div></div>'
        )
    if state == "unavailable" or not analysis:
        return (
            '<div class="score-analysis" data-state="unavailable"><i></i><div>'
            '<strong>Preset analysis unavailable</strong>'
            '<span>You can still enter tempo, meter and quantization manually.</span>'
            '</div></div>'
        )
    confidence = html.escape(str(analysis.get("confidence") or "low").title())
    pickup = float(analysis.get("pickup_beats") or 0.0)
    pickup_copy = "no pickup" if pickup <= 0 else f"{pickup:g}-beat pickup"
    key_signature = html.escape(str(analysis.get("key_signature") or "C major"))
    flags = len(analysis.get("review_flags") or [])
    return (
        '<div class="score-analysis" data-state="ready"><i></i><div>'
        f'<strong>Score prepared · {flags} review flag{"s" if flags != 1 else ""}</strong>'
        f'<span>{float(analysis["tempo_bpm"]):.1f} BPM · {html.escape(str(analysis["time_signature"]))} · '
        f'{key_signature} · {html.escape(str(analysis["quantization"]))} mixed grid · '
        f'{pickup_copy} · {confidence} confidence</span>'
        '</div></div>'
    )


def _instrument_lock_html(instruments: list[str] | None = None) -> str:
    selected = list(dict.fromkeys(instruments or []))
    total = len(INSTRUMENT_NAMES)
    if not selected:
        return (
            '<div class="instrument-lock-status" data-state="automatic">'
            '<span class="instrument-lock-icon" aria-hidden="true"></span><div>'
            '<strong>Automatic instrument discovery</strong>'
            f'<small>{total} groups available · add a lock when the lineup is known.</small>'
            '</div></div>'
        )
    count = len(selected)
    noun = "group" if count == 1 else "groups"
    return (
        '<div class="instrument-lock-status" data-state="locked">'
        '<span class="instrument-lock-icon" aria-hidden="true"></span><div>'
        f'<strong>{count} instrument {noun} locked</strong>'
        f'<small>{total} available · hard constraint across every five-second window.</small>'
        '</div></div>'
    )


def _merge_instrument_groups(
    *groups: list[str] | None,
) -> tuple[list[str], str, dict[str, Any]]:
    """Merge family selectors into MuScriptor's automatic hard lock."""

    selected: list[str] = []
    seen: set[str] = set()
    for group in groups:
        for name in group or []:
            if name in seen:
                continue
            selected.append(name)
            seen.add(name)
    return (
        selected,
        _instrument_lock_html(selected),
        gr.update(interactive=bool(selected)),
    )


def _clear_instrument_groups() -> tuple[Any, ...]:
    """Clear every visible family selector and return to automatic discovery."""

    empty_categories = tuple([] for _ in INSTRUMENT_CATEGORY_CHOICES)
    return (
        *empty_categories,
        [],
        _instrument_lock_html(),
        gr.update(interactive=False),
    )


def suggest_score_settings(session: dict[str, Any] | None) -> tuple[Any, ...]:
    session = session or {}
    if session.get("state") != "complete" or not session.get("tracks"):
        return (
            session,
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            _score_analysis_html(state="unavailable"),
            gr.update(interactive=False),
        )
    analysis = estimate_score_settings(session["tracks"])
    updated_session = dict(session)
    updated_session["score_analysis"] = analysis
    return (
        updated_session,
        analysis["tempo_bpm"],
        analysis["time_signature"],
        analysis["quantization"],
        analysis["key_signature"],
        analysis["pickup_beats"],
        _score_analysis_html(analysis),
        gr.update(interactive=True),
    )


def begin_notation(
    session: dict[str, Any] | None,
    current_viewer: dict[str, Any] | str | None = None,
) -> tuple[dict[str, Any], dict[str, Any]]:
    session = session or {}
    viewer = normalize_viewer_value(current_viewer or session.get("viewer"))
    if session.get("state") != "complete":
        return viewer, gr.update(interactive=False)
    updated = dict(viewer)
    updated.update(
        {
            "state": "generating_score",
            "status": "Engraving the full score and instrument parts…",
            "progress": 0,
        }
    )
    return updated, gr.update(interactive=False, value="Generating score…")


def reset_score_outputs(audio: Any, decode_preset: str = "fast") -> tuple[Any, ...]:
    """Clear the previous session immediately, before ZeroGPU enters its queue."""

    path = _audio_path(audio)
    if not path:
        raise gr.Error("Upload or record audio before starting transcription.")
    audio_name = Path(path).name if path else ""
    duration = _audio_duration(path)
    if duration <= 0:
        raise gr.Error("The audio duration could not be read.")
    if duration > MAX_AUDIO_SECONDS:
        raise gr.Error(
            f"This Space accepts audio up to {MAX_AUDIO_SECONDS / 60:.1f} minutes."
        )
    original_audio_url = _served_audio_url(path)
    runtime_plan = _transcription_runtime_plan(duration, decode_preset)
    if not runtime_plan["fits"]:
        supported = _short_duration(runtime_plan["max_supported_seconds"])
        raise gr.Error(
            f"{str(runtime_plan['preset']).title()} decoding cannot fit this recording "
            f"inside a {GPU_DURATION_CAP}s ZeroGPU call. Choose Fast or Balanced, "
            f"or trim this preset to about {supported}."
        )
    if runtime_plan["batch_size"] == 1:
        prelude_available = MODEL is not None and _accepts_keyword(
            MODEL.transcribe, "prelude_forcing"
        )
        mode = "sequential continuity" if prelude_available else "compatible sequential"
    else:
        mode = f"throughput batch {runtime_plan['batch_size']}"
    queue_status = (
        f"Request sent · {runtime_plan['preset'].title()} · "
        f"ZeroGPU {TRANSCRIPTION_GPU_LABEL} up to "
        f"{runtime_plan['requested_seconds']}s · {mode}"
    )
    viewer = _viewer_payload(
        state="queued",
        status=queue_status,
        progress=0.01,
        audio_name=audio_name,
        elapsed=0,
        duration=duration,
        tracks=[],
        original_audio=original_audio_url,
    )
    return (
        path,
        None,
        None,
        None,
        None,
        None,
        None,
        viewer,
        {},
        gr.update(interactive=False),
        _score_analysis_html(),
        120,
        "4/4",
        "1/16",
        "C major",
        0,
        "readable",
    )


def finalize_notation(
    session: dict[str, Any] | None,
    tempo_bpm: float,
    time_signature: str,
    quantization: str,
    key_signature: str,
    pickup_beats: float,
    cleanup_profile: str,
    show_solfege: bool,
) -> tuple[str | None, str | None, list[str] | None, str | None, dict[str, Any], dict[str, Any]]:
    session = session or {}
    viewer = normalize_viewer_value(session.get("viewer"))
    if session.get("state") != "complete":
        return None, None, None, None, viewer, gr.update(interactive=False)
    try:
        result = generate_notation(
            session["tracks"],
            title=session.get("title") or "Transcription MuScriptor",
            tempo_bpm=tempo_bpm,
            time_signature=time_signature,
            quantization=quantization,
            key_signature=key_signature,
            pickup_beats=pickup_beats,
            cleanup_profile=cleanup_profile,
            show_solfege=show_solfege,
            midi_files=session.get("midi_files"),
            timing_analysis=session.get("score_analysis"),
        )
        score_page_urls = [
            url
            for page in result.svg_pages
            if (url := _served_file_url(str(page), inline=True))
        ]
        musicxml_url = _served_file_url(str(result.musicxml))
        bundle_url = _served_file_url(str(result.bundle))
        report_url = _served_file_url(str(result.report))
        score_pdf_url = _served_file_url(str(result.pdf), inline=True) if result.pdf else ""
        score_parts = []
        for part in result.parts:
            part_pdf_url = _served_file_url(str(part.pdf), inline=True) if part.pdf else ""
            part_svg_urls = [
                url
                for page in part.svg_pages
                if (url := _served_file_url(str(page), inline=True))
            ]
            score_parts.append(
                {
                    "track_id": part.track_id,
                    "track_key": part.track_key,
                    "name": part.name,
                    "pdf": part_pdf_url,
                    "pdf_name": part.pdf.name if part.pdf else "",
                    "musicxml": _served_file_url(str(part.musicxml)),
                    "musicxml_name": part.musicxml.name,
                    "pages": part_svg_urls,
                    "page_count": len(part.svg_pages),
                    "render_error": part.render_error,
                }
            )
        part_files = [
            str(path)
            for part in result.parts
            for path in (part.pdf, part.musicxml)
            if path is not None
        ]
        updated = dict(viewer)
        if result.warnings:
            status = (
                f"MusicXML exports are ready · {len(result.warnings)} visual rendering "
                f"warning{'s' if len(result.warnings) != 1 else ''}"
            )
        else:
            part_count = len(result.parts)
            status = (
                f"Full score and {part_count} instrument "
                f"{'part' if part_count == 1 else 'parts'} are ready"
            )
        review_flags = list(result.diagnostics.get("review_flags") or [])
        engraved_settings = result.diagnostics.get("settings") or {}
        if review_flags:
            status += f" · {len(review_flags)} item{'s' if len(review_flags) != 1 else ''} to review"
        updated.update(
            {
                "state": "ready",
                "status": status,
                "score_svg": "" if score_page_urls else result.preview_svg,
                "score_pdf_url": score_pdf_url,
                "score_parts": score_parts,
                "score_page_urls": score_page_urls,
                "score_pages": len(result.svg_pages),
                "notation": {
                    "tempo": round(float(engraved_settings.get("tempo_bpm") or tempo_bpm), 1),
                    "time_signature": engraved_settings.get("time_signature") or time_signature,
                    "quantization": engraved_settings.get("quantization") or quantization,
                    "grid_family": "binary + ternary",
                    "key_signature": engraved_settings.get("key_signature") or key_signature,
                    "pickup_beats": round(float(engraved_settings.get("pickup_beats") or 0), 2),
                    "cleanup_profile": engraved_settings.get("cleanup_profile") or cleanup_profile,
                    "show_solfege": bool(show_solfege),
                    "part_count": len(result.parts),
                    "review_flags": review_flags,
                    "pdf": score_pdf_url,
                    "pdf_name": result.pdf.name if result.pdf else "",
                    "musicxml": musicxml_url,
                    "musicxml_name": result.musicxml.name,
                    "bundle_name": result.bundle.name,
                    "bundle": bundle_url,
                    "report": report_url,
                    "report_name": result.report.name,
                    "export_stem": result.export_stem,
                    "warnings": list(result.warnings),
                    "svg": (
                        score_page_urls[0]
                        if score_page_urls
                        else ""
                    ),
                },
            }
        )
        return (
            str(result.musicxml),
            str(result.pdf) if result.pdf else None,
            part_files,
            str(result.bundle),
            updated,
            gr.update(interactive=True, value="Regenerate full score + parts"),
        )
    except Exception as exc:
        updated = dict(viewer)
        updated["state"] = "notation_error"
        updated["status"] = f"MIDI is ready, but the score could not be generated: {type(exc).__name__}: {exc}"
        return (
            None,
            None,
            None,
            None,
            updated,
            gr.update(interactive=True, value="Retry score generation"),
        )


APP_CSS = (ROOT / "frontend" / "app.css").read_text(encoding="utf-8")

INSTRUMENT_CHOICES = [(_display_name(name), name) for name in INSTRUMENT_NAMES]
INSTRUMENT_CATEGORY_CHOICES = _instrument_category_choices(INSTRUMENT_NAMES)

TIME_SIGNATURE_CHOICES = ["2/2", "2/4", "3/4", "4/4", "5/4", "6/8", "7/8", "9/8", "12/8"]
KEY_SIGNATURE_CHOICES = [
    *(f"{tonic} major" for tonic in ("C♭", "G♭", "D♭", "A♭", "E♭", "B♭", "F", "C", "G", "D", "A", "E", "B", "F♯", "C♯")),
    *(f"{tonic} minor" for tonic in ("A♭", "E♭", "B♭", "F", "C", "G", "D", "A", "E", "B", "F♯", "C♯", "G♯", "D♯", "A♯")),
]

KYUTAI_URL = "https://kyutai.org/"
MIRELO_URL = "https://www.mirelo.ai/"
MODEL_PAGE_URL = f"https://huggingface.co/{MODEL_ID}"
DEMUCS_URL = "https://github.com/adefossez/demucs"

HEADER = f"""
<header class="workbench-nav">
  <a class="wordmark" href="https://github.com/muscriptor/muscriptor" target="_blank" rel="noreferrer">
    <span class="wordmark-glyph" aria-hidden="true"><i></i><i></i><i></i><i></i></span>
    <span><strong>MuScriptor</strong><small>Audio → multitrack MIDI</small></span>
  </a>
  <div class="runtime-status" aria-label="Runtime configuration">
    <span class="model-status"><i></i> Large · best quality</span>
    <span class="gpu-status">ZeroGPU {TRANSCRIPTION_GPU_LABEL}</span>
  </div>
</header>
"""

INTRO = f"""
<section class="workbench-intro">
  <div class="intro-statement">
    <span class="overline">Multi-instrument music transcription</span>
    <h1>Turn a recording into<br><em>playable MIDI tracks.</em></h1>
    <p>Transcribe notes by instrument, compare the result with the original, then export MIDI or engrave a score.</p>
  </div>
  <nav class="project-references" aria-label="MuScriptor research partners and model">
    <a href="{KYUTAI_URL}" target="_blank" rel="noreferrer" title="Visit Kyutai">
      <span>Research lab</span><strong>Kyutai</strong><i aria-hidden="true">↗</i>
    </a>
    <a href="{MIRELO_URL}" target="_blank" rel="noreferrer" title="Visit Mirelo">
      <span>Audio AI</span><strong>Mirelo</strong><i aria-hidden="true">↗</i>
    </a>
    <a href="{MODEL_PAGE_URL}" target="_blank" rel="noreferrer" title="Open the MuScriptor Large model card">
      <span>Model card</span><strong>Large · 1.3B</strong><i aria-hidden="true">↗</i>
    </a>
  </nav>
</section>
"""

WORKBENCH_HEADING = f"""
<header class="panel-heading">
  <div><span>New session</span><h2>Prepare the transcription</h2></div>
  <p>Upload, microphone or public YouTube video · up to {MAX_AUDIO_SECONDS / 60:.0f} minutes</p>
</header>
"""

RESULT_HEADING = """
<header class="panel-heading studio-heading">
  <div><span>Review & export</span><h2>Transcription studio</h2></div>
  <p>Original audio, playable MIDI tracks and score parts in one workspace</p>
</header>
"""

SOURCE_HEADING = """
<div class="control-heading"><span class="control-icon" aria-hidden="true"><svg viewBox="0 0 24 24"><path d="M4 14v-4m4 8V6m4 14V4m4 13V7m4 7v-4"/></svg></span><div><strong>Source audio</strong><small>Upload, record or import audio, then keep the full mix or load an isolated stem into this player.</small></div></div>
"""

YOUTUBE_IMPORT_HEADING = f"""
<div class="youtube-import-heading">
  <span class="youtube-import-icon" aria-hidden="true"><svg viewBox="0 0 24 24"><rect x="3.5" y="6" width="17" height="12" rx="3"/><path d="m10 9 5 3-5 3z"/></svg></span>
  <div><strong>Import from YouTube</strong><small>Best-effort import from one public video; direct upload remains available if YouTube challenges the Space.</small></div>
  <b>Video · {MAX_AUDIO_SECONDS / 60:.0f} min max</b>
</div>
"""

STEM_WORKBENCH_HEADING = f"""
<div class="stem-workbench-heading">
  <div><strong>Optional source isolation</strong><small>Keep the original mix or ask Demucs for four broad stems before transcription.</small></div>
  <a href="{DEMUCS_URL}" target="_blank" rel="noreferrer">Demucs v4 <span aria-hidden="true">↗</span></a>
</div>
"""

STEM_SOURCE_HEADING = """
<div class="stem-source-heading">
  <div><span>Player source</span><strong>What should MuScriptor listen to?</strong></div>
  <small>Choosing a source loads it into the audio player above.</small>
</div>
"""

DECODE_HEADING = """
<div class="control-heading"><span class="control-icon" aria-hidden="true"><svg viewBox="0 0 24 24"><path d="M4 7h10m4 0h2M4 17h2m4 0h10M14 4v6M6 14v6"/></svg></span><div><strong>Transcription controls</strong><small>Lock known instruments for consistency; Fast · Beam 1 is the recommended default.</small></div></div>
"""

INSTRUMENT_LOCK_HEADING = """
<div class="instrument-lock-heading">
  <div><strong>Instrument lock <span>Recommended when known</span></strong><small>Selecting one or more groups immediately restricts every five-second window. Leave all clear for automatic discovery.</small></div>
  <b>Optional</b>
</div>
"""

SCORE_HEADING = """
<div class="control-heading"><span class="control-icon score-icon" aria-hidden="true"><svg viewBox="0 0 24 24"><path d="M6 3h9l3 3v15H6zM9 10h6m-6 4h6m-6 4h4"/></svg></span><div><strong>Prepare score</strong><small>Review the suggested tempo, meter, key, pickup and rhythmic detail before engraving.</small></div></div>
"""

FOOTNOTE = f"""
<div class="footnote">
  <span>Ephemeral audio processing · no recording retention</span>
  <span>MuScriptor Large · ZeroGPU {TRANSCRIPTION_GPU_LABEL}</span>
</div>
"""

with gr.Blocks(title="MuScriptor Studio") as demo:
    with gr.Column(elem_id="app-shell"):
        gr.HTML(HEADER, elem_id="masthead")
        gr.HTML(INTRO, elem_id="intro")

        with gr.Column(elem_id="input-workbench"):
            gr.HTML(WORKBENCH_HEADING)
            with gr.Column(elem_id="source-panel"):
                gr.HTML(SOURCE_HEADING)
                audio_input = gr.Audio(
                    label="Upload or record audio",
                    sources=["upload", "microphone"],
                    type="filepath",
                    elem_id="audio-input",
                )
                stem_session_state = gr.State(value={})
                with gr.Column(elem_id="youtube-import-panel"):
                    gr.HTML(YOUTUBE_IMPORT_HEADING, elem_id="youtube-import-heading")
                    with gr.Row(elem_id="youtube-import-row"):
                        youtube_url_input = gr.Textbox(
                            label="YouTube video URL",
                            placeholder="https://www.youtube.com/watch?v=…",
                            lines=1,
                            max_lines=1,
                            show_label=False,
                            container=False,
                            elem_id="youtube-url-input",
                        )
                        youtube_import_button = gr.Button(
                            "Load audio",
                            variant="secondary",
                            elem_id="youtube-import-button",
                        )
                    youtube_import_status_output = gr.HTML(
                        value=_youtube_import_status_html(),
                        elem_id="youtube-import-status-output",
                    )
                with gr.Column(elem_id="stem-workbench"):
                    gr.HTML(STEM_WORKBENCH_HEADING, elem_id="stem-workbench-heading")
                    with gr.Row(elem_id="demucs-quality-row"):
                        gr.HTML(
                            '<div class="demucs-quality-copy"><strong>Separation quality</strong>'
                            '<small>Shift averaging · chunk overlap</small></div>',
                            elem_id="demucs-quality-heading",
                        )
                        demucs_preset_input = gr.Radio(
                            choices=[("Fast", "fast"), ("Balanced", "balanced"), ("Quality", "quality")],
                            value=DEMUCS_DEFAULT_PRESET,
                            label=None,
                            show_label=False,
                            container=False,
                            interactive=True,
                            elem_id="demucs-quality-preset",
                        )
                    with gr.Row(elem_id="stem-separation-row"):
                        stem_status_output = gr.HTML(
                            value=_stem_status_html(),
                            elem_id="stem-status-output",
                        )
                        separate_stems_button = gr.Button(
                            "Separate into stems",
                            variant="secondary",
                            interactive=False,
                            elem_id="separate-stems-button",
                        )
                    gr.HTML(STEM_SOURCE_HEADING, elem_id="stem-source-heading")
                    stem_source_input = gr.Radio(
                        choices=[(DEMUCS_STEM_LABELS["full"], "full")],
                        value="full",
                        label=None,
                        show_label=False,
                        container=False,
                        interactive=False,
                        elem_id="stem-source-selector",
                    )
                    # Retain every generated stem in Gradio's managed cache;
                    # user-facing selection happens through the source cards.
                    stem_files_output = gr.File(
                        label="Demucs stems",
                        file_count="multiple",
                        visible=False,
                    )

            with gr.Row(elem_id="settings-grid", equal_height=False):
                with gr.Column(scale=2, min_width=420, elem_id="decode-panel"):
                    gr.HTML(DECODE_HEADING)
                    with gr.Row(elem_id="transcription-controls-layout", equal_height=False):
                        with gr.Column(scale=3, min_width=300, elem_id="instrument-lock-panel"):
                            gr.HTML(INSTRUMENT_LOCK_HEADING, elem_id="instrument-lock-heading")
                            instrument_category_inputs: list[gr.CheckboxGroup] = []
                            with gr.Column(elem_id="instrument-categories"):
                                for category, choices in INSTRUMENT_CATEGORY_CHOICES:
                                    instrument_category_inputs.append(
                                        gr.CheckboxGroup(
                                            choices=choices,
                                            value=[],
                                            label=category,
                                            interactive=True,
                                            elem_classes=["instrument-category"],
                                            elem_id=f"instrument-category-{_slug(category)}",
                                        )
                                    )

                        with gr.Column(scale=2, min_width=220, elem_id="decode-settings-panel"):
                            decode_preset_input = gr.Radio(
                                choices=[
                                    ("Fast · Beam 1", "fast"),
                                    ("Balanced · Beam 2", "balanced"),
                                    ("High quality · Beam 4", "quality"),
                                    ("Creative · Sampling", "creative"),
                                ],
                                value="fast",
                                label="Decode preset",
                                info="Beam presets are deterministic. Creative uses sampling with Beam 1.",
                                interactive=True,
                                elem_id="decode-preset",
                            )
                            with gr.Column(elem_id="decoding-options"):
                                temperature_input = gr.Slider(
                                    0.2,
                                    1.4,
                                    value=1.0,
                                    step=0.1,
                                    label="Temperature",
                                    info="Only used when creative decoding is enabled.",
                                    interactive=False,
                                )

                    # Consolidate every visible family into the model's single
                    # lock, then show its state as a footer for both columns.
                    instrument_input = gr.CheckboxGroup(
                        choices=INSTRUMENT_CHOICES,
                        value=[],
                        label="Selected instrument groups",
                        interactive=True,
                        visible=False,
                        elem_id="instrument-lock",
                    )
                    with gr.Row(elem_id="instrument-lock-footer"):
                        instrument_lock_output = gr.HTML(
                            value=_instrument_lock_html(),
                            elem_id="instrument-lock-output",
                        )
                        clear_instrument_lock_button = gr.Button(
                            "Clear lock",
                            variant="secondary",
                            interactive=False,
                            elem_id="clear-instrument-lock-button",
                        )

                with gr.Column(scale=1, min_width=280, elem_id="notation-panel"):
                    gr.HTML(SCORE_HEADING)
                    with gr.Row(elem_id="score-primary-settings"):
                        tempo_input = gr.Number(value=120, minimum=20, maximum=300, label="Tempo (BPM)")
                        time_signature_input = gr.Dropdown(
                            TIME_SIGNATURE_CHOICES,
                            value="4/4",
                            label="Meter",
                            allow_custom_value=True,
                        )
                    with gr.Row(elem_id="score-musical-settings"):
                        key_signature_input = gr.Dropdown(
                            KEY_SIGNATURE_CHOICES,
                            value="C major",
                            label="Concert key",
                            allow_custom_value=True,
                            elem_id="key-signature-input",
                        )
                        pickup_input = gr.Number(
                            value=0,
                            minimum=0,
                            maximum=4.75,
                            step=0.25,
                            label="Pickup (beats)",
                            elem_id="pickup-input",
                        )
                    quantization_input = gr.Radio(
                        [("1/8", "1/8"), ("1/16", "1/16"), ("1/32", "1/32")],
                        value="1/16",
                        label="Rhythm detail",
                        info="Each level accepts both straight notes and triplets.",
                        interactive=True,
                        elem_id="score-quantization",
                    )
                    # A single musician-oriented cleanup pipeline keeps this
                    # technical implementation choice out of the UI. The
                    # value remains stateful so existing event wiring and API
                    # manifests stay explicit and reproducible.
                    cleanup_profile_input = gr.State(value="readable")
                    solfege_input = gr.Checkbox(
                        label="Show solfège names (Do, Ré, Mi)",
                        value=False,
                        interactive=True,
                    )
                    score_analysis_output = gr.HTML(
                        value=_score_analysis_html(),
                        elem_id="score-analysis-output",
                    )
                    generate_score_button = gr.Button(
                        "Generate full score + parts",
                        variant="secondary",
                        interactive=False,
                        elem_id="generate-score-button",
                    )

            with gr.Row(elem_id="run-row"):
                transcribe_button = gr.Button(
                    "Transcribe audio",
                    variant="primary",
                    size="lg",
                    elem_id="transcribe-button",
                )
                gr.HTML(
                    f'<p class="cold-note"><span></span><strong>ZeroGPU {TRANSCRIPTION_GPU_LABEL} on demand</strong> · the first run may briefly queue.</p>'
                )

        with gr.Column(elem_id="transcription-panel"):
            gr.HTML(RESULT_HEADING)
            session_state = gr.State(value={})
            viewer = StudioViewer(elem_id="studio-viewer")
            original_audio_output = gr.Audio(
                label="Original audio playback",
                interactive=False,
                visible="hidden",
                elem_id="original-audio-output",
            )
            # Keep file components as private event outputs so Gradio retains
            # and serves generated artifacts. All user-facing downloads live
            # in the studio toolbar and per-track / per-score controls.
            full_midi_output = gr.File(label="Full MIDI", visible=False)
            track_midi_output = gr.File(label="MIDI by instrument", file_count="multiple", visible=False)
            musicxml_output = gr.File(label="Full score MusicXML", visible=False)
            pdf_output = gr.File(label="Full score PDF", visible=False)
            part_score_output = gr.File(label="Scores by instrument", file_count="multiple", visible=False)
            bundle_output = gr.File(label="All exports (.zip)", visible=False)

        gr.HTML(FOOTNOTE)

    decode_preset_input.change(
        _temperature_control,
        inputs=[decode_preset_input],
        outputs=[temperature_input],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
    )

    youtube_import_outputs = [
        audio_input,
        stem_session_state,
        stem_source_input,
        stem_status_output,
        separate_stems_button,
        stem_files_output,
        youtube_import_status_output,
        youtube_import_button,
    ]
    youtube_preflight = youtube_import_button.click(
        begin_youtube_import,
        inputs=[youtube_url_input],
        outputs=[youtube_import_status_output, youtube_import_button],
        queue=False,
        show_progress="hidden",
    )
    youtube_preflight.then(
        import_youtube_audio,
        inputs=[youtube_url_input],
        outputs=youtube_import_outputs,
        api_name="import_youtube",
        queue=False,
        show_progress="hidden",
    )
    youtube_submit_preflight = youtube_url_input.submit(
        begin_youtube_import,
        inputs=[youtube_url_input],
        outputs=[youtube_import_status_output, youtube_import_button],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
    )
    youtube_submit_preflight.then(
        import_youtube_audio,
        inputs=[youtube_url_input],
        outputs=youtube_import_outputs,
        queue=False,
        show_progress="hidden",
        api_visibility="private",
    )

    audio_input.input(
        register_source_audio,
        inputs=[audio_input],
        outputs=[
            stem_session_state,
            stem_source_input,
            stem_status_output,
            separate_stems_button,
            stem_files_output,
        ],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
        trigger_mode="always_last",
    )
    audio_input.input(
        reset_youtube_import_status,
        inputs=[],
        outputs=[youtube_import_status_output],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
        trigger_mode="always_last",
    )
    audio_input.clear(
        clear_source_audio,
        inputs=[],
        outputs=[
            stem_session_state,
            stem_source_input,
            stem_status_output,
            separate_stems_button,
            stem_files_output,
        ],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
    )
    audio_input.clear(
        reset_youtube_import_status,
        inputs=[],
        outputs=[youtube_import_status_output],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
    )
    stem_preflight = separate_stems_button.click(
        begin_stem_separation,
        inputs=[audio_input, stem_session_state, demucs_preset_input],
        outputs=[stem_status_output, separate_stems_button, stem_source_input],
        queue=False,
        show_progress="hidden",
    )
    stem_preflight.then(
        separate_audio_stems,
        inputs=[audio_input, stem_session_state, demucs_preset_input],
        outputs=[
            stem_session_state,
            stem_source_input,
            stem_status_output,
            separate_stems_button,
            stem_files_output,
        ],
        api_name="separate_stems",
        queue=True,
        concurrency_limit=1,
        concurrency_id="muscriptor-gpu",
        show_progress="hidden",
    )
    stem_source_input.input(
        select_stem_source,
        inputs=[stem_source_input, stem_session_state],
        outputs=[audio_input, stem_session_state, stem_status_output],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
    )

    for category_input in instrument_category_inputs:
        category_input.input(
            _merge_instrument_groups,
            inputs=instrument_category_inputs,
            outputs=[instrument_input, instrument_lock_output, clear_instrument_lock_button],
            queue=False,
            show_progress="hidden",
            api_visibility="private",
        )
    clear_instrument_lock_button.click(
        _clear_instrument_groups,
        inputs=[],
        outputs=[
            *instrument_category_inputs,
            instrument_input,
            instrument_lock_output,
            clear_instrument_lock_button,
        ],
        queue=False,
        show_progress="hidden",
        api_visibility="private",
    )

    preflight = transcribe_button.click(
        reset_score_outputs,
        inputs=[audio_input, decode_preset_input],
        outputs=[
            original_audio_output,
            full_midi_output,
            track_midi_output,
            musicxml_output,
            pdf_output,
            part_score_output,
            bundle_output,
            viewer,
            session_state,
            generate_score_button,
            score_analysis_output,
            tempo_input,
            time_signature_input,
            quantization_input,
            key_signature_input,
            pickup_input,
            cleanup_profile_input,
        ],
        queue=False,
        show_progress="hidden",
    )
    transcription = preflight.then(
        transcribe_audio,
        inputs=[
            audio_input,
            instrument_input,
            decode_preset_input,
            temperature_input,
        ],
        outputs=[
            original_audio_output,
            full_midi_output,
            track_midi_output,
            viewer,
            session_state,
            generate_score_button,
        ],
        api_name="transcribe",
        queue=True,
        concurrency_limit=1,
        concurrency_id="muscriptor-gpu",
        show_progress="full",
    )
    transcription.then(
        suggest_score_settings,
        inputs=[session_state],
        outputs=[
            session_state,
            tempo_input,
            time_signature_input,
            quantization_input,
            key_signature_input,
            pickup_input,
            score_analysis_output,
            generate_score_button,
        ],
        queue=False,
        show_progress="hidden",
        api_name="analyze_score_settings",
    )
    notation_preflight = generate_score_button.click(
        begin_notation,
        inputs=[session_state, viewer],
        outputs=[viewer, generate_score_button],
        queue=False,
        show_progress="hidden",
    )
    notation_preflight.then(
        finalize_notation,
        inputs=[
            session_state,
            tempo_input,
            time_signature_input,
            quantization_input,
            key_signature_input,
            pickup_input,
            cleanup_profile_input,
            solfege_input,
        ],
        outputs=[musicxml_output, pdf_output, part_score_output, bundle_output, viewer, generate_score_button],
        api_name="generate_score",
        queue=False,
        show_progress="hidden",
    )

demo.queue(default_concurrency_limit=1, max_size=8)


if __name__ == "__main__":
    demo.launch(
        css=APP_CSS,
        theme=gr.themes.Base(primary_hue="teal", neutral_hue="zinc"),
        js="""() => {
          document.documentElement.classList.add('dark');
          document.body.classList.add('dark');
          document.body.style.background = '#090b0f';
          try { localStorage.setItem('theme', 'dark'); } catch (_) {}
        }""",
        footer_links=[],
    )