youngsemicolon
Claude Opus 5.5
YuE2-Jam: run m-a-p/YuE2-3B instead of MiniMax Music 3, rebrand for YuE2
de0fe6c | """YuE2-3B inference engine for ZeroGPU. | |
| Wraps the official `yue2_infer` package shipped in the m-a-p/YuE2-3B repo: one AR/NAR backbone | |
| plans a chord-annotated ABC score, writes 25 Hz semantic tokens, turns them into acoustic latents | |
| by flow matching, and YuE2-Vae decodes 48 kHz stereo. Everything here runs at import time (model | |
| load) except `generate_wav`, which must be called from inside a `@spaces.GPU` function. | |
| """ | |
| import re | |
| import subprocess | |
| import sys | |
| import time | |
| import spaces # noqa: F401 must precede torch (ZeroGPU patches CUDA init) | |
| import numpy as np | |
| import torch | |
| MODEL_ID = "m-a-p/YuE2-3B" | |
| VAE_ID = "m-a-p/YuE2-Vae" | |
| # The inference package lives in the model repo. Its metadata hard-pins huggingface-hub 0.x, which | |
| # Gradio 6 cannot run with, so it goes in without dependencies (the real ones are in requirements.txt). | |
| WHEEL_URL = f"https://huggingface.co/{MODEL_ID}/resolve/main/yue2_infer-0.1.5-py3-none-any.whl" | |
| try: | |
| import yue2 # noqa: F401 | |
| except ImportError: | |
| subprocess.run([sys.executable, "-m", "pip", "install", "--no-deps", "-q", WHEEL_URL], check=True) | |
| from yue2 import YuE2Pipeline # noqa: E402 | |
| from yue2.modeling_vae import YuE2VAE # noqa: E402 | |
| from yue2.modeling_yue2 import YuE2ForCausalLM # noqa: E402 | |
| from yue2.protocol import CONTEXT, SongRequest, token_prefixes # noqa: E402 | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # YuE2Pipeline.__init__ caps the process at a 24 GiB memory fraction, which forces a real CUDA init | |
| # under ZeroGPU's startup emulation; skip only that cap. | |
| _set_memory_fraction = torch.cuda.set_per_process_memory_fraction | |
| torch.cuda.set_per_process_memory_fraction = lambda *args, **kwargs: None | |
| try: | |
| PIPE = YuE2Pipeline.from_pretrained(MODEL_ID, vae=VAE_ID, device=DEVICE, progress=False) | |
| finally: | |
| torch.cuda.set_per_process_memory_fraction = _set_memory_fraction | |
| # Weights are loaded here rather than lazily by the pipeline: transformers 5 takes `dtype` (the | |
| # package passes `torch_dtype`), and the VAE stays resident next to the backbone, whereas the | |
| # package's decode() moves the 7 GB backbone to the CPU and back for every song. | |
| PIPE._model = YuE2ForCausalLM.from_pretrained( | |
| PIPE.model_dir, local_files_only=True, dtype=torch.bfloat16 | |
| ).eval().to(DEVICE) | |
| VAE = YuE2VAE.from_pretrained(PIPE.vae_dir, decoder_only=True, device=DEVICE, local_files_only=True) | |
| SAMPLE_RATE = int(VAE.config.sample_rate) | |
| FRAME_RATE = SAMPLE_RATE / VAE.config.downsampling_ratio # semantic tokens per second of audio | |
| ABC_MAX_TOKENS = PIPE.generation_config.abc.max_tokens | |
| ODE_STEPS = PIPE.generation_config.ode_steps | |
| MAX_AUDIO_SECONDS = 240.0 | |
| MAX_STYLE_CHARS, MAX_LYRICS_CHARS = 1500, 10000 | |
| FADE_SECONDS = 3.0 | |
| _TAG_LINE = re.compile(r"^\s*\[[^\]]*\]\s*$") | |
| def is_instrumental(lyrics: str) -> bool: | |
| """True when the lyrics hold only section tags, i.e. there are no words to sing.""" | |
| return not any(line.strip() and not _TAG_LINE.match(line) for line in (lyrics or "").splitlines()) | |
| def _cot(lyrics: str) -> str: | |
| # Full planning writes a melody+chords score whose lead is a Vocal voice, which YuE2 then sings; | |
| # instrumentals skip the symbolic plan. | |
| return "off" if is_instrumental(lyrics) else "full" | |
| def max_audio_seconds(duration: float) -> float: | |
| """Hard cap on the rendered length. YuE2 picks the song length from the lyrics (the composer | |
| sizes them to `duration`); the cap only bounds GPU time, with headroom for intros and outros.""" | |
| return min(float(duration) * 1.3 + 15.0, MAX_AUDIO_SECONDS) | |
| def _gpu_time(audio_seconds: float, abc_tokens: float) -> float: | |
| # Fitted on ZeroGPU with this same package (hugging-apps/yue2-instrumental-cot-lora): ~104 tok/s | |
| # while planning, ~0.265 s per audio-second for semantic decoding + VAE, ~0.0023 s per | |
| # audio-second per ODE step, ~11 s fixed (worker start + VAE). | |
| return 11.0 + abc_tokens / 104.0 + audio_seconds * (0.265 + 0.0023 * ODE_STEPS) | |
| def estimate_gpu_seconds(duration: float) -> int: | |
| """ZeroGPU reservation: worst case (ABC plan at its token cap, song at its length cap) + 10%.""" | |
| return int(min(_gpu_time(max_audio_seconds(duration), ABC_MAX_TOKENS) * 1.1, 300)) | |
| def eta_seconds(duration: float) -> int: | |
| """Typical wall time for a song of ~`duration` seconds (plans run 1.6k-2.7k ABC tokens).""" | |
| return int(_gpu_time(float(duration) * 1.1, 2200)) | |
| def validate(style: str, lyrics: str) -> None: | |
| """CPU-side input validation (raise ValueError with a user-facing message). | |
| Runs before the ZeroGPU call: exceptions raised inside the GPU worker lose their message. | |
| """ | |
| style = (style or "").strip() | |
| lyrics = (lyrics or "").strip() | |
| if not style: | |
| raise ValueError("A style description is required.") | |
| if not lyrics: | |
| raise ValueError("Lyrics are required (section tags only, e.g. [Instrumental], make an instrumental).") | |
| if len(style) > MAX_STYLE_CHARS: | |
| raise ValueError(f"Keep the style under {MAX_STYLE_CHARS} characters.") | |
| if len(lyrics) > MAX_LYRICS_CHARS: | |
| raise ValueError(f"Keep the lyrics under {MAX_LYRICS_CHARS} characters.") | |
| prompt_tokens = len(token_prefixes(SongRequest(style=style, lyrics=lyrics, cot=_cot(lyrics)), PIPE.tokenizer)) | |
| if prompt_tokens + ABC_MAX_TOKENS + 3 + int(MAX_AUDIO_SECONDS * FRAME_RATE) > CONTEXT: | |
| raise ValueError("The lyrics are too long for YuE2's context window; please shorten them.") | |
| def generate_wav(style: str, lyrics: str, duration: float, seed: int): | |
| """Run plan -> semantic tokens -> flow matching -> VAE. Must run inside @spaces.GPU. | |
| Returns (int16 stereo ndarray [samples, 2], sample_rate, audio_seconds, wall_seconds). | |
| """ | |
| style, lyrics = (style or "").strip(), (lyrics or "").strip() | |
| start = time.time() | |
| plan = PIPE.plan(style=style, lyrics=lyrics, cot=_cot(lyrics), seed=int(seed)) | |
| max_tokens = int(max_audio_seconds(duration) * FRAME_RATE) | |
| semantic = PIPE.generate_semantic(plan, sampling={"max_tokens": max_tokens}) | |
| latents = PIPE.synthesize(semantic) | |
| audio = VAE.decode_tiled( | |
| torch.from_numpy(latents).T.contiguous().unsqueeze(0), core_frames=1024, halo_frames=16, output_device="cpu" | |
| ) | |
| wav = audio[0].clamp(-1.0, 1.0).T.contiguous().numpy() | |
| if semantic.truncated: # stopped at the length cap: fade out rather than cut mid-phrase | |
| n = min(len(wav), int(FADE_SECONDS * SAMPLE_RATE)) | |
| wav[-n:] *= np.linspace(1.0, 0.0, n, dtype=np.float32)[:, None] | |
| print( | |
| f"[engine] cot={plan.request.cot} abc_tokens={len(plan.abc_ids)} (truncated={plan.truncated}) " | |
| f"semantic_tokens={len(semantic.tokens)}/{max_tokens} (truncated={semantic.truncated}) " | |
| f"abc={plan.timing.get('seconds', 0):.1f}s semantic={semantic.timing['seconds']:.1f}s", | |
| flush=True, | |
| ) | |
| return (wav * 32767.0).astype(np.int16), SAMPLE_RATE, len(wav) / SAMPLE_RATE, time.time() - start | |