Upload 8 files

app.py

@@ -57,6 +57,23 @@ GPU_BASE_SECONDS = 35
 GPU_SECONDS_PER_CHUNK = 9
 GPU_DURATION_CAP = int(os.environ.get("MAX_GPU_SECONDS", "600"))
 
+# 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_SHIFTS = 1
+DEMUCS_OVERLAP = 0.25
+DEMUCS_GPU_DURATION_CAP = int(os.environ.get("DEMUCS_MAX_GPU_SECONDS", "600"))
+DEMUCS_STEM_ORDER = ("vocals", "drums", "bass", "other")
+DEMUCS_STEM_LABELS = {
+    "full": "Full mix",
+    "vocals": "Vocals",
+    "drums": "Drums",
+    "bass": "Bass",
+    "other": "Other",
+}
+
 COLORS = (
     "#35c8ff",
     "#f6a623",
@@ -144,6 +161,8 @@ INSTRUMENT_CATEGORY_MEMBERS: tuple[tuple[str, tuple[str, ...]], ...] = (
 
 MODEL: Any | None = None
 MODEL_ERROR = ""
+DEMUCS_SEPARATOR: Any | None = None
+DEMUCS_ERROR = ""
 NoteStartEvent: Any = None
 NoteEndEvent: Any = None
 ProgressEvent: Any = None
@@ -217,6 +236,42 @@ if os.environ.get("MUSCRIPTOR_SKIP_MODEL_LOAD") != "1":
         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
@@ -240,6 +295,325 @@ def _audio_duration(path: str | None) -> float:
         return 0.0
 
 
+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 = "",
+) -> str:
+    safe_name = html.escape(audio_name or "the source audio")
+    label = html.escape(DEMUCS_STEM_LABELS.get(selected, _display_name(selected)))
+    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…",
+            "ZeroGPU is isolating vocals, drums, bass and the remaining accompaniment.",
+        ),
+        "complete": (
+            "Four stems ready",
+            "Full mix is still loaded. Choose a source below before starting 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"])
+    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)} · 4 stems</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,
+) -> tuple[str, dict[str, Any], dict[str, Any]]:
+    path = _stem_source_path(session, "full") or _audio_path(audio)
+    if not path:
+        return (
+            _stem_status_html("error", message="Add or record audio before running Demucs."),
+            gr.update(interactive=False, value="Separate into stems"),
+            _stem_selector_update(session, interactive=False),
+        )
+    return (
+        _stem_status_html("working", audio_name=Path(path).name),
+        gr.update(interactive=False, value="Separating stems…"),
+        _stem_selector_update(session, interactive=False),
+    )
+
+
+def _estimate_demucs_gpu_duration(
+    audio: Any,
+    session: dict[str, Any] | None = None,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
+    path = _stem_source_path(session, "full") or _audio_path(audio)
+    duration = _audio_duration(path) or 15.0
+    return min(DEMUCS_GPU_DURATION_CAP, max(60, math.ceil(35 + duration * 0.85)))
+
+
+@spaces.GPU(duration=_estimate_demucs_gpu_duration)
+def separate_audio_stems(
+    audio: Any,
+    session: dict[str, Any] | None,
+    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),
+        )
+
+    started = time.perf_counter()
+    separator = DEMUCS_SEPARATOR
+    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, DEMUCS_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(callback=on_progress, callback_arg={"audio_name": audio_name})
+        _, 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,
+            "duration": duration,
+            "elapsed": time.perf_counter() - started,
+        }
+        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),
+            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(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,
@@ -931,6 +1305,7 @@ INSTRUMENT_CATEGORY_CHOICES = _instrument_category_choices(INSTRUMENT_NAMES)
 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">
@@ -969,7 +1344,7 @@ INTRO = f"""
 WORKBENCH_HEADING = f"""
 <header class="panel-heading">
   <div><span>New session</span><h2>Prepare the transcription</h2></div>
-  <p>WAV, MP3, FLAC or OGG · up to {MAX_AUDIO_SECONDS / 60:.0f} minutes</p>
+  <p>Upload or microphone · WAV, MP3, FLAC or OGG · up to {MAX_AUDIO_SECONDS / 60:.0f} minutes</p>
 </header>
 """
 
@@ -981,7 +1356,14 @@ RESULT_HEADING = """
 """
 
 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 once; the same Gradio player remains available for A/B comparison.</small></div></div>
+<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 or record audio, then keep the full mix or load an isolated stem into this player.</small></div></div>
+"""
+
+STEM_WORKBENCH_HEADING = f"""
+<div class="stem-workbench-heading">
+  <div><strong>Prepare the transcription source</strong><small>Optional separation into four broad audio stems before MuScriptor.</small></div>
+  <a href="{DEMUCS_URL}" target="_blank" rel="noreferrer">Demucs v4 <span aria-hidden="true">↗</span></a>
+</div>
 """
 
 DECODE_HEADING = """
@@ -1016,11 +1398,40 @@ with gr.Blocks(title="MuScriptor Studio") as demo:
             with gr.Column(elem_id="source-panel"):
                 gr.HTML(SOURCE_HEADING)
                 audio_input = gr.Audio(
-                    label="Drop audio here",
-                    sources=["upload"],
+                    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="stem-workbench"):
+                    gr.HTML(STEM_WORKBENCH_HEADING, elem_id="stem-workbench-heading")
+                    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",
+                        )
+                    stem_source_input = gr.Radio(
+                        choices=[(DEMUCS_STEM_LABELS["full"], "full")],
+                        value="full",
+                        label="Source to transcribe",
+                        info="Selecting a source loads it into the Gradio player above. Then press Transcribe audio.",
+                        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"):
@@ -1152,6 +1563,63 @@ with gr.Blocks(title="MuScriptor Studio") as demo:
         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,
+        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,
+        api_visibility="private",
+    )
+    stem_preflight = separate_stems_button.click(
+        begin_stem_separation,
+        inputs=[audio_input, stem_session_state],
+        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],
+        outputs=[
+            stem_session_state,
+            stem_source_input,
+            stem_status_output,
+            separate_stems_button,
+            stem_files_output,
+        ],
+        api_name="separate_stems",
+        concurrency_limit=1,
+        concurrency_id="muscriptor-gpu",
+        show_progress="full",
+    )
+    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,
+        api_visibility="private",
+    )
+
     for category_input in instrument_category_inputs:
         category_input.change(
             _merge_instrument_groups,
@@ -1198,7 +1666,7 @@ with gr.Blocks(title="MuScriptor Studio") as demo:
         ],
         api_name="transcribe",
         concurrency_limit=1,
-        concurrency_id="muscriptor-large",
+        concurrency_id="muscriptor-gpu",
         show_progress="full",
     )
     transcription.then(

requirements.txt

@@ -1,6 +1,7 @@
 gradio==6.20.0
 spaces==0.51.0
 muscriptor==0.2.1
+demucs==4.1.0
 numpy==1.26.4
 einops==0.8.2
 mido==1.3.3