"""Dynamics Needed - Gradio demo for structure-driven drum-velocity prediction. Upload a drum-MIDI groove (flat or existing velocities), pick a model and the groove's musical context, and get back a version whose note velocities are predicted from *structure and timing alone* - plus a before/after velocity plot and an audio preview rendered through a General MIDI SoundFont. Runtime notes ------------- * ``import spaces`` MUST precede any torch import so its ZeroGPU hijack lands. The models here are tiny and CPU-only, so a single no-op ``@spaces.GPU`` function satisfies the ZeroGPU requirement without ever requesting a GPU (nothing in the hot path is decorated, so no visitor quota is burned). * The whole ``drum_dynamics`` package is vendored next to this file at deploy time, and the six ready-to-load model files live under ``models/``. """ from __future__ import annotations import spaces # noqa: F401 (must come before torch; see module docstring) import os HERE = os.path.dirname(os.path.abspath(__file__)) SF2 = os.path.join(HERE, "sf", "FluidR3_GM.sf2") SR = 44100 def _preseed_partitura_soundfont() -> None: """Stop partitura from FTP-downloading a soundfont at import time. partitura's top-level ``__init__`` imports its audio-export module, which - when fluidsynth is importable and its bundled asset is missing - fetches a MuseScore soundfont over FTP. HF Spaces block FTP egress, so that call hangs and the app crashes on ``import partitura``. Pre-seed the asset (symlink our bundled GM soundfont) so partitura finds it and skips the download. Must run BEFORE ``import partitura``. """ import importlib.util spec = importlib.util.find_spec("partitura") # locate without importing if not spec or not spec.origin: return assets = os.path.join(os.path.dirname(spec.origin), "assets") target = os.path.join(assets, "MuseScore_General.sf3") if os.path.exists(target): return try: os.makedirs(assets, exist_ok=True) if os.path.exists(SF2): os.symlink(SF2, target) else: open(target, "wb").close() # placeholder; partitura's synth is unused except OSError: pass _preseed_partitura_soundfont() import tempfile import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import gradio as gr 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")) LEVELS = ENGINE.levels() STYLES = LEVELS["styles"] if os.path.exists(SF2): set_soundfont(SF2) MODEL_LABELS = { "Transformer - MDN (temperature)": "mdn", "Transformer - Categorical": "categorical", "LightGBM (deterministic, less dynamic)": "lgbm", } # --- ZeroGPU registration hook (never called on the hot path) ------------------ @spaces.GPU(duration=1) def _zerogpu_noop(): # pragma: no cover - satisfies the "needs one GPU fn" rule """No-op so the Space qualifies as ZeroGPU; real inference runs on CPU.""" return "ok" # --- helpers ------------------------------------------------------------------ def _load_notes(midi_path: str): """Return the partitura performance + its editable note-dict list.""" perf = partitura.load_performance_midi(midi_path) pp = perf.performedparts[0] return perf, pp 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 fl = fluidsynth.Synth(samplerate=SR) sfid = fl.sfload(SF2) fl.program_select(9, sfid, 128, 0) # GM percussion set on channel 9 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(timeline[:-1], timeline[1:]): if is_on: fl.noteon(9, pitch, vel) else: fl.noteoff(9, pitch) n = int(max(0, (t1 - t0)) * SR) if n: audio.extend(fl.get_samples(n)[::2]) # interleaved stereo -> mono audio.extend(fl.get_samples(SR)[::2]) # let the last hits ring out fl.delete() return SR, np.asarray(audio, dtype=np.int16) 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 # --- main callback ------------------------------------------------------------ def humanize(midi_file, model_label, style, bpm, time_signature, beat_type, temperature, blend, seed, want_audio): if midi_file is None: raise gr.Error("Upload a drum-MIDI file first (or click an example).") model = MODEL_LABELS[model_label] perf, pp = _load_notes(midi_file) if not pp.notes: raise gr.Error("No notes found in that MIDI file.") 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} for i, nd in enumerate(pp.notes) ] request = { "model": model, "seed": int(seed), "temperature": float(temperature), "blend": float(blend), "bpm": float(bpm), "time_signature": time_signature, "style": style, "beat_type": beat_type, "notes": notes, } pred = ENGINE.predict(request) # {note_index: new_velocity} for i, nd in enumerate(pp.notes): if i in pred: nd["velocity"] = int(pred[i]) 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 = _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 GENRES = LEVELS["genres"] _EX_DIR = os.path.join(HERE, "examples") EXAMPLES = [ [os.path.join(_EX_DIR, "funk_138_beat_4-4.midi")], [os.path.join(_EX_DIR, "jazz-funk_116_beat_4-4.midi")], [os.path.join(_EX_DIR, "soul_102_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 Drum programming often lands notes at a **flat, robotic velocity**. This demo predicts a human-sounding velocity for every note from **structure and timing 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** (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") as demo: gr.Markdown(DESC) with gr.Row(): with gr.Column(scale=1): midi_in = gr.File(label="① Upload a drum MIDI", file_types=[".mid", ".midi"], type="filepath") gr.Markdown("