"""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": "", "parameters": { "prompt": "optional extra text" } } Response: { "audio_base64": "", "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 importlib import io import os import subprocess import sys import tempfile import traceback import wave from typing import Any import numpy as np import torch def _ensure_pkg(module: str, *pip_names: str) -> None: try: importlib.import_module(module) return except ImportError: pass subprocess.check_call( [sys.executable, "-m", "pip", "install", "--no-cache-dir", *pip_names], ) importlib.invalidate_caches() importlib.import_module(module) # Diffusers exposes dummy SpectrogramDiffusionPipeline objects unless note-seq # is present *before* that symbol is first resolved. Install extras first. _ensure_pkg("note_seq", "note-seq", "pretty_midi") _ensure_pkg("onnxruntime", "onnxruntime") def _pipeline_classes(): try: from diffusers import MidiProcessor, SpectrogramDiffusionPipeline except ImportError: # moved under deprecated/ in newer Diffusers from diffusers.pipelines.deprecated.spectrogram_diffusion import ( MidiProcessor, SpectrogramDiffusionPipeline, ) dummy = "dummy" in getattr(SpectrogramDiffusionPipeline, "__module__", "") dummy = dummy or bool(getattr(SpectrogramDiffusionPipeline, "_backends", None)) if not dummy: return MidiProcessor, SpectrogramDiffusionPipeline for key in list(sys.modules): if key == "diffusers" or key.startswith("diffusers."): del sys.modules[key] from diffusers import MidiProcessor, SpectrogramDiffusionPipeline return MidiProcessor, SpectrogramDiffusionPipeline MidiProcessor, SpectrogramDiffusionPipeline = _pipeline_classes() def _patch_t5_cache_position() -> None: """Newer Transformers T5Attention requires cache_position; this pipeline never passes it.""" import inspect try: from transformers.models.t5.modeling_t5 import T5Attention except ImportError: return orig = T5Attention.forward if "cache_position" not in inspect.signature(orig).parameters: return def wrapped(self, *args, **kwargs): # type: ignore[no-untyped-def] if kwargs.get("cache_position") is None: hidden = args[0] if args else kwargs.get("hidden_states") if hidden is not None: kwargs["cache_position"] = torch.arange(hidden.shape[1], device=hidden.device) return orig(self, *args, **kwargs) T5Attention.forward = wrapped # type: ignore[method-assign] _patch_t5_cache_position() 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 if _has_diffusers_index(path): # model_index.json names encoders as library `spectrogram_diffusion`. # The shim package lives next to this handler in the snapshot. sys.path.insert(0, os.path.abspath(path)) 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]: try: payload = data if isinstance(data, dict) else {"inputs": data} midi_bytes = _midi_bytes(payload) params = payload.get("parameters") if isinstance(payload.get("parameters"), dict) else {} kwargs: dict[str, Any] = {} steps = params.get("num_inference_steps") if steps is not None: kwargs["num_inference_steps"] = max(1, int(steps)) with tempfile.NamedTemporaryFile(suffix=".mid", delete=True) as tmp: tmp.write(midi_bytes) tmp.flush() tokens = self.processor(tmp.name) if not tokens: raise ValueError("MIDI produced no spectrogram tokens (clip too short?)") 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", } except Exception as err: raise RuntimeError( f"{type(err).__name__}: {err}; {_payload_preview(data)}\n{traceback.format_exc()}" ) from err def _has_diffusers_index(path: str) -> bool: return bool(path) and os.path.isfile(os.path.join(path, "model_index.json")) def _payload_preview(data: Any) -> str: if isinstance(data, dict): keys = ",".join(sorted(str(k) for k in data.keys())) raw = data.get("inputs", data.get("midi_base64")) return f"payload_keys={keys} inputs_type={type(raw).__name__}" return f"payload_type={type(data).__name__}" def _midi_bytes(payload: dict[str, Any]) -> bytes: raw = payload.get("inputs") if raw is None: raw = 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 > 1e-6: arr = arr * (0.89 / 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()