import { existsSync, mkdirSync, readdirSync, readFileSync, writeFileSync, renameSync, unlinkSync } from 'node:fs' import { join } from 'node:path' import { createJob, emitJob, getJob, restoreMusicJob, type Job } from './jobs' import { saveTrack, findYue2Track } from './library' import { sharedGpuHeaders } from './sharedGpu' import type { MusicJobParams } from './musicChain' function settings() { const config = useRuntimeConfig() return { url: String(config.comfyControlUrl || process.env.COMFY_CONTROL_URL || '').replace(/\/$/, ''), token: String(config.comfyControlToken || process.env.COMFY_CONTROL_TOKEN || '') } } export function yue2Configured() { return Boolean(settings().url) } function pendingRoot() { return join(String(useRuntimeConfig().libraryDir || process.env.LIBRARY_DIR || '/data/library'), 'yue2-pending') } export function yue2PendingAlive(jobId: string) { if (!jobId) return false return existsSync(join(pendingRoot(), `${jobId}.json`)) } async function request(path: string, body?: unknown) { const { url, token } = settings() if (!url) throw new Error('YuE2 host is not configured. Set COMFY_CONTROL_URL.') const response = await fetch(`${url}/yue2/${path}`, { method: body === undefined ? 'GET' : 'POST', headers: { ...(token ? { Authorization: `Bearer ${token}` } : {}), ...(body === undefined ? {} : { ...sharedGpuHeaders(), 'Content-Type': 'application/json' }) }, body: body === undefined ? undefined : JSON.stringify(body), signal: AbortSignal.timeout(60_000) }) if (!response.ok) { const detail = await response.json().catch(() => ({})) as { message?: string; error?: string } throw Object.assign(new Error(detail.message || detail.error || `YuE2 host returned ${response.status}`), { statusCode: response.status }) } return response } function persist(job: Job) { mkdirSync(pendingRoot(), { recursive: true }) const path = join(pendingRoot(), `${job.id}.json`) writeFileSync(path + '.tmp', JSON.stringify({ id: job.id, clientId: job.clientId, startedAt: job.startedAt, library: job.library, trackId: job.trackId })) renameSync(path + '.tmp', path) } async function settle(job: Job) { const { onLiveVideoSettled } = await import('./studioQueue') await onLiveVideoSettled(job) const path = join(pendingRoot(), `${job.id}.json`) if (existsSync(path)) unlinkSync(path) } export async function cancelYue2Job(job: Job) { await request(`jobs/${job.id}/cancel`, {}) job.status = 'cancelled' emitJob(job, { type: 'error', error: 'Cancelled', message: 'Cancelled' }) } async function watch(job: Job) { let failures = 0 while (job.status !== 'cancelled') { try { const state = await (await request(`jobs/${job.id}`)).json() as { status: string; message: string; error?: string; stage?: string; progress?: number; step?: number; maxStep?: number; duration?: number } if (job.status === 'cancelled') { await settle(job); return } job.musicActivity = { checkedAt: Date.now(), running: ['starting', 'running', 'cancelling'].includes(state.status) } if (state.status === 'error' || state.status === 'cancelled') { job.status = state.status === 'cancelled' ? 'cancelled' : 'error' job.error = state.error || state.message if (/out of memory/i.test(job.error || '')) job.error += ' YuE2 does not switch engines automatically.' emitJob(job, { type: 'error', error: job.error, message: job.error }) await settle(job) return } if (state.status === 'complete') { job.saving = true emitJob(job, { type: 'status', message: 'Saving audio to library…', progress: 98 }) const lib = job.library! lib.audioExt = 'wav' job.trackId ||= findYue2Track(lib.ownerKey, job.id)?.id if (!job.trackId) { const audio = Buffer.from(await (await request(`jobs/${job.id}/audio`)).arrayBuffer()) const track = await saveTrack({ ownerKey: lib.ownerKey, folderId: lib.folderId, name: lib.name, tags: lib.tags || lib.prompt, lyrics: lib.lyrics || '', duration: state.duration || lib.duration || 60, seed: lib.seed, steps: 0, cfg: 0, instrumental: lib.instrumental === true, engine: 'yue2', audio, ext: 'wav', comfyFilename: `yue2-${job.id}.wav` }) job.trackId = track.id persist(job) } job.status = 'complete'; job.saving = false emitJob(job, { type: 'complete', message: lib.folderLocked ? 'Saved to the locked folder.' : 'Track ready', progress: 100, trackId: job.trackId, folderLocked: lib.folderLocked, audioExt: 'wav' }) await settle(job) return } failures = 0 job.status = 'running' const percent = Number(state.progress || 0) emitJob(job, { type: 'progress', message: `${state.message}${state.maxStep ? ` · ${state.step}/${state.maxStep}` : ''}`, progress: Math.min(95, Math.max(1, percent)), step: state.step || 0, maxStep: state.maxStep || 0 }) } catch (error) { failures++ job.saving = false emitJob(job, { type: 'status', message: `YuE2 connection/save check failed; retrying: ${error instanceof Error ? error.message : String(error)}` }) if (failures >= 10) { job.status = 'error'; job.error = 'YuE2 host unreachable or audio save failed. Job files are retained for recovery.' emitJob(job, { type: 'error', error: job.error, message: job.error }) const { onLiveVideoSettled } = await import('./studioQueue') await onLiveVideoSettled(job) return } } await new Promise(resolve => setTimeout(resolve, 2000)) } await settle(job) } export function startYue2Job(params: MusicJobParams) { const job = createJob('music') job.yue2 = true job.library = { ...params, prompt: params.tags, engine: 'yue2', aspect: 'audio', width: 0, height: 0, hideThumbnail: false, turbo: false, sound: true } persist(job) setTimeout(() => { void (async () => { try { if (job.status === 'cancelled') { await settle(job); return } await request('jobs', { id: job.id, tags: params.tags, lyrics: params.lyrics, duration: params.duration, seed: params.seed }) await watch(job) } catch (error) { const statusCode = (error as { statusCode?: number })?.statusCode if (statusCode && statusCode >= 400 && statusCode < 500) { job.status = 'error'; job.error = error instanceof Error ? error.message : String(error) emitJob(job, { type: 'error', error: job.error, message: job.error }) await settle(job) return } emitJob(job, { type: 'status', message: `Checking YuE2 submission: ${error instanceof Error ? error.message : String(error)}` }) await watch(job) } })() }, 0) return job } export function resumeYue2Jobs() { if (!existsSync(pendingRoot())) return for (const file of readdirSync(pendingRoot()).filter(file => /^[a-zA-Z0-9-]+\.json$/.test(file))) { try { const record = JSON.parse(readFileSync(join(pendingRoot(), file), 'utf8')) if (getJob(record.id)) continue const job = restoreMusicJob(record) job.yueGp = false job.yue2 = true job.message = 'Reconnecting to YuE2' void watch(job) } catch { /* Preserve invalid records for diagnosis. */ } } }