Upload folder using huggingface_hub

app.py

@@ -67,6 +67,7 @@ import partitura
 
 from drum_dynamics.serve.models import Engine
 from drum_dynamics.viz.playback import set_soundfont, _ensure_fluidsynth_discoverable
+from drum_dynamics.core.midi import drum_name
 
 # --- load once at module scope ------------------------------------------------
 ENGINE = Engine.load(os.path.join(HERE, "models"))
@@ -97,8 +98,14 @@ def _load_notes(midi_path: str):
     return perf, pp
 
 
-def _render_audio(pp) -> tuple[int, np.ndarray]:
-    """Render a performed part's drum notes to a mono int16 waveform."""
+def _note_events(pp):
+    """Snapshot the current notes as (on_sec, off_sec, pitch, velocity) tuples."""
+    return [(float(nd["note_on"]), float(nd["note_off"]),
+             int(nd["midi_pitch"]), int(nd["velocity"])) for nd in pp.notes]
+
+
+def _render_events(events) -> tuple[int, np.ndarray]:
+    """Render (on, off, pitch, velocity) drum events to a mono int16 waveform."""
     _ensure_fluidsynth_discoverable()  # no-op off macOS; helps on local dev
     import fluidsynth  # lazy: keeps the app importable without the native lib
 
@@ -106,14 +113,14 @@ def _render_audio(pp) -> tuple[int, np.ndarray]:
     sfid = fl.sfload(SF2)
     fl.program_select(9, sfid, 128, 0)  # GM percussion set on channel 9
 
-    events = []  # (time_sec, is_on, pitch, velocity)
-    for nd in pp.notes:
-        events.append((float(nd["note_on"]), True, int(nd["midi_pitch"]), int(nd["velocity"])))
-        events.append((float(nd["note_off"]), False, int(nd["midi_pitch"]), 0))
-    events.sort(key=lambda e: e[0])
+    timeline = []  # (time_sec, is_on, pitch, velocity)
+    for on, off, pitch, vel in events:
+        timeline.append((on, True, pitch, vel))
+        timeline.append((off, False, pitch, 0))
+    timeline.sort(key=lambda e: e[0])
 
     audio = []
-    for (t0, is_on, pitch, vel), (t1, *_rest) in zip(events[:-1], events[1:]):
+    for (t0, is_on, pitch, vel), (t1, *_rest) in zip(timeline[:-1], timeline[1:]):
         if is_on:
             fl.noteon(9, pitch, vel)
         else:
@@ -121,24 +128,47 @@ def _render_audio(pp) -> tuple[int, np.ndarray]:
         n = int(max(0, (t1 - t0)) * SR)
         if n:
             audio.extend(fl.get_samples(n)[::2])  # interleaved stereo -> mono
-    # let the last hits ring out
-    audio.extend(fl.get_samples(SR)[::2])
+    audio.extend(fl.get_samples(SR)[::2])  # let the last hits ring out
     fl.delete()
     return SR, np.asarray(audio, dtype=np.int16)
 
 
