app.py

"""YuE2 Studio: lyrics-to-song demo using the original YuE2 inference package."""
import spaces
import json
import secrets
import tempfile
import time
from pathlib import Path

import gradio as gr
import torch
from yue2 import YuE2Pipeline
from yue2.modeling_vae import YuE2VAE

MODEL_ID = "m-a-p/YuE2-3B"
# YuE2 sets a process-wide allocator limit that bypasses ZeroGPU's CUDA shim.
_original_limit = torch.cuda.set_per_process_memory_fraction
try:
    torch.cuda.set_per_process_memory_fraction = lambda *args, **kwargs: None
    pipe = YuE2Pipeline.from_pretrained(MODEL_ID, device="cuda", backend="torch", progress=True)
finally:
    torch.cuda.set_per_process_memory_fraction = _original_limit
# Eagerly register both models with ZeroGPU before its startup packing phase.
pipe._load_model()
pipe._vae = YuE2VAE.from_pretrained(pipe.vae_dir, decoder_only=True, device="cpu", local_files_only=True).to("cuda")
print("YuE2 model and audio decoder loaded", flush=True)

LYRICS = """[Verse]
The city fades behind the train
A thousand lights like summer rain
I leave a little doubt behind
And chase the morning in my mind

[Chorus]
Here we go, into the blue
Every road is something new
With an open heart and room to grow
Follow where the wild winds blow

[Outro]
Into the blue, here we go"""
STYLE = "Uplifting indie pop, warm expressive vocal, bright guitars, melodic bass, driving drums, spacious chorus"

def gpu_duration(style, lyrics, mode="full", seed=-1, abc="", *args, **kwargs):
    return min(240, max(90, 90 + len(lyrics or "") // 20))

@spaces.GPU(duration=gpu_duration)
def generate(style: str, lyrics: str, mode: str = "full", seed: int = -1, abc: str = "") -> tuple:
    """Generate a YuE2 song from musical style, sectioned lyrics, and an optional edited ABC score; return audio, score, settings, and generation status."""
    style, lyrics, abc = (style or "").strip(), (lyrics or "").strip(), (abc or "").strip()
    if not style or not lyrics:
        raise gr.Error("Enter a musical style and lyrics before generating.")
    if len(style) > 1000 or len(lyrics) > 6000 or len(abc) > 50000:
        raise gr.Error("Use at most 1,000 style characters, 6,000 lyric characters, and 50,000 score characters.")
    if mode not in {"full", "melody", "off"}:
        raise gr.Error("Choose a valid score planning mode.")
    if abc and mode == "off":
        raise gr.Error("Select melody or melody + chords when supplying a score.")
    if seed is None or int(seed) != seed or not -1 <= seed <= 2147483647:
        raise gr.Error("Seed must be an integer from -1 to 2147483647.")
    used_seed = secrets.randbelow(2147483647) if int(seed) == -1 else int(seed)
    start = time.perf_counter()
    song = pipe(style=style, lyrics=lyrics, cot=mode, seed=used_seed, abc=abc or None)
    output = Path(tempfile.mkdtemp(prefix="yue2-studio-"))
    audio = output / "song.flac"
    song.save(audio)
    score = song.abc or ""
    score_path = output / "score.abc"
    score_path.write_text(score, encoding="utf-8")
    settings_path = output / "settings.json"
    settings_path.write_text(json.dumps({"model": MODEL_ID, "style": style, "lyrics": lyrics, "cot": mode, "seed": used_seed, "abc": abc or None, "truncated": song.truncated}, indent=2), encoding="utf-8")
    elapsed = time.perf_counter() - start
    status = f"Generated in {elapsed:.1f}s · seed {used_seed} · 48 kHz stereo"
    if any(song.truncated.values()):
        status += " · Model reached its length limit; try shorter lyrics."
    print(status, flush=True)
    return str(audio), score, [str(audio), str(score_path), str(settings_path)], status

CSS = """
.gradio-container {max-width:1200px!important; margin:auto!important}
#hero {padding:32px;border:1px solid #cbd5e1;border-radius:20px;background:linear-gradient(120deg,#0f172a,#312e81);margin-bottom:24px}
#hero h1 {color:#fff!important;font-size:44px!important;letter-spacing:-1.5px;margin-bottom:8px}
#hero p {color:#c7d2fe!important;font-size:16px}
#generate {min-height:52px;font-weight:700}
"""
with gr.Blocks(
    title="YuE2 Studio",
    theme=gr.themes.Soft(primary_hue="indigo"),
    css=CSS,
    delete_cache=(3600, 86400),
) as demo:
    gr.HTML('<section id="hero"><p>YuE2 / MUSIC STUDIO</p><h1>Give your words a soundtrack.</h1><p>Write the lyrics. Set the mood. Create a song with vocals and accompaniment.</p></section>')
    with gr.Row():
        with gr.Column(scale=3):
            style = gr.Textbox(label="01 / Musical direction", value=STYLE, lines=3, max_length=1000)
            lyrics = gr.Textbox(label="02 / Lyrics", value=LYRICS, lines=14, max_length=6000)
            with gr.Accordion("Score planning & reproducibility", open=False):
                mode = gr.Radio([("Melody + chords", "full"), ("Melody only", "melody"), ("Direct generation", "off")], value="full", label="Planning mode")
                seed = gr.Number(value=-1, precision=0, minimum=-1, maximum=2147483647, label="Seed (-1 for a fresh take)")
                abc = gr.Textbox(label="Optional ABC score", lines=8, info="Paste or edit a generated score to guide your next take. For melody-only covers, remove chord symbols first.")
            button = gr.Button("Generate song →", variant="primary", elem_id="generate")
        with gr.Column(scale=2):
            gr.Markdown("### 03 / Your recording\nYour generated song appears here. Queue and GPU time depend on availability.")
            audio = gr.Audio(label="Listen", type="filepath", interactive=False)
            status = gr.Textbox(label="Generation status", value="Ready for your first take", interactive=False)
            files = gr.File(label="Download audio, score & settings", file_count="multiple")
            score = gr.Textbox(label="Generated ABC score", lines=9, interactive=False)
            gr.Markdown("Copy the score into **Optional ABC score**, adjust the style or notes, then generate another take.")
    button.click(generate, [style, lyrics, mode, seed, abc], [audio, score, files, status], api_name="generate", concurrency_limit=1, concurrency_id="music")
    gr.Examples([[STYLE, LYRICS], ["Intimate acoustic folk, soft male vocal, fingerpicked guitar, warm upright bass", LYRICS]], inputs=[style, lyrics], outputs=[audio, score, files, status], fn=generate, cache_examples=True, cache_mode="lazy")
    with gr.Accordion("About this demo / ComfyUI workflow", open=False):
        gr.Markdown("This Space runs the original [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B) inference package. [Comfy-Org/YuE2](https://huggingface.co/Comfy-Org/YuE2) repackages that model for ComfyUI; its single-file checkpoint is not loaded by this Gradio app. Download the official ComfyUI workflow below to use those files in ComfyUI.\n\nModel weights: **CC BY-NC 4.0 (noncommercial)**. Outputs are AI-generated. Use lyrics and scores you have permission to use. Files are temporary and cleared after 24 hours or a restart.")
        gr.File(value="comfyui-workflow.json", label="Official ComfyUI text-to-music workflow", interactive=False)

if __name__ == "__main__":
    demo.queue(max_size=12, default_concurrency_limit=1).launch(
        mcp_server=True,
        ssr_mode=False,
    )