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)