app.py

# ============================================================================
#  app.py — ZeroGPU deployment via gradio.Server
#
#  gradio.Server extends FastAPI, so:
#    - @app.get("/")           serves the custom HTML editor (unchanged)
#    - @app.api()              wraps the GPU transcription with Gradio's queue
#                              + ZeroGPU allocation (@spaces.GPU)
#    - @app.post("/analyze")   and @app.post("/pmi") are plain FastAPI routes
#                              (CPU only — key estimation / PMI, no GPU)
#

# ============================================================================
import os, base64, tempfile, time

import spaces                      # ZeroGPU
from gradio import Server
from gradio.data_classes import FileData
from fastapi import UploadFile, File, Form
from fastapi.responses import FileResponse, JSONResponse
import pretty_midi

import pmi_core

# NOTE: heavy imports (torch, demucs, YourMT3+) happen lazily inside pipeline.py
# so the web server starts fast and the model is only touched inside the GPU call.
from pipeline import chroma_from_audio, demucs_stem, estimate_bpm, notes_from_audio


# ---------------------------------------------------------------------------
# helpers (identical to the FastAPI version)
# ---------------------------------------------------------------------------
def notes_to_midi_bytes(notes):
    pm = pretty_midi.PrettyMIDI()
    inst = pretty_midi.Instrument(program=0)
    for n in notes:
        s, e, p = n[0], n[1], n[2]
        conf = n[3] if len(n) > 3 else 1.0
        vel = max(1, min(127, int(round(conf * 127))))
        inst.notes.append(pretty_midi.Note(velocity=vel, pitch=int(p),
                                            start=float(s), end=float(e)))
    pm.instruments.append(inst)
    t = tempfile.NamedTemporaryFile(delete=False, suffix=".mid"); t.close()
    pm.write(t.name)
    data = open(t.name, "rb").read()
    os.unlink(t.name)
    return data


def _save_temp_midi(raw_bytes):
    t = tempfile.NamedTemporaryFile(delete=False, suffix=".mid")
    t.write(raw_bytes); t.flush(); t.close()
    return t.name


def _parse_forced(tonic, mode):
    if tonic is None or tonic == "" or mode is None or mode == "":
        return None
    try:
        return (int(tonic) % 12, "minor" if str(mode).lower().startswith("min") else "major")
    except (ValueError, TypeError):
        return None


# ---------------------------------------------------------------------------
# the GPU work: demucs + YourMT3+ + post-processing, ONE @spaces.GPU call.
#
# Why one call and not two (separation / transcription split): each entry into
# a @spaces.GPU function pays a fixed toll -- a process fork, CUDA re-attach,
# a pickle round-trip, and a fresh pass through the node queue -- and each
# entry runs the caller through the quota gate again. Splitting doubles all of
# that per clip, which is exactly what burned through visitors' free quota.
# Quota is charged on ACTUAL seconds used, not on the declared duration; the
# declaration is only (a) an admission check against the caller's remaining
# quota and (b) a queue-priority hint (shorter ranks higher). So the right
# declaration is a realistic estimate with modest headroom -- never a fat
# safety margin, which locks callers out long before their quota is spent.
# ---------------------------------------------------------------------------
# Set after every successful GPU call and read by _gpu_duration to decide
# whether the next request can reasonably expect a warm worker. Lives in the
# main process (the web server), which is where the duration callable runs.
_last_gpu_ok = 0.0
WARM_WINDOW = 120.0        # seconds; conservative -- covers back-to-back testing

def _resolve(audio_path):
    """FileData arrives as an object with .path, a dict, or a plain string."""
    if isinstance(audio_path, dict):
        return audio_path.get("path", "")
    if hasattr(audio_path, "path"):
        return audio_path.path
    return audio_path


def _audio_seconds(path):
    try:
        import soundfile as _sf
        with _sf.SoundFile(path) as f:
            return len(f) / float(f.samplerate)
    except Exception:
        return 30.0                      # unknown -> assume a modest clip

def _gpu_duration(audio_path, instrument="vocals"):
 
    sec = _audio_seconds(_resolve(audio_path))
    if (time.monotonic() - _last_gpu_ok) < WARM_WINDOW:
        return int(max(5, min(120, 0.2 * sec)))   
    return int(max(15, min(120, 0.25 * sec)))     
