aigen/scripts/yue2-host.mjs

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 non-empty lyrics are required.')
  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/, ''))
        const localModel = process.env.YUE2_MODEL || installation.model || 'm-a-p/YuE2-3B'
        const localVae = process.env.YUE2_VAE || installation.vae || 'm-a-p/YuE2-Vae'
        writeFileSync(join(dir(job.id), 'request.json'), JSON.stringify({
          ...request,
          model: localModel,
          vae: localVae
        }))
        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)
    }
  }
}