From 85ef0873647f1ad0409da53239f34f22b1a7d0ac Mon Sep 17 00:00:00 2001 From: Towsty Date: Mon, 14 Sep 2026 18:26:55 -0500 Subject: [PATCH] Add YuE2 host, worker, and setup scripts. Co-authored-by: Cursor --- scripts/setup-yue2.ps1 | 23 +++++++ scripts/yue2-host.mjs | 134 +++++++++++++++++++++++++++++++++++++++++ scripts/yue2-worker.py | 104 ++++++++++++++++++++++++++++++++ 3 files changed, 261 insertions(+) create mode 100644 scripts/setup-yue2.ps1 create mode 100644 scripts/yue2-host.mjs create mode 100644 scripts/yue2-worker.py diff --git a/scripts/setup-yue2.ps1 b/scripts/setup-yue2.ps1 new file mode 100644 index 0000000..61f5110 --- /dev/null +++ b/scripts/setup-yue2.ps1 @@ -0,0 +1,23 @@ +# Standalone YuE2 only. Never install into Comfy's Python environment. +[CmdletBinding()] +param( + [string]$Root = 'C:\Users\ianjm\Development\YuE2', + [string]$Uv = 'uv' +) +$ErrorActionPreference = 'Stop' +$Root = [IO.Path]::GetFullPath($Root) +function Run([string]$Exe, [string[]]$Arguments) { + & $Exe @Arguments + if ($LASTEXITCODE -ne 0) { throw "$Exe failed ($LASTEXITCODE)" } +} +if (!(Test-Path -LiteralPath $Root)) { + Run git @('clone', 'https://github.com/multimodal-art-projection/YuE.git', $Root) +} +Run $Uv @('venv', '--allow-existing', '--python', '3.12', (Join-Path $Root '.venv')) +$python = Join-Path $Root '.venv\Scripts\python.exe' +Run $Uv @('pip', 'install', '--python', $python, 'torch', 'torchaudio', '--index-url', 'https://download.pytorch.org/whl/cu128') +Run $Uv @('pip', 'install', '--python', $python, '-e', $Root) +Run $Uv @('pip', 'install', '--python', $python, 'soundfile', 'psutil') +Run $python @('-c', 'import torch; from yue2 import YuE2Pipeline; assert torch.cuda.is_available(); print(torch.__version__, torch.cuda.get_device_name(0))') +@{ model='m-a-p/YuE2-3B'; vae='m-a-p/YuE2-Vae'; python=$python } | ConvertTo-Json | Set-Content -LiteralPath (Join-Path $Root 'aigen-ready.json') -Encoding utf8 +Write-Host "YuE2 ready at $Root. Set YUE2_ROOT / YUE2_PYTHON on the GPU host if needed. Weights download from Hugging Face on first song. No Comfy environment was modified." diff --git a/scripts/yue2-host.mjs b/scripts/yue2-host.mjs new file mode 100644 index 0000000..46582d9 --- /dev/null +++ b/scripts/yue2-host.mjs @@ -0,0 +1,134 @@ +import { spawn } from 'node:child_process' +import { existsSync, mkdirSync, readFileSync, writeFileSync, renameSync, appendFileSync, readdirSync } from 'node:fs' +import { join, resolve } from 'node:path' +import { fileURLToPath } from 'node:url' + +export function validateYue2Request(body) { + const duration = body.duration ?? 60 + if (!Number.isInteger(duration) || duration < 30 || duration > 150) throw new Error('YuE2 duration must be 30–150 seconds.') + if (!Number.isInteger(body.seed) || body.seed < 0 || body.seed > 2147483647) throw new Error('Invalid seed.') + if (!/^[a-zA-Z0-9-]{12,80}$/.test(body.id || '')) throw new Error('Invalid job ID.') + const tags = String(body.tags || '').trim() + const lyrics = String(body.lyrics || '').trim() + if (!tags || tags.length > 2000 || !lyrics || lyrics.length > 8000) throw new Error('Genre tags and one lyric section are required.') + const sections = [...lyrics.matchAll(/\[([^\]]+)\]\s*([^\[]*)/gs)].filter(m => m[2].trim()) + if (sections.length !== 1) throw new Error('YuE2 requires one non-empty lyric section.') + return { id: body.id, duration, seed: body.seed, tags, lyrics } +} + +/** One isolated Python process per song. A process exit is the GPU release boundary. */ +export function createYue2Host({ prepare, leaseValid, spawnProcess = spawn, root, python, dataDir } = {}) { + const repo = resolve(root || process.env.YUE2_ROOT || 'C:\\Users\\ianjm\\Development\\YuE2') + const executable = python || process.env.YUE2_PYTHON || join(repo, '.venv', 'Scripts', 'python.exe') + const data = resolve(dataDir || process.env.YUE2_JOBS_DIR || join(repo, 'aigen-jobs')) + const worker = fileURLToPath(new URL('./yue2-worker.py', import.meta.url)) + let active = null + // A worker whose parent died exits within two seconds. Block new GPU owners + // across a host restart until that watchdog has had time to run. + const restartHoldUntil = existsSync(data) && readdirSync(data).some(id => { + try { return ['running', 'starting', 'cancelling'].includes(JSON.parse(readFileSync(join(data, id, 'status.json'), 'utf8')).status) } + catch { return false } + }) ? Date.now() + 10000 : 0 + const dir = id => { + if (!/^[a-zA-Z0-9-]{12,80}$/.test(id || '')) throw new Error('Invalid job ID.') + return join(data, id) + } + const persist = job => { + const target = join(dir(job.id), 'status.json') + writeFileSync(target + '.tmp', JSON.stringify(job)) + renameSync(target + '.tmp', target) + } + const read = id => { + if (active?.job.id === id) return { ...active.job } + const path = join(dir(id), 'status.json') + if (!existsSync(path)) return null + const job = JSON.parse(readFileSync(path, 'utf8')) + if (['running', 'starting', 'cancelling'].includes(job.status)) { + job.status = 'error'; job.error = 'YuE2 host restarted. Intermediate files are preserved.' + } + return job + } + return { + busy: () => Boolean(active) || Date.now() < restartHoldUntil, + configured: () => existsSync(executable) && existsSync(join(repo, 'aigen-ready.json')), + read, + audio(id) { + return read(id)?.status === 'complete' ? join(dir(id), 'audio.wav') : null + }, + async start(body, lease) { + const request = validateYue2Request(body) + const previous = read(request.id) + if (previous) return previous // POST retries cannot generate the same song twice. + if (active || Date.now() < restartHoldUntil) throw Object.assign(new Error('YuE2 is already running or releasing a previous worker.'), { statusCode: 409 }) + if (!existsSync(executable) || !existsSync(join(repo, 'aigen-ready.json'))) throw new Error('YuE2 is not installed. Run scripts/setup-yue2.ps1 on the GPU host.') + const job = { id: request.id, status: 'starting', message: 'Preparing GPU for YuE2', progress: 0, startedAt: Date.now(), checkedAt: Date.now() } + active = { job, child: null, cancelled: false } + const run = active + mkdirSync(dir(job.id), { recursive: true }) + persist(job) + try { + await prepare() + if (run.cancelled || !leaseValid(lease)) throw new Error('YuE2 start cancelled or GPU reservation expired.') + const installation = JSON.parse(readFileSync(join(repo, 'aigen-ready.json'), 'utf8').replace(/^\uFEFF/, '')) + writeFileSync(join(dir(job.id), 'request.json'), JSON.stringify({ + ...request, + model: process.env.YUE2_MODEL || installation.model || 'm-a-p/YuE2-3B', + vae: process.env.YUE2_VAE || installation.vae || 'm-a-p/YuE2-Vae' + })) + const child = spawnProcess(executable, ['-u', worker, '--root', repo, '--request', join(dir(job.id), 'request.json')], + { cwd: repo, windowsHide: true, shell: false, stdio: ['ignore', 'pipe', 'pipe'], env: { ...process.env, PYTHONUTF8: '1', TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD: '1' } }) + run.child = child + job.status = 'running'; job.message = 'Loading YuE2'; persist(job) + let tail = '' + let lines = '' + const log = chunk => { appendFileSync(join(dir(job.id), 'worker.log'), chunk); tail = (tail + chunk.toString()).slice(-6000) } + child.stderr.on('data', log) + child.stdout.on('data', chunk => { + log(chunk); lines += chunk.toString() + const parts = lines.split(/\r?\n/); lines = parts.pop().slice(-64000) + for (const line of parts) { + if (!line.startsWith('AIGEN_EVENT ')) continue + try { + const event = JSON.parse(line.slice(12)) + for (const key of ['stage', 'message', 'progress', 'step', 'maxStep', 'error', 'duration']) if (event[key] !== undefined) job[key] = event[key] + job.checkedAt = Date.now(); persist(job) + } catch { /* Invalid log lines are not state transitions. */ } + } + }) + const watchdog = setInterval(() => { + if (!leaseValid(lease)) { job.error = 'GPU reservation expired; YuE2 stopped.'; child.kill() } + }, 5000) + let finished = false + const finish = (code, error) => { + if (finished) return + finished = true + clearInterval(watchdog) + job.status = run.cancelled ? 'cancelled' : !error && code === 0 && existsSync(join(dir(job.id), 'audio.wav')) ? 'complete' : 'error' + if (job.status === 'error') job.error ||= error?.message || tail || `YuE2 exited with code ${code}` + job.message = job.status === 'complete' ? 'Audio ready' : job.status === 'cancelled' ? 'Cancelled' : job.error + job.checkedAt = Date.now(); persist(job) + if (active === run) active = null + } + child.once('error', error => finish(null, error)) + child.once('close', code => finish(code)) + return { ...job } + } catch (error) { + job.status = run.cancelled ? 'cancelled' : 'error'; job.error = error.message; persist(job) + if (active === run) active = null + throw error + } + }, + async cancel(id) { + if (active?.job.id !== id) return read(id) + const run = active + run.cancelled = true + run.job.status = 'cancelling' + persist(run.job) + if (run.child) { + const child = run.child + await new Promise(resolve => { child.once('close', resolve); child.kill() }) + } + return read(id) + } + } +} diff --git a/scripts/yue2-worker.py b/scripts/yue2-worker.py new file mode 100644 index 0000000..43195cd --- /dev/null +++ b/scripts/yue2-worker.py @@ -0,0 +1,104 @@ +"""Headless YuE2 adapter: plan → generate_semantic → synthesize → decode. + +No Comfy imports. No score editor, covers, or auto-fallback to YuEGP. +Stdout AIGEN_EVENT lines are consumed by the device-local host agent. +""" +import argparse +import gc +import json +import os +from pathlib import Path +import sys +import time +import threading + + +def emit(**event): + print('AIGEN_EVENT ' + json.dumps(event), flush=True) + + +def normalize_lyrics(text): + # Same one-section rule as YuEGP field validation. + import re + sections = re.findall(r'\[([^\]]+)\]\s*([^\[]*)', text.strip(), re.S) + sections = [(name, words.strip()) for name, words in sections if words.strip()] + if len(sections) != 1: + raise ValueError('YuE2 requires one non-empty lyric section. Combine the lyrics under one heading.') + name, words = sections[0] + return '[' + re.sub(r'\W+', '', name).lower() + ']\n' + words + '\n\n' + + +def main(): + import psutil + parent = psutil.Process(os.getppid()) + + def parent_watchdog(): + while parent.is_running(): + time.sleep(2) + os._exit(2) # The host died: never leave an orphan consuming the GPU. + + threading.Thread(target=parent_watchdog, daemon=True).start() + parser = argparse.ArgumentParser() + parser.add_argument('--root', required=True) + parser.add_argument('--request', required=True) + cli = parser.parse_args() + request = json.loads(Path(cli.request).read_text(encoding='utf-8-sig')) + root = Path(cli.root).resolve() + output = Path(cli.request).resolve().parent + duration = int(request.get('duration', 60)) + if not 30 <= duration <= 150: + raise ValueError('Duration must be 30–150 seconds.') + lyrics = normalize_lyrics(request['lyrics']) + style = ' '.join(str(request['tags']).split()) + seed = int(request['seed']) + model = request.get('model') or os.environ.get('YUE2_MODEL') or 'm-a-p/YuE2-3B' + vae = request.get('vae') or os.environ.get('YUE2_VAE') or 'm-a-p/YuE2-Vae' + os.chdir(root) + if str(root) not in sys.path: + sys.path.insert(0, str(root)) + emit(stage='loading', progress=1, message='Loading YuE2', model=model, vae=vae, duration=duration) + import torch + import soundfile as sf + from yue2 import YuE2Pipeline + if not torch.cuda.is_available(): + raise RuntimeError('YuE2 requires a CUDA GPU; CPU fallback is disabled.') + pipe_kwargs = dict(style=style, lyrics=lyrics, cot='full', seed=seed) + # Duration is kept for library metadata and validation. Upstream one-shot + # requests do not take a seconds field; song length follows the plan. + with YuE2Pipeline.from_pretrained(model, vae=vae, device='cuda') as pipe: + emit(stage='plan', message='Planning melody and chords', progress=5) + plan = pipe.plan(**pipe_kwargs) + emit(stage='semantic', message='Generating semantic tokens', progress=25) + semantic = pipe.generate_semantic(plan) + emit(stage='synthesize', message='Synthesizing acoustic latents', progress=55) + latents = pipe.synthesize(semantic) + emit(stage='decode', message='Decoding audio', progress=80) + audio = pipe.decode(latents) + # Context exit unloads the pipeline. Clear any residual CUDA cache. + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + wave = audio + if hasattr(audio, 'detach'): + wave = audio.detach().cpu().numpy() + import numpy as np + wave = np.asarray(wave) + if wave.ndim == 1: + pass + elif wave.shape[0] <= 8 and wave.shape[0] < wave.shape[-1]: + wave = wave.T + sample_rate = 48000 + target = output / 'audio.wav' + sf.write(str(target), wave, sample_rate, subtype='PCM_16') + info = sf.info(str(target)) + if info.frames <= 0: + raise RuntimeError('YuE2 produced empty audio.') + emit(stage='complete', message='Audio ready', progress=100, duration=info.duration) + + +if __name__ == '__main__': + try: + main() + except Exception as error: + emit(stage='error', message=str(error), error=str(error)) + raise