@spaces.GPU(duration=_gpu_duration)
def _transcribe_gpu(audio_path, instrument="vocals"):
    """Only the steps that actually need the GPU: demucs separation and YourMT3+
    inference. Chroma extraction and MIDI encoding are pure CPU work and run
    outside this window, where they cost the visitor no quota.

    Beat tracking stays inside despite being CPU work: notes_from_audio needs
    the tempo, and a second GPU entry to hand it back in would cost far more
    (fork, CUDA re-attach, queue, quota gate) than the ~0.5 s it takes."""
    stem_path = demucs_stem(audio_path, instrument)   # GPU: separation
    tempo = estimate_bpm(stem_path)                   # CPU, but needed inline:
                                                      # notes_from_audio needs it
    notes = notes_from_audio(stem_path, tempo=tempo)  # GPU: YourMT3+
    # Return the stem as BYTES, not a path: the GPU worker is a separate,
    # short-lived process and its return value crosses a pickle boundary, so a
    # path would rely on the file still being readable from the main process.
    # Bytes are self-contained.
    return open(stem_path, "rb").read(), notes


# ---------------------------------------------------------------------------
# Both excerpts in ONE GPU call.
## ---------------------------------------------------------------------------
def _pair_duration(audio_a, audio_b, instrument_a="vocals", instrument_b="vocals"):
    total = _audio_seconds(_resolve(audio_a)) + _audio_seconds(_resolve(audio_b))
    if (time.monotonic() - _last_gpu_ok) < WARM_WINDOW:
        return int(max(20, min(120, 8 + 0.2 * total)))
    return int(max(28, min(120, 15 + 0.25 * total)))


@spaces.GPU(duration=_pair_duration)
def _transcribe_pair_gpu(audio_a, audio_b, instrument_a="vocals", instrument_b="vocals"):
    """Separate + transcribe two excerpts in one GPU window. The model is moved
    to the device once by the first call into yourmt3_transcribe and stays there
    for the second, so the second excerpt costs only its own compute."""
    out = []
    for path, inst in ((audio_a, instrument_a), (audio_b, instrument_b)):
        stem_path = demucs_stem(_resolve(path), inst)
        tempo = estimate_bpm(stem_path)
        notes = notes_from_audio(stem_path, tempo=tempo)
        out.append((open(stem_path, "rb").read(), notes))
    return out[0], out[1]


# ---------------------------------------------------------------------------
# Warm-up at startup, on CPU, where time is free. Loading the YourMT3+
# checkpoint takes ~10-15 s; done lazily it lands inside the FIRST visitor's
# GPU window and is billed to THEIR quota. Doing it here means the GPU window
# only ever pays for the cheap cpu->cuda move, not the disk load. demucs
# weights are fetched here too, so no download ever happens on GPU time.
# ---------------------------------------------------------------------------
try:
    from pipeline import _load_model
    _load_model()                                  # YourMT3+ -> CPU, cached
except Exception as _e:
    print(f"[warmup] YourMT3+ preload failed (will retry lazily): {_e}", flush=True)
try:
    import demucs.pretrained as _dp
    _dp.get_model(name="htdemucs")                 # default (vocals) weights
except Exception as _e:
    print(f"[warmup] demucs preload failed (will retry lazily): {_e}", flush=True)


# ---------------------------------------------------------------------------
# gradio.Server (a FastAPI app with Gradio's API engine on top)
# ---------------------------------------------------------------------------
app = Server()


@app.get("/")
def index():
    return FileResponse("transcription-editor.html")


