app.py

import sys

sys.stdout.reconfigure(line_buffering=True)

try:
    import spaces
except ImportError:
    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
    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 shutil
import tempfile
import threading

import gradio as gr
import numpy as np
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint

from songgen import (
    SongGenDualTrackForConditionalGeneration,
    SongGenMixedForConditionalGeneration,
    SongGenProcessor,
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

MODES = {
    "Mixed": SongGenMixedForConditionalGeneration,
    "Dual-track (vocals + accompaniment)": SongGenDualTrackForConditionalGeneration,
}
CHECKPOINTS = {
    "Mixed": "LiuZH-19/SongGen_mixed_pro",
    "Dual-track (vocals + accompaniment)": "LiuZH-19/SongGen_interleaving_A_V",
}

state = {name: {"model": None, "processor": None} for name in MODES}
xcodec_ready = False
xcodec_error = None


def fetch_xcodec_assets():
    """Grab the X-Codec checkpoint and its config. SongGen's XCodecModel
    hardcodes this exact folder rather than taking a path argument, so the
    files have to be here before any SongGen model is constructed."""
    ckpt_dir = os.path.join(
        os.path.dirname(__file__), "songgen", "xcodec_wrapper", "xcodec_infer", "ckpts", "general_more"
    )
    os.makedirs(ckpt_dir, exist_ok=True)
    for filename in ("xcodec_hubert_general_audio_v2.pth", "config_hubert_general.yaml"):
        dest = os.path.join(ckpt_dir, filename)
        if not os.path.exists(dest):
            src = hf_hub_download(repo_id="ZhenYe234/xcodec", filename=filename)
            shutil.copy(src, dest)


def load_xcodec_assets_bg():
    """Fetch the shared X-Codec assets in the background. Pure file I/O, so
    unlike the actual SongGen models it's safe to do off the main thread."""
    global xcodec_ready, xcodec_error
    try:
        fetch_xcodec_assets()
    except Exception as e:
        xcodec_error = str(e)
    finally:
        xcodec_ready = True


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


model_card = ModelCard(
    name="SongGen",
    description=(
        "Single-stage auto-regressive transformer for text-to-song generation. "
        "Give it lyrics and a description of the music, plus an optional reference "
        "voice, and it generates a full song."
    ),
    author="Zihan Liu, Shuangrui Ding, Zhixiong Zhang, Xiaoyi Dong, Pan Zhang, Yuhang Zang, Yuhang Cao, Dahua Lin, Jiaqi Wang",
    tags=["text-to-music", "song-generation"],
)


def get_mode(mode_name):
    """Build a mode's model + processor on first use, then reuse them.
    Loading MERT's remote code touches CUDA, so this can only happen inside
    an @spaces.GPU call, not at startup."""
    entry = state[mode_name]
    if entry["model"] is None:
        if not xcodec_ready:
            raise gr.Error("Still preparing shared assets, please wait a moment and try again.")
        if xcodec_error:
            raise gr.Error(f"Failed to prepare shared assets: {xcodec_error}")
        repo_id = CHECKPOINTS[mode_name]
        model = MODES[mode_name].from_pretrained(repo_id, attn_implementation="sdpa").to(DEVICE)
        processor = SongGenProcessor(repo_id, DEVICE)
        entry["model"] = model
        entry["processor"] = processor
    return entry["model"], entry["processor"]


def write_wav(audio_arr, sample_rate):
    """Save a numpy audio array to a temp wav file and return its path."""
    path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    sf.write(path, audio_arr, sample_rate)
    return path


def silent_wav(sample_rate, seconds=0.1):
    """A short silent placeholder for the vocals-stem output when the
    selected mode doesn't produce a separate vocal track."""
    return write_wav(np.zeros(int(seconds * sample_rate), dtype=np.float32), sample_rate)


@spaces.GPU(duration=180)
@torch.inference_mode()
def process_fn(mode_name, description, lyrics, ref_voice_path, extract_vocals, temperature):
    """Generate a song from a description and lyrics, optionally guided by a reference voice."""
    model, processor = get_mode(mode_name)

    model_inputs = processor(
        text=description,
        lyrics=lyrics,
        ref_voice_path=ref_voice_path if ref_voice_path else None,
        separate=extract_vocals,
    )
    generation = model.generate(**model_inputs, do_sample=True, temperature=temperature)
    sample_rate = model.config.sampling_rate

    if mode_name == "Mixed":
        song_path = write_wav(generation.cpu().numpy().squeeze(), sample_rate)
        vocals_path = silent_wav(sample_rate)
    else:
        acc_array = generation[0].cpu().numpy().squeeze()
        vocal_array = generation[1].cpu().numpy().squeeze()
        min_len = min(vocal_array.shape[0], acc_array.shape[0])
        acc_array = acc_array[:min_len]
        vocal_array = vocal_array[:min_len]

        song_path = write_wav(vocal_array + acc_array, sample_rate)
        vocals_path = write_wav(vocal_array, sample_rate)

    return song_path, vocals_path


with gr.Blocks() as demo:
    input_components = [
        gr.Dropdown(
            choices=list(MODES.keys()),
            value="Mixed",
            label="Mode",
            info="Mixed outputs one finished track; Dual-track also gives you an isolated vocal stem",
        ),
        gr.Textbox(
            label="Music Description",
            info="Genre, mood, instruments (e.g. \"upbeat pop song with piano and drums\")",
        ),
        gr.Textbox(
            label="Lyrics",
            info="English lyrics for the song (required)",
        ),
        gr.Audio(type="filepath", label="Reference Voice (optional)").harp_required(False),
        gr.Checkbox(
            value=True,
            label="Extract vocals from reference",
            info="(default: True, per repo example) turn off only if the reference audio is already an isolated vocal, not a full mix",
        ),
        gr.Slider(
            minimum=0.1, maximum=2.0, step=0.1, value=1.0,
            label="Creativity",
            info="(default: 1.0, per repo config) sampling temperature — higher is more varied/unpredictable, lower is safer/more repetitive",
        ),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Song").set_info(
            "The generated song. In Dual-track mode this is vocals + accompaniment combined."
        ),
        gr.Audio(type="filepath", label="Isolated Vocals").set_info(
            "Vocal-only stem. Only populated in Dual-track mode; silent in Mixed mode."
        ),
    ]

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

if __name__ == "__main__":
    demo.queue().launch(pwa=True)