Add custom Inference Endpoint handler (SpectrogramDiffusionPipeline)

handler.py

@@ -0,0 +1,141 @@
+"""Custom Hugging Face Inference Endpoint handler.
+
+Loads Google's Multi-Instrument Spectrogram Diffusion
+(`google/music-spectrogram-diffusion`) from the endpoint's local snapshot
+and turns a MIDI file into a 16 kHz WAV.
+
+This file must live at the *root* of the Hub repo the endpoint points at,
+alongside `requirements.txt`. `google/music-spectrogram-diffusion` has no
+`handler.py` and no Transformers `config.json`, so a Custom task against
+that repo falls back to `pipeline()` and crashes. Fork the model, add these
+two files, then set the endpoint to that fork with task=custom.
+
+Request body (Band's `/midi-to-audio` client):
+
+    { "inputs": "<base64 .mid>", "parameters": { "prompt": "optional extra text" } }
+
+Response:
+
+    { "audio_base64": "<base64 wav>", "sample_rate": 16000, "format": "wav" }
+
+The prompt is accepted so Band can send clip text; this pipeline does not
+condition on it. Keep the field so a later handler can without a client change.
+"""
+from __future__ import annotations
+
+import base64
+import io
+import os
+import wave
+from typing import Any
+
+import numpy as np
+import torch
+
+try:
+    from diffusers import MidiProcessor, SpectrogramDiffusionPipeline
+except ImportError:  # moved under deprecated/ in newer Diffusers
+    from diffusers.pipelines.deprecated.spectrogram_diffusion import (
+        MidiProcessor,
+        SpectrogramDiffusionPipeline,
+    )
+
+
+SAMPLE_RATE = 16000
+FALLBACK_MODEL_ID = os.environ.get("HF_MSD_MODEL_ID", "google/music-spectrogram-diffusion")
+
+
+class EndpointHandler:
+    def __init__(self, path: str = "") -> None:
+        # Inference Endpoints pass the local snapshot (`/repository`), which
+        # has `model_index.json` + decoder/encoder/melgan weights. Load that,
+        # not the Hub id — otherwise the replica re-downloads and ignores the
+        # fork this handler was deployed from.
+        model_path = path if _has_diffusers_index(path) else FALLBACK_MODEL_ID
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.float16 if device == "cuda" else torch.float32
+        self.pipe = SpectrogramDiffusionPipeline.from_pretrained(
+            model_path,
+            torch_dtype=dtype,
+        )
+        try:
+            self.pipe.to(device)
+        except Exception:
+            # MelGAN is ONNX and may refuse `.to()`. Move the PyTorch modules.
+            for name in ("notes_encoder", "continuous_encoder", "decoder"):
+                module = getattr(self.pipe, name, None)
+                if module is not None:
+                    module.to(device)
+        self.processor = MidiProcessor()
+        self.device = device
+
+    def __call__(self, data: Any) -> dict[str, Any]:
+        payload = data if isinstance(data, dict) else {}
+        midi_bytes = _midi_bytes(payload)
+        params = payload.get("parameters") if isinstance(payload.get("parameters"), dict) else {}
+        steps = params.get("num_inference_steps")
+        kwargs: dict[str, Any] = {}
+        if steps is not None:
+            kwargs["num_inference_steps"] = max(1, int(steps))
+
+        tokens = self.processor(midi_bytes)
+        audio = self.pipe(tokens, **kwargs)
+        samples = _to_mono_float(audio)
+        wav = _wav_bytes(samples, SAMPLE_RATE)
+        return {
+            "audio_base64": base64.b64encode(wav).decode("ascii"),
+            "sample_rate": SAMPLE_RATE,
+            "format": "wav",
+        }
+
+
+def _has_diffusers_index(path: str) -> bool:
+    return bool(path) and os.path.isfile(os.path.join(path, "model_index.json"))
+
+
+def _midi_bytes(payload: dict[str, Any]) -> bytes:
+    raw = payload.get("inputs", payload.get("midi_base64", ""))
+    if isinstance(raw, list) and raw:
+        raw = raw[0]
+    if isinstance(raw, dict):
+        raw = raw.get("midi_base64") or raw.get("inputs") or raw.get("data") or ""
+    if isinstance(raw, (bytes, bytearray)):
+        midi = bytes(raw)
+    elif isinstance(raw, str) and raw:
+        text = raw.strip()
+        if text.lower().startswith("data:") and "," in text:
+            text = text.split(",", 1)[1]
+        midi = base64.b64decode(text)
+    else:
+        raise ValueError("missing base64 MIDI in inputs")
+    if len(midi) < 8 or midi[:4] != b"MThd":
+        raise ValueError("inputs is not a MIDI file")
+    return midi
+
+
+def _to_mono_float(audio: Any) -> np.ndarray:
+    if hasattr(audio, "audios") and audio.audios is not None:
+        arr = np.asarray(audio.audios, dtype=np.float32)
+    elif isinstance(audio, (list, tuple)) and audio:
+        arr = np.asarray(audio[0], dtype=np.float32)
+    else:
+        arr = np.asarray(audio, dtype=np.float32)
+    arr = np.squeeze(arr)
+    if arr.ndim > 1:
+        arr = arr.reshape(-1)
+    peak = float(np.max(np.abs(arr))) if arr.size else 0.0
+    if peak > 1.0:
+        arr = arr / peak
+    return arr
+
+
+def _wav_bytes(samples: np.ndarray, rate: int) -> bytes:
+    pcm = np.clip(samples, -1.0, 1.0)
+    pcm = (pcm * 32767.0).astype(np.int16)
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(rate)
+        wf.writeframes(pcm.tobytes())
+    return buf.getvalue()