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'
'
f'
{heading}{detail}
'
)
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'
'
f"{title}{detail}"
f'
{html.escape(DEMUCS_MODEL_NAME)} · {preset_label}
'
)
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 (
'
'
'Automatic notation presets'
'Tempo, meter, key, pickup and rhythm grid will be estimated from the decoded notes.'
'
'
)
if state == "unavailable" or not analysis:
return (
'
'
'Preset analysis unavailable'
'You can still enter tempo, meter and quantization manually.'
'
'
)
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 (
'
Choosing a source loads it into the audio player above.
"""
DECODE_HEADING = """
Transcription controlsLock known instruments for consistency; Fast · Beam 1 is the recommended default.
"""
INSTRUMENT_LOCK_HEADING = """
Instrument lock Recommended when knownSelecting one or more groups immediately restricts every five-second window. Leave all clear for automatic discovery.
Optional
"""
SCORE_HEADING = """
Prepare scoreReview the suggested tempo, meter, key, pickup and rhythmic detail before engraving.
"""
FOOTNOTE = f"""
Ephemeral audio processing · no recording retentionMuScriptor Large · ZeroGPU {TRANSCRIPTION_GPU_LABEL}
"""
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(
'