app.py

import sys
sys.stdout.reconfigure(line_buffering=True)

try:
    import spaces
except ImportError:
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                func = args[0]
                return func

import os
import tempfile
import threading
from pathlib import Path
import torch
import torchaudio
import soundfile as sf
import gradio as gr
from huggingface_hub import hf_hub_download
from safetensors import torch as sft

HF_TOKEN = os.environ.get("HF_TOKEN")

import hyperparameters as hp
from conditioning.condition_type import ConditionType
from conditioning.conditioning_method import ConditioningMethod
from conditioning.prompt_processor import InterleavedContextPromptProcessor
from conditioning.t5embedder import T5EmbedderGPU
from models.lightning_musicgen import LightningMusicgen
from pyharp import ModelCard, build_endpoint

REPO_ID = "teamup-tech/STAGE-checkpoints"
SAMPLE_RATE = 32_000
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

drums_model: LightningMusicgen | None = None
bass_model: LightningMusicgen | None = None
models_on_device = False
model_loading = True
model_error: str | None = None

# from hyperparameters.py
def build_params(encodec_weights: str, lm_weights: str) -> hp.MusicgenParams:
    return hp.MusicgenParams(
        encodec_params=hp.EncodecParams(
            sample_rate=32_000,
            seanet_params=hp.SeaNetParams(128, 64, (8, 5, 4, 4), False, True),
            quantizer_params=hp.QuantizerParams(128, 4, 2048),
            sum_loss_mulitiplier=0,
            weights=encodec_weights,
        ),
        prompt_processor_params=hp.PromptProcessorParams(
            keep_only_valid_steps=True,
            model_class=InterleavedContextPromptProcessor,
            context_dropout=0.1,
        ),
        conditioning_params=hp.ConditioningParams(
            embedder_types={ConditionType.DESCRIPTION: T5EmbedderGPU},
            conditioning_methods={
                ConditionType.DESCRIPTION: ConditioningMethod.CROSS_ATTENTION
            },
            conditioning_dropout=0.5,
        ),
        lm_params=hp.PretrainedSmallLmParams(sep_token=2049, weights=lm_weights),
    )


def load_checkpoint(params: hp.MusicgenParams, ckp_path: str) -> LightningMusicgen:
    model = LightningMusicgen(params)  # model_class strings in hyperparameters.py use the old "stage.*" prefix
    sft.load_model(model, ckp_path)
    return model.cpu().eval()


def load_models():
    """Download all weights and build both models on CPU. Background thread only — no CUDA."""
    global drums_model, bass_model, model_loading, model_error
    try:
        print("Downloading shared weights...")
        encodec_path = hf_hub_download(REPO_ID, "encodec_32khz.pt", token=HF_TOKEN)
        lm_path = hf_hub_download(REPO_ID, "lm-small-weights.pt", token=HF_TOKEN)
        params = build_params(encodec_path, lm_path)

        print("Building drums model (CPU)...")
        drums_ckp = hf_hub_download(REPO_ID, "stage-drums.safetensors", token=HF_TOKEN)
        drums_model = load_checkpoint(params, drums_ckp)
        print("Drums model ready.")

        print("Building bass model (CPU)...")
        bass_ckp = hf_hub_download(REPO_ID, "stage-bass.safetensors", token=HF_TOKEN)
        bass_model = load_checkpoint(params, bass_ckp)
        print("Bass model ready.")
    except Exception as e:
        model_error = str(e)
        print(f"Load error: {e}")
    finally:
        model_loading = False


threading.Thread(target=load_models, daemon=True).start()


model_card = ModelCard(
    name="STAGE",
    description=(
        "Single Stem Accompaniment Generation. Provide an audio mix or click track "
        "and STAGE generates a coherent drums or bass stem to accompany it."
    ),
    author="Giorgio Strano, Vansh Chaudhary, Derek Tran",
    tags=["music-generation", "accompaniment", "stems"],
)


@spaces.GPU(duration=120)
@torch.inference_mode()
def process_fn(
    input_audio_path: str,
    instrument: str,
    gen_seconds: int,
    description: str,
) -> str:
    global drums_model, bass_model, models_on_device

    if model_loading:
        raise gr.Error("Model is still loading — please wait a moment and try again.")
    if model_error:
        raise gr.Error(f"Model failed to load: {model_error}")

    if not models_on_device:
        drums_model = drums_model.to(DEVICE)  # type: ignore[union-attr]
        bass_model = bass_model.to(DEVICE)    # type: ignore[union-attr]
        models_on_device = True

    model = drums_model if instrument == "Drums" else bass_model

    audio_np, orig_sr = sf.read(input_audio_path, always_2d=True)
    context = torch.from_numpy(audio_np.T).float()
    context = torchaudio.functional.resample(context, orig_sr, SAMPLE_RATE)
    if context.shape[0] > 1:
        context = context.mean(dim=0, keepdim=True)
    context = context.reshape(1, 1, -1).to(DEVICE)

    generated = model.generate(
        n_samples=1,
        gen_seconds=gen_seconds,
        prompt=None,
        context=context,
        style=None,
        beat=None,
        description=[description.strip() or None],
    )

    # generated: (1, 1, T) — squeeze to (1, T) and duplicate to stereo
    audio_out = generated.squeeze(0).cpu().float()
    if audio_out.shape[0] == 1:
        audio_out = audio_out.repeat(2, 1)

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        out_path = f.name
    sf.write(out_path, audio_out.T.numpy(), samplerate=SAMPLE_RATE)
    return out_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(
            type="filepath",
            label="Context Track",
        ).harp_required(True),
        gr.Dropdown(
            choices=["Drums", "Bass"],
            value="Drums",
            label="Instrument",
            info="Which accompaniment stem to generate.",
        ),
        gr.Slider(
            minimum=5,
            maximum=20,
            step=1,
            value=10,
            label="Length (seconds)",
            info="Duration of the generated stem (default: 10 sec, per paper).",
        ),
        gr.Textbox(
            value="",
            label="Style Description",
            info="Optional — describe the mood or style (e.g. 'lo-fi chill groove').",
            placeholder="lo-fi chill groove with soft kick...",
        ),
    ]
    output_components = [
        gr.Audio(
            type="filepath",
            label="Generated Stem",
        ).set_info(
            "The generated accompaniment stem. Mix it back with your context track."
        ),
    ]

    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

demo.queue().launch(pwa=True)