app.py

import gradio as gr
import torchaudio
import torch
import numpy as np
from huggingface_hub import hf_hub_download

# ✅ Map of model names to files on Hugging Face
RAVE_MODELS = {
    "Guitar": "guitar_iil_b2048_r48000_z16.ts",
    "Soprano Sax": "sax_soprano_franziskaschroeder_b2048_r48000_z20.ts",
    "Organ (Archive)": "organ_archive_b2048_r48000_z16.ts",
    "Organ (Bach)": "organ_bach_b2048_r48000_z16.ts",
    "Voice Multivoice": "voice-multi-b2048-r48000-z11.ts",
    "Birds Dawn Chorus": "birds_dawnchorus_b2048_r48000_z8.ts",
    "Magnets": "magnets_b2048_r48000_z8.ts",
    "Whale Songs": "humpbacks_pondbrain_b2048_r48000_z20.ts"
}

MODEL_CACHE = {}

def load_rave_model(model_name):
    """Load TorchScript RAVE model from Hugging Face Hub."""
    if model_name in MODEL_CACHE:
        return MODEL_CACHE[model_name]

    model_file = hf_hub_download(
        repo_id="Intelligent-Instruments-Lab/rave-models",
        filename=RAVE_MODELS[model_name]
    )

    model = torch.jit.load(model_file, map_location="cpu")
    model.eval()
    MODEL_CACHE[model_name] = model
    return model

def apply_rave(audio, model_name):
    """Apply selected RAVE model to uploaded audio."""
    model = load_rave_model(model_name)

    # Convert numpy audio to torch tensor
    audio_tensor = torch.tensor(audio[0]).unsqueeze(0)  # [1, samples]
    sr = audio[1]

    # ✅ Resample if needed (most RAVE models expect 48kHz)
    if sr != 48000:
        audio_tensor = torchaudio.functional.resample(audio_tensor, sr, 48000)
        sr = 48000

    with torch.no_grad():
        # ✅ TorchScript models have encode & decode methods
        z = model.encode(audio_tensor)
        processed_audio = model.decode(z)

    return (processed_audio.squeeze().cpu().numpy(), sr)

# 🎛 Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🎛 RAVE Style Transfer on Stems")
    gr.Markdown("Upload audio, pick a RAVE model, and get a transformed version.")

    with gr.Row():
        audio_input = gr.Audio(type="numpy", label="Upload Audio", sources=["upload", "microphone"])
        model_selector = gr.Dropdown(list(RAVE_MODELS.keys()), label="Select Style", value="Guitar")

    with gr.Row():
        output_audio = gr.Audio(type="numpy", label="Transformed Audio")

    process_btn = gr.Button("Apply Style Transfer")
    process_btn.click(fn=apply_rave, inputs=[audio_input, model_selector], outputs=output_audio)

demo.launch()