Blame for app.py

"""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)