YuE2-Jam / app.py
youngsemicolon
Claude Opus 5.5
YuE2-Jam: run m-a-p/YuE2-3B instead of MiniMax Music 3, rebrand for YuE2
de0fe6c
Raw History Blame Contribute Delete
15.4 kB
"""YuE2-Jam: describe a song in plain English, an LLM writes the lyrics and the style prompt,
YuE2 plans a score and sings it, and the song can be shared to a community feed backed by an HF
bucket mounted at /data."""
import os
# Cache dirs must be set before any HF/torch import (ZeroGPU: ~/.cache is not writable).
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
for _v in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"):
os.environ.pop(_v, None)
import spaces # noqa: E402 must precede torch
import base64 # noqa: E402
import io # noqa: E402
import json # noqa: E402
import random # noqa: E402
import subprocess # noqa: E402
import time # noqa: E402
import traceback # noqa: E402
import uuid # noqa: E402
from datetime import datetime, timezone # noqa: E402
import numpy as np # noqa: E402
import scipy.io.wavfile # noqa: E402
import torch # noqa: E402
from fastapi.responses import FileResponse, HTMLResponse # noqa: E402
import gradio as gr # noqa: E402
from gradio import Server # noqa: E402
import composer # noqa: E402
import engine # noqa: E402 (loads YuE2-3B + YuE2-Vae onto the GPU)
_HERE = os.path.dirname(os.path.abspath(__file__))
BUCKET_ID = os.environ.get("COMMUNITY_BUCKET", "Mike0021/yue2-jam-community")
BUCKET_URL = f"https://huggingface.co/buckets/{BUCKET_ID}/resolve"
DATA_DIR = os.environ.get("COMMUNITY_DIR", "/data")
SONGS_DIR = os.path.join(DATA_DIR, "songs")
MAX_SEED = int(np.iinfo(np.int32).max)
MAX_DURATION = 180.0
FEED_LIMIT = 60
AUDIO_MIME = {"mp3": "audio/mpeg", "wav": "audio/wav"}
print(f"[startup] model ready: sr={engine.SAMPLE_RATE} frame_rate={engine.FRAME_RATE}", flush=True)
print(f"[startup] community bucket: {BUCKET_ID} mounted at {DATA_DIR} (exists={os.path.isdir(DATA_DIR)})", flush=True)
# ── Cover art (Z-Image-Turbo), optional ──────────────────────────────────────
try:
from diffusers import FlowMatchEulerDiscreteScheduler, ZImagePipeline
_zimage = ZImagePipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16)
_zimage.to("cuda")
print("[startup] Z-Image-Turbo loaded for cover art", flush=True)
except Exception as e: # noqa: BLE001
_zimage = None
print(f"[startup] Z-Image-Turbo unavailable, songs will ship without cover art: {e}", flush=True)
def _render_cover(word: str) -> bytes | None:
if _zimage is None or not word:
return None
try:
t0 = time.time()
_zimage.scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
image = _zimage(
prompt=f"{word} studio photography close-up black background",
height=768, width=768,
guidance_scale=0.0,
num_inference_steps=9,
generator=torch.Generator("cuda").manual_seed(random.randint(1, 1_000_000)),
max_sequence_length=512,
).images[0]
buf = io.BytesIO()
image.save(buf, format="PNG", optimize=True)
print(f"[cover] '{word}' in {time.time() - t0:.1f}s ({len(buf.getvalue()) // 1024}KB)", flush=True)
return buf.getvalue()
except Exception as e: # noqa: BLE001
print(f"[cover] failed: {e}", flush=True)
return None
# ── GPU job: song (+ optional cover) in a single ZeroGPU allocation ───────────
def _gpu_seconds(style, lyrics, duration, seed, cover_word, *args, **kwargs):
return engine.estimate_gpu_seconds(duration) + (12 if cover_word else 0)
@spaces.GPU(duration=_gpu_seconds)
def _gpu_job(style, lyrics, duration, seed, cover_word):
# Errors are returned as data: exceptions raised in the ZeroGPU worker reach the caller
# stripped down to their class name.
try:
wav, sr, audio_seconds, wall = engine.generate_wav(style, lyrics, duration, seed)
except Exception as e: # noqa: BLE001
print(f"[gpu ERROR] {type(e).__name__}: {e}\n{traceback.format_exc()}", flush=True)
return {"error": f"{type(e).__name__}: {e}"}
print(f"[gpu] {audio_seconds:.1f}s of audio in {wall:.1f}s (seed={seed})", flush=True)
cover = _render_cover(cover_word) if cover_word else None
return {"wav": wav, "sr": sr, "seconds": audio_seconds, "cover": cover}
class GenerationError(RuntimeError):
pass
def _encode_audio(wav: np.ndarray, sr: int) -> tuple[bytes, str]:
"""MP3 at 192 kbps (a 3-minute 48 kHz stereo WAV is ~35 MB), WAV if ffmpeg fails. Returns (bytes, ext)."""
try:
out = subprocess.run(
["ffmpeg", "-hide_banner", "-loglevel", "error", "-f", "s16le", "-ar", str(sr), "-ac", str(wav.shape[1]),
"-i", "pipe:0", "-codec:a", "libmp3lame", "-b:a", "192k", "-f", "mp3", "pipe:1"],
input=np.ascontiguousarray(wav).tobytes(), capture_output=True, check=True, timeout=120,
).stdout
if out:
return out, "mp3"
except (OSError, subprocess.SubprocessError) as e:
print(f"[encode] mp3 failed, falling back to WAV: {e}", flush=True)
buf = io.BytesIO()
scipy.io.wavfile.write(buf, sr, wav)
return buf.getvalue(), "wav"
def _friendly_error(err: Exception) -> str:
msg = (str(err) or "").lower()
if any(h in msg for h in ("gpu limit", "quota", "no gpu", "could not allocate", "gpu is busy", "too many", "concurrent")):
return "The shared GPU is at capacity right now. Please wait a minute and retry."
if "out of memory" in msg or "oom" in msg:
return "Generation ran out of GPU memory. Try a shorter duration."
if isinstance(err, GenerationError):
return "Generation failed on the GPU. Please try again (a shorter duration helps if it repeats)."
if isinstance(err, (ValueError, RuntimeError)) and str(err):
return str(err)
return f"Generation failed ({type(err).__name__}). Please try again."
# ── Community feed (bucket-backed, cached in memory) ─────────────────────────
_feed: list[dict] = []
def _public_meta(meta: dict) -> dict:
keys = ("id", "title", "tags", "description", "lyrics", "style", "duration", "audio_url", "thumb_url", "created_at", "seed")
return {k: meta.get(k) for k in keys}
def _load_feed():
if not os.path.isdir(SONGS_DIR):
print("[feed] no songs dir yet, starting empty", flush=True)
return
t0 = time.time()
for song_id in os.listdir(SONGS_DIR):
meta_path = os.path.join(SONGS_DIR, song_id, "meta.json")
if not os.path.isfile(meta_path):
continue
try:
with open(meta_path) as f:
meta = json.load(f)
meta["audio_url"] = f"{BUCKET_URL}/songs/{song_id}/{meta.get('audio_file') or song_id + '.wav'}"
if os.path.isfile(os.path.join(SONGS_DIR, song_id, "thumb.png")):
meta["thumb_url"] = f"{BUCKET_URL}/songs/{song_id}/thumb.png"
else:
meta["thumb_url"] = None
_feed.append(_public_meta(meta))
except Exception as e: # noqa: BLE001
print(f"[feed] skipping {song_id}: {e}", flush=True)
_feed.sort(key=lambda s: s.get("created_at") or "", reverse=True)
print(f"[feed] loaded {len(_feed)} songs in {time.time() - t0:.1f}s", flush=True)
_load_feed()
def _share(audio: bytes, ext: str, cover: bytes | None, meta: dict) -> dict:
"""Persist a song to the bucket mount and prepend it to the in-memory feed."""
song_id = uuid.uuid4().hex[:12]
song_dir = os.path.join(SONGS_DIR, song_id)
os.makedirs(song_dir, exist_ok=True)
audio_file = f"{song_id}.{ext}"
with open(os.path.join(song_dir, audio_file), "wb") as f:
f.write(audio)
thumb_url = None
if cover:
with open(os.path.join(song_dir, "thumb.png"), "wb") as f:
f.write(cover)
thumb_url = f"{BUCKET_URL}/songs/{song_id}/thumb.png"
meta = {
**meta,
"id": song_id,
"audio_file": audio_file,
"audio_url": f"{BUCKET_URL}/songs/{song_id}/{audio_file}",
"thumb_url": thumb_url,
"created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
}
with open(os.path.join(song_dir, "meta.json"), "w") as f:
json.dump(meta, f, indent=2, ensure_ascii=False)
public = _public_meta(meta)
_feed.insert(0, public)
print(f"[share] {song_id} '{meta.get('title')}' -> {meta['audio_url']}", flush=True)
return public
# ── gr.Server app ────────────────────────────────────────────────────────────
app = Server(title="YuE2-Jam")
def _clamp_duration(d) -> float:
try:
d = float(d)
except (TypeError, ValueError):
d = 60.0
return max(10.0, min(MAX_DURATION, d))
def _resolve_seed(seed) -> int:
try:
seed = int(seed)
except (TypeError, ValueError):
seed = -1
return random.randint(0, MAX_SEED) if seed < 0 else min(seed, MAX_SEED)
def _run_song(*, description, title, tags, style, lyrics, duration, seed, community, cover_word):
"""Shared tail of /create and /generate: GPU job, encode, optional share. Returns result dict."""
engine.validate(style, lyrics)
out = _gpu_job(style, lyrics, duration, seed, cover_word)
if out.get("error"):
msg = out["error"]
if "out of memory" in msg.lower():
raise GenerationError("out of memory")
raise GenerationError(msg)
wav, sr, audio_seconds, cover = out["wav"], out["sr"], out["seconds"], out["cover"]
audio, ext = _encode_audio(wav, sr)
result = {
"title": title,
"tags": tags,
"lyrics": lyrics,
"style": style,
"seed": seed,
"duration": round(audio_seconds, 1),
"created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
"audio": f"data:{AUDIO_MIME[ext]};base64," + base64.b64encode(audio).decode(),
}
if cover:
result["thumbnail"] = "data:image/png;base64," + base64.b64encode(cover).decode()
if community:
try:
shared = _share(audio, ext, cover, {
"title": title, "tags": tags, "description": description, "lyrics": lyrics,
"style": style, "duration": round(audio_seconds, 1), "seed": seed,
})
result["community_url"] = shared["audio_url"]
result["id"] = shared["id"]
result["created_at"] = shared["created_at"]
except Exception as e: # noqa: BLE001
print(f"[share] failed: {e}", flush=True)
result["share_error"] = "Song generated, but sharing to the community feed failed."
return result
@app.api(name="create", concurrency_limit=None, time_limit=900)
def create(
description: str,
audio_duration: float = 60.0,
seed: int = -1,
community: bool = False,
instrumental: bool = False,
) -> str:
"""One-box: describe a song -> LLM writes title/style/lyrics -> YuE2 plans a score and sings it.
Streams JSON status messages ({"status": ...}) and finally the result
({audio, title, tags, lyrics, style, seed, duration, created_at, thumbnail?, community_url?})."""
try:
description = (description or "").strip()
if not description:
raise ValueError("Describe the song you want first.")
duration = _clamp_duration(audio_duration)
seed = _resolve_seed(seed)
yield json.dumps({"status": "composing", "message": "Writing lyrics & style…"})
t0 = time.time()
song = composer.compose(description, duration, instrumental=bool(instrumental))
print(f"[create] composed '{song['title']}' in {time.time() - t0:.1f}s | tags={song['tags']}", flush=True)
cover_word = composer.visual_word(song["title"], song["tags"], song["lyrics"], description) if _zimage else ""
yield json.dumps({
"status": "generating",
"message": f"YuE2 is scoring & recording β€œ{song['title']}”…",
"title": song["title"], "tags": song["tags"], "lyrics": song["lyrics"],
"eta": engine.eta_seconds(duration),
})
result = _run_song(
description=description, title=song["title"], tags=song["tags"], style=song["style"],
lyrics=song["lyrics"], duration=duration, seed=seed, community=bool(community), cover_word=cover_word,
)
yield json.dumps(result)
except Exception as e: # noqa: BLE001
print(f"[create ERROR] {type(e).__name__}: {e}\n{traceback.format_exc()}", flush=True)
raise gr.Error(_friendly_error(e), print_exception=False) from e
@app.api(name="generate", concurrency_limit=None, time_limit=900)
def generate(
style: str,
lyrics: str,
audio_duration: float = 60.0,
seed: int = -1,
title: str = "",
tags: str = "",
community: bool = False,
cover_word: str = "",
) -> str:
"""Advanced: generate from an explicit style prompt + section-tagged lyrics (no LLM).
style: comma-separated descriptors (language, genre, singer, instruments, mood, BPM), e.g. "English, warm
piano pop, expressive female voice, acoustic piano, light drums, 88 BPM". lyrics: [Verse]/[Chorus]/... tags
each on their own line; tags with no words render an instrumental. audio_duration: target length in seconds
(YuE2 sets the length from the lyrics; this caps it). cover_word: optional visual noun for the Z-Image
cover art. Returns the same JSON result shape as /create."""
try:
duration = _clamp_duration(audio_duration)
seed = _resolve_seed(seed)
lyrics = composer.normalize_lyrics(lyrics)
title = (title or "").strip()[:80] or "Untitled"
tags = (tags or "").strip()[:120]
result = _run_song(
description="", title=title, tags=tags, style=" ".join((style or "").split()), lyrics=lyrics,
duration=duration, seed=seed, community=bool(community),
cover_word=(cover_word or "").strip()[:40] if _zimage else "",
)
return json.dumps(result)
except Exception as e: # noqa: BLE001
print(f"[generate ERROR] {type(e).__name__}: {e}\n{traceback.format_exc()}", flush=True)
raise gr.Error(_friendly_error(e), print_exception=False) from e
@app.api(name="community", concurrency_limit=8)
def community() -> str:
"""Newest community songs (JSON list), served from memory. Each entry carries created_at (UTC ISO)."""
return json.dumps(_feed[:FEED_LIMIT], ensure_ascii=False)
@app.get("/", response_class=HTMLResponse)
async def homepage():
with open(os.path.join(_HERE, "index.html"), encoding="utf-8") as f:
return f.read()
@app.get("/logo.svg", include_in_schema=False)
async def logo():
return FileResponse(os.path.join(_HERE, "logo.svg"), media_type="image/svg+xml")
demo = app
if __name__ == "__main__":
# mcp_server=True exposes /create, /generate, /community as MCP tools (schema comes
# straight from each function's type hints + docstring, already written for the REST API).
demo.launch(show_error=True, mcp_server=True)