handler.py

"""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 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()