add beam search quality control

README.md

@@ -22,4 +22,4 @@ The checkpoint is gated under CC BY-NC 4.0. This Space requires an `HF_TOKEN` se
 
 ## API
 
-The Gradio endpoint is `/transcribe`. It accepts an audio file, optional exact instrument group names, a sampling toggle, and temperature. It returns the combined MIDI, all isolated track files, the browser visualization payload, and a machine-readable manifest.
+The Gradio endpoint is `/transcribe`. It accepts an audio file, optional exact instrument group names, a sampling toggle, temperature, and a beam width from 1–4. It returns the combined MIDI, all isolated track files, the browser visualization payload, and a machine-readable manifest.

app.py

@@ -140,12 +140,24 @@ def _audio_duration(path: str | None) -> float | None:
         return None
 
 
-def _estimate_gpu_duration(audio: Any, *args: Any, **kwargs: Any) -> int:
+def _estimate_gpu_duration(
+    audio: Any,
+    instruments: list[str] | None = None,
+    use_sampling: bool = False,
+    temperature: float = 1.0,
+    beam_size: int = 1,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
     """Estimate allocation from 5-second chunks; accepts Gradio extras."""
 
     duration = _audio_duration(_audio_path(audio))
     chunks = max(1, math.ceil((duration or 15.0) / 5.0))
-    return min(GPU_DURATION_CAP, GPU_BASE_SECONDS + chunks * GPU_SECONDS_PER_CHUNK)
+    beam = max(1, min(4, int(beam_size or 1)))
+    return min(
+        GPU_DURATION_CAP,
+        GPU_BASE_SECONDS + chunks * GPU_SECONDS_PER_CHUNK * beam,
+    )
 
 
 def _display_name(name: str) -> str:
@@ -348,6 +360,7 @@ def transcribe_audio(
     instruments: list[str] | None,
     use_sampling: bool,
     temperature: float,
+    beam_size: int,
 ) -> Iterator[tuple[str | None, list[str] | None, str, dict[str, Any]]]:
     """Stream chunk progress, then return combined and isolated MIDI files."""
 
@@ -407,6 +420,7 @@ def transcribe_audio(
         return
 
     requested = list(dict.fromkeys(instruments or []))
+    beam = max(1, min(4, int(beam_size or 1)))
     started = time.perf_counter()
     events: list[NoteStartEvent | NoteEndEvent | ProgressEvent] = []
     last_completed = -1
@@ -417,6 +431,7 @@ def transcribe_audio(
             instruments=requested or None,
             use_sampling=bool(use_sampling),
             temperature=float(temperature),
+            beam_size=beam,
             batch_size=1,
         )
         for event in stream:
@@ -598,18 +613,21 @@ EXAMPLE_ROWS = [
         ["violin", "acoustic_piano"],
         False,
         1.0,
+        1,
     ],
     [
         str(EXAMPLE_ROOT / "livery-stable-blues-jazz-band.mp3"),
         ["trumpet", "clarinet", "trombone", "drums"],
         False,
         1.0,
+        1,
     ],
     [
         str(EXAMPLE_ROOT / "double-tracked-pop-rock.mp3"),
         [],
         False,
         1.0,
+        1,
     ],
 ]
 
@@ -663,6 +681,17 @@ with gr.Blocks(title="MuScriptor — Audio to MIDI") as demo:
                         step=0.1,
                         label="Temperature",
                     )
+                    beam_size_input = gr.Slider(
+                        minimum=1,
+                        maximum=4,
+                        value=1,
+                        step=1,
+                        label="Beam width",
+                        info=(
+                            "1 is faster. Widths 2–4 use slower deterministic "
+                            "beam search; sampling and temperature apply only at 1."
+                        ),
+                    )
                 transcribe_button = gr.Button(
                     "Transcribe audio",
                     variant="primary",
@@ -700,6 +729,7 @@ with gr.Blocks(title="MuScriptor — Audio to MIDI") as demo:
                     instrument_input,
                     use_sampling_input,
                     temperature_input,
+                    beam_size_input,
                 ],
                 outputs=[
                     full_midi_output,
@@ -729,6 +759,7 @@ with gr.Blocks(title="MuScriptor — Audio to MIDI") as demo:
             instrument_input,
             use_sampling_input,
             temperature_input,
+            beam_size_input,
         ],
         outputs=[
             full_midi_output,