# transcription: wrapped by Gradio's queue + ZeroGPU. concurrency_limit=1
# because a single ZeroGPU slice serves one transcription at a time.
# Returns base64 strings so the existing front-end decoding still works.
@app.api(name="transcribe", concurrency_limit=1)
def transcribe(audio_path: FileData, instrument: str = "vocals") -> dict:
    """audio_path: a file uploaded via @gradio/client handle_file().
    instrument: which stem to isolate -- "vocals" (default, the validated path),
    "guitar" or "piano". Only single-line melodies are meaningful downstream:
    PMI aligns one melodic line, and its thresholds were calibrated on vocals.
    Returns {'midi': b64, 'vocal': b64, 'chroma': {...}}; the 'vocal' key holds
    whichever stem was requested."""
    path = _resolve(audio_path)
    if instrument not in ("vocals", "guitar", "piano"):
        instrument = "vocals"

    # --- GPU window: separation + transcription only --------------------
    global _last_gpu_ok
    stem_bytes, notes = _transcribe_gpu(path, instrument)
    _last_gpu_ok = time.monotonic()      # a worker was warm as of now

    # --- CPU (no quota): encoding + chroma for the completion hints ------
    midi_bytes = notes_to_midi_bytes(notes)
    _st = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    _st.write(stem_bytes); _st.flush(); _st.close()
    try:
        chroma = chroma_from_audio(_st.name)
    finally:
        try: os.unlink(_st.name)
        except OSError: pass

    return {
        "midi":  base64.b64encode(midi_bytes).decode(),
        "vocal": base64.b64encode(stem_bytes).decode(),
        "chroma": {
            "data":     base64.b64encode(chroma["data"]).decode(),
            "energy":   base64.b64encode(chroma["energy"]).decode(),
            "n_frames": chroma["n_frames"],
            "sr":       chroma["sr"],
            "hop":      chroma["hop"],
        },
    }


@app.api(name="transcribe_pair", concurrency_limit=1)
def transcribe_pair(audio_a: FileData, audio_b: FileData,
                    instrument_a: str = "vocals",
                    instrument_b: str = "vocals") -> dict:
    """Both excerpts in one call, so ZeroGPU counts one run instead of two.
    Returns {'a': {...}, 'b': {...}} with each side shaped exactly like the
    single-excerpt response, so the front end can reuse the same decoding."""
    global _last_gpu_ok
    if instrument_a not in ("vocals", "guitar", "piano"): instrument_a = "vocals"
    if instrument_b not in ("vocals", "guitar", "piano"): instrument_b = "vocals"

    # --- the single GPU window -------------------------------------------
    (stem_a, notes_a), (stem_b, notes_b) = _transcribe_pair_gpu(
        audio_a, audio_b, instrument_a, instrument_b)
    _last_gpu_ok = time.monotonic()

    # --- CPU (no quota): encoding + chroma -------------------------------
    def _pack(stem_bytes, notes):
        t = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        t.write(stem_bytes); t.flush(); t.close()
        try:
            chroma = chroma_from_audio(t.name)
        finally:
            try: os.unlink(t.name)
            except OSError: pass
        return {
            "midi":  base64.b64encode(notes_to_midi_bytes(notes)).decode(),
            "vocal": base64.b64encode(stem_bytes).decode(),
            "chroma": {
                "data":     base64.b64encode(chroma["data"]).decode(),
                "energy":   base64.b64encode(chroma["energy"]).decode(),
                "n_frames": chroma["n_frames"],
                "sr":       chroma["sr"],
                "hop":      chroma["hop"],
            },
        }

    return {"a": _pack(stem_a, notes_a), "b": _pack(stem_b, notes_b)}


# key analysis — CPU only, plain FastAPI route (no GPU, no gradio client needed)
@app.post("/analyze")
async def analyze_api(midi: UploadFile = File(...)):
    p = _save_temp_midi(await midi.read())
    try:
        return JSONResponse(pmi_core.analyze_one_midi(p))
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=400)
    finally:
        try: os.unlink(p)
        except OSError: pass


# PMI — CPU only, plain FastAPI route
@app.post("/pmi")
async def pmi_api(midi_a: UploadFile = File(...), midi_b: UploadFile = File(...),
                  tonic_a: str = Form(None), mode_a: str = Form(None),
                  tonic_b: str = Form(None), mode_b: str = Form(None)):
    pa = _save_temp_midi(await midi_a.read())
    pb = _save_temp_midi(await midi_b.read())
    try:
        result = pmi_core.pmi_from_two_midis(
            pa, pb,
            forced_A=_parse_forced(tonic_a, mode_a),
            forced_B=_parse_forced(tonic_b, mode_b),
        )
        return JSONResponse(result)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=400)
    finally:
        for p in (pa, pb):
            try: os.unlink(p)
            except OSError: pass


if __name__ == "__main__":
    app.launch()