-def _velocity_plot(before: list[int], after: list[int]):
-    """Stem plot of velocity per note, original vs predicted."""
-    fig, ax = plt.subplots(figsize=(11, 3.2))
-    x = np.arange(len(before))
-    ax.plot(x, before, ".", ms=4, alpha=0.5, label="original", color="#9aa0a6")
-    ax.plot(x, after, ".", ms=5, alpha=0.9, label="predicted", color="#4f7cff")
-    ax.set_xlabel("note (in time order of input)")
-    ax.set_ylabel("velocity")
-    ax.set_ylim(0, 128)
-    ax.legend(loc="upper right", framealpha=0.9)
-    ax.grid(True, alpha=0.2)
-    fig.tight_layout()
+def _drumroll_figure(na_before, na_after):
+    """Stacked drum-rolls (velocity as colour) for original vs humanized.
+
+    Only the drum pieces actually present are shown, on a shared 0-127 colour
+    scale so the change in dynamics is directly comparable between the two.
+    """
+    pitches = sorted({int(p) for p in na_before["pitch"]}, reverse=True)  # high on top
+    row_of = {p: i for i, p in enumerate(pitches)}
+    n_ticks = int(max(na_before["onset_tick"] + na_before["duration_tick"])) + 1
+    min_w = max(1, n_ticks // 200)  # widen hits so short drums stay visible
+
+    def build(na):
+        roll = np.zeros((len(pitches), n_ticks))
+        for pt, on, dur, vel in zip(na["pitch"], na["onset_tick"],
+                                    na["duration_tick"], na["velocity"]):
+            r = row_of[int(pt)]
+            s = int(on)
+            e = min(n_ticks, s + max(int(dur), min_w))
+            roll[r, s:e] = np.maximum(roll[r, s:e], vel / 127.0)
+        return roll
+
+    fig, axes = plt.subplots(2, 1, figsize=(11, 6.2), sharex=True)
+    im = None
+    for ax, na, title in ((axes[0], na_before, "Original"),
+                          (axes[1], na_after, "Humanized")):
+        im = ax.imshow(build(na), aspect="auto", origin="lower", interpolation="none",
+                       cmap="magma", vmin=0.0, vmax=1.0)
+        ax.set_title(title, fontsize=10, loc="left")
+        ax.set_yticks(range(len(pitches)))
+        ax.set_yticklabels([drum_name(p) for p in pitches], fontsize=7)
+        ax.set_ylabel("drum piece", fontsize=8)
+    axes[1].set_xlabel("time (MIDI ticks)")
+    cbar = fig.colorbar(im, ax=axes, fraction=0.03, pad=0.02)
+    cbar.set_label("velocity")
+    cbar.set_ticks(np.linspace(0, 1, 8))
+    cbar.set_ticklabels(np.linspace(0, 127, 8, dtype=int))
     return fig
 
 
@@ -153,7 +183,8 @@ def humanize(midi_file, model_label, style, bpm, time_signature, beat_type,
     if not pp.notes:
         raise gr.Error("No notes found in that MIDI file.")
 
-    before = [int(nd["velocity"]) for nd in pp.notes]
+    na_before = pp.note_array().copy()          # original, for the roll
+    events_before = _note_events(pp)            # original, for the audio
     notes = [
         {"index": i, "onset_sec": float(nd["note_on"]), "pitch": int(nd["midi_pitch"]),
          "velocity": int(nd["velocity"]), "selected": True}
@@ -169,27 +200,61 @@ def humanize(midi_file, model_label, style, bpm, time_signature, beat_type,
     for i, nd in enumerate(pp.notes):
         if i in pred:
             nd["velocity"] = int(pred[i])
-    after = [int(nd["velocity"]) for nd in pp.notes]
+    na_after = pp.note_array().copy()           # humanized, for the roll
+    events_after = _note_events(pp)             # humanized, for the audio
 
     out_midi = os.path.join(tempfile.mkdtemp(), "humanized.mid")
     partitura.save_performance_midi(perf, out_midi)
 
-    fig = _velocity_plot(before, after)
-    audio = _render_audio(pp) if want_audio and os.path.exists(SF2) else None
-    return out_midi, fig, audio
+    fig = _drumroll_figure(na_before, na_after)
+    do_audio = bool(want_audio) and os.path.exists(SF2)
+    audio_before = _render_events(events_before) if do_audio else None
+    audio_after = _render_events(events_after) if do_audio else None
+    return out_midi, fig, audio_before, audio_after
 
 
-# --- example rows: [midi, model, style, bpm, time_sig, beat_type, temp, blend, seed, audio]
+GENRES = LEVELS["genres"]
 _EX_DIR = os.path.join(HERE, "examples")
 EXAMPLES = [
-    [os.path.join(_EX_DIR, "rock_65_beat_4-4.midi"),
-     "Transformer - MDN (temperature)", "rock", 65, "4-4", "beat", 1.0, 1.0, 42, True],
-    [os.path.join(_EX_DIR, "pop_138_beat_4-4.midi"),
-     "LightGBM (fast, deterministic)", "pop", 138, "4-4", "beat", 1.0, 1.0, 42, True],
-    [os.path.join(_EX_DIR, "soul_105_beat_4-4.midi"),
-     "Transformer - Categorical", "soul", 105, "4-4", "beat", 1.0, 1.0, 42, True],
+    [os.path.join(_EX_DIR, "rock_65_beat_4-4.midi")],
+    [os.path.join(_EX_DIR, "pop_138_beat_4-4.midi")],
+    [os.path.join(_EX_DIR, "soul_105_beat_4-4.midi")],
 ]
 
+
+def _autofill_from_file(path):
+    """Best-effort: read musical context from an E-GMD-style filename.
+
+    Fills style / bpm / time-signature / beat-type when the name carries them
+    (e.g. ``rock_65_beat_4-4.midi``); leaves any field it can't parse untouched.
+    Returns gr.update()s in the order (style, bpm, time_signature, beat_type).
+    """
+    import re
+
+    keep = (gr.update(), gr.update(), gr.update(), gr.update())
+    if not path:
+        return keep
+    tokens = re.split(r"[_\s]+", os.path.splitext(os.path.basename(path))[0].lower())
+
+    bpm = next((int(t) for t in tokens if t.isdigit() and 30 <= int(t) <= 300), None)
+    ts = next((t for t in tokens if re.fullmatch(r"\d-\d", t)), None)
+    beat = next((t for t in tokens if t in ("beat", "fill")), None)
+    style = None
+    for t in tokens:
+        if t in STYLES:
+            style = t
+            break
+        if t.split("-")[0] in GENRES:
+            style = t.split("-")[0]
+            break
+
+    return (
+        gr.update(value=style) if style else gr.update(),
+        gr.update(value=bpm) if bpm else gr.update(),
+        gr.update(value=ts) if ts else gr.update(),
+        gr.update(value=beat) if beat else gr.update(),
+    )
+
 DESC = """\
 # 🥁 Dynamics Needed
 
@@ -199,18 +264,24 @@ alone** - never from the note's own loudness - using models trained on the
 [E-GMD](https://magenta.tensorflow.org/datasets/e-gmd) dataset of real
 performances.
 
-Upload a drum-MIDI groove (or pick an example), set its musical context, and
-compare the result by ear and by plot. Same engine that drives the REAPER
+Upload a drum-MIDI groove **or** pick an example, set its musical context, and
+compare the result **by ear** (original vs humanized audio) and **by eye**
+(a drum-roll where colour is velocity). Same engine that drives the REAPER
 plugin, behind a web UI.
 """
 
 
-with gr.Blocks(title="Dynamics Needed", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Dynamics Needed") as demo:
     gr.Markdown(DESC)
     with gr.Row():
         with gr.Column(scale=1):
-            midi_in = gr.File(label="Drum MIDI", file_types=[".mid", ".midi"], type="filepath")
-            model_in = gr.Dropdown(list(MODEL_LABELS), value=list(MODEL_LABELS)[0], label="Model")
+            midi_in = gr.File(label="① Upload a drum MIDI", file_types=[".mid", ".midi"],
+                              type="filepath")
+            gr.Markdown("<div style='text-align:center;opacity:.7'>— or —</div>")
+            gr.Examples(examples=EXAMPLES, inputs=[midi_in],
+                        label="Pick an example groove (E-GMD, CC-BY 4.0)")
+            model_in = gr.Dropdown(list(MODEL_LABELS), value=list(MODEL_LABELS)[0],
+                                   label="② Model")
             style_in = gr.Dropdown(STYLES, value="rock", label="Style / genre",
                                    info="Conditions the transformer; the genre is the part before '/'.")
             with gr.Row():
@@ -224,21 +295,24 @@ with gr.Blocks(title="Dynamics Needed", theme=gr.themes.Soft()) as demo:
                 blend_in = gr.Slider(0.0, 1.0, value=1.0, step=0.05, label="Blend",
                                      info="1.0 = fully predicted, 0.0 = keep original velocities.")
                 seed_in = gr.Number(value=42, label="Seed", precision=0)
-            audio_chk = gr.Checkbox(value=True, label="Render audio preview")
-            run_btn = gr.Button("Predict velocities", variant="primary")
+            audio_chk = gr.Checkbox(value=True, label="Render audio previews")
+            run_btn = gr.Button("③ Predict velocities", variant="primary")
         with gr.Column(scale=1):
-            midi_out = gr.File(label="Humanized MIDI")
-            plot_out = gr.Plot(label="Velocity: original vs predicted")
-            audio_out = gr.Audio(label="Audio preview", type="numpy")
+            with gr.Row():
+                audio_before_out = gr.Audio(label="Original (as uploaded)", type="numpy")
+                audio_after_out = gr.Audio(label="Humanized", type="numpy")
+            plot_out = gr.Plot(label="Drum-roll · colour = velocity (original vs humanized)")
+            midi_out = gr.File(label="Humanized MIDI (download)")
+
+    # Auto-fill musical context from the filename (examples and named uploads).
+    ctx_out = [style_in, bpm_in, ts_in, beat_in]
+    midi_in.change(_autofill_from_file, inputs=[midi_in], outputs=ctx_out)
 
     inputs = [midi_in, model_in, style_in, bpm_in, ts_in, beat_in,
               temp_in, blend_in, seed_in, audio_chk]
-    outputs = [midi_out, plot_out, audio_out]
+    outputs = [midi_out, plot_out, audio_before_out, audio_after_out]
     run_btn.click(humanize, inputs=inputs, outputs=outputs)
 
-    gr.Examples(examples=EXAMPLES, inputs=inputs, outputs=outputs, fn=humanize,
-                cache_examples=False, label="Example grooves (E-GMD, CC-BY 4.0)")
-
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(mcp_server=True, theme=gr.themes.Soft())