102 lines
5.0 KiB
JavaScript
102 lines
5.0 KiB
JavaScript
import test from 'node:test'
|
|
import assert from 'node:assert/strict'
|
|
import { EventEmitter } from 'node:events'
|
|
import { PassThrough } from 'node:stream'
|
|
import { mkdtempSync, writeFileSync, readFileSync, existsSync } from 'node:fs'
|
|
import { tmpdir } from 'node:os'
|
|
import { join } from 'node:path'
|
|
import ts from 'typescript'
|
|
import { createYue2Host, validateYue2Request } from '../scripts/yue2-host.mjs'
|
|
import { createYueGpHost, validateYueGpRequest } from '../scripts/yuegp-host.mjs'
|
|
|
|
const request = { id: 'test-song-123456', tags: 'pop, warm vocals', lyrics: '[Verse 1]\nA quiet morning, a new day', seed: 0 }
|
|
|
|
test('YuE2 defaults are 60 seconds and allow multiple lyric sections', () => {
|
|
assert.deepEqual(validateYue2Request(request), { ...request, duration: 60 })
|
|
assert.deepEqual(
|
|
validateYue2Request({ ...request, lyrics: '[Verse]\nA\n[Chorus]\nB' }).lyrics,
|
|
'[Verse]\nA\n[Chorus]\nB'
|
|
)
|
|
assert.throws(() => validateYue2Request({ ...request, lyrics: '' }))
|
|
assert.throws(() => validateYue2Request({ ...request, duration: 10 }))
|
|
assert.deepEqual(validateYue2Request({ ...request, duration: 15 }).duration, 15)
|
|
assert.throws(() => validateYue2Request({ ...request, id: '../escape' }))
|
|
})
|
|
|
|
test('yue still validates as YuEGP and stays separate from yue2', () => {
|
|
assert.equal(validateYueGpRequest(request).profile, 1)
|
|
assert.equal(validateYue2Request(request).duration, 60)
|
|
assert.notEqual(createYueGpHost.name, createYue2Host.name)
|
|
})
|
|
|
|
function fixture() {
|
|
const root = mkdtempSync(join(tmpdir(), 'aigen-yue2-'))
|
|
const python = join(root, 'python.exe')
|
|
writeFileSync(python, '')
|
|
writeFileSync(join(root, 'aigen-ready.json'), JSON.stringify({ model: 'm-a-p/YuE2-3B', vae: 'm-a-p/YuE2-Vae' }))
|
|
const child = new EventEmitter()
|
|
child.stdout = new PassThrough(); child.stderr = new PassThrough()
|
|
child.kill = () => { setImmediate(() => child.emit('close', 1)); return true }
|
|
const calls = []
|
|
const dataDir = join(root, 'jobs')
|
|
const host = createYue2Host({ root, python, dataDir, prepare: async () => {}, leaseValid: () => true,
|
|
spawnProcess: (...args) => { calls.push(args); return child } })
|
|
return { host, child, calls, root, dataDir }
|
|
}
|
|
|
|
test('host launches yue2-worker, preserves seed 0 and lyrics, and deduplicates POST retries', async () => {
|
|
const f = fixture()
|
|
await f.host.start(request, 'lease')
|
|
await f.host.start(request, 'lease')
|
|
assert.equal(f.calls.length, 1)
|
|
const [, args, opts] = f.calls[0]
|
|
assert.ok(String(args[1]).endsWith('yue2-worker.py'))
|
|
assert.equal(opts.windowsHide, true)
|
|
assert.equal(opts.shell, false)
|
|
const posted = JSON.parse(readFileSync(join(f.dataDir, request.id, 'request.json')))
|
|
assert.equal(posted.seed, 0)
|
|
assert.equal(posted.lyrics, request.lyrics)
|
|
assert.equal(posted.model, 'm-a-p/YuE2-3B')
|
|
assert.equal(posted.vae, 'm-a-p/YuE2-Vae')
|
|
await assert.rejects(f.host.start({ ...request, id: 'second-song-12345' }, 'lease'), /already running/)
|
|
await f.host.cancel(request.id)
|
|
assert.equal(f.host.busy(), false)
|
|
assert.equal(f.host.read(request.id).status, 'cancelled')
|
|
})
|
|
|
|
test('OOM stays on yue2 with no engine switch', async () => {
|
|
const f = fixture()
|
|
await f.host.start(request, 'lease')
|
|
f.child.stdout.write('AIGEN_EVENT {"stage":"error","error":"CUDA out of memory"}\n')
|
|
f.child.emit('close', 1)
|
|
assert.match(f.host.read(request.id).error, /out of memory/)
|
|
assert.equal(f.calls.length, 1)
|
|
})
|
|
|
|
test('music workflow refuses yue and yue2 while ACE graphs still build', () => {
|
|
const source = readFileSync(new URL('../server/utils/musicWorkflow.ts', import.meta.url), 'utf8')
|
|
const code = ts.transpileModule(source, { compilerOptions: { target: ts.ScriptTarget.ES2022, module: ts.ModuleKind.CommonJS } }).outputText
|
|
const exports = {}
|
|
new Function('require', 'exports', 'useRuntimeConfig', code)(id => {
|
|
if (id.endsWith('.json')) return { default: JSON.parse(readFileSync(new URL(`../server/assets/${id.split('/').pop()}`, import.meta.url))) }
|
|
if (id === '~/utils/music') return { parseMusicBpm: () => 120 }
|
|
throw new Error(id)
|
|
}, exports, () => ({}))
|
|
const params = { tags: 'pop', lyrics: '[Verse]\nhello', duration: 60, steps: 8, seed: 0, cfg: 1, lyricsStrength: 0.9 }
|
|
assert.throws(() => exports.buildMusicWorkflow({ ...params, engine: 'yue' }), /standalone YuEGP/)
|
|
assert.throws(() => exports.buildMusicWorkflow({ ...params, engine: 'yue2' }), /standalone YuE2/)
|
|
for (const engine of ['ace-step', 'ace-step-1.5']) {
|
|
const graph = exports.buildMusicWorkflow({ ...params, engine })
|
|
assert.ok(Object.values(graph).some(n => n.class_type === 'KSampler' && n.inputs.seed === 0))
|
|
}
|
|
assert.equal(existsSync(new URL('../server/assets/workflow_yue.json', import.meta.url)), false)
|
|
})
|
|
|
|
test('musicChain is YuE2-only', () => {
|
|
const source = readFileSync(new URL('../server/utils/musicChain.ts', import.meta.url), 'utf8')
|
|
assert.match(source, /startYue2Job/)
|
|
assert.match(source, /from '\.\/yue2'/)
|
|
assert.doesNotMatch(source, /startYueGpJob/)
|
|
assert.doesNotMatch(source, /buildMusicWorkflow/)
|
|
})
|