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)