439 lines
15 KiB
TypeScript
439 lines
15 KiB
TypeScript
import { createJob, emitJob, type Job } from '~/server/utils/jobs'
|
||
import { extractAudio, fetchHistory, fetchHistoryAll, findHistoryAudio, freeComfyVram, purgeComfyArtifacts, queuePrompt } from '~/server/utils/comfy'
|
||
import { comfyWsUrl, comfyFetch } from '~/server/utils/comfy'
|
||
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
|
||
import { downloadComfyAudio, saveTrack } from '~/server/utils/library'
|
||
import { buildMusicWorkflow, musicFilenamePrefix, assertMusicEngineNodes } from '~/server/utils/musicWorkflow'
|
||
import { NODE_LABELS } from '~/server/utils/workflow'
|
||
import { musicExtension } from '~/utils/music'
|
||
import type { MusicEngine } from '~/utils/music'
|
||
|
||
export type MusicJobParams = {
|
||
ownerKey: string
|
||
folderId: string
|
||
name: string
|
||
tags: string
|
||
lyrics: string
|
||
duration: number
|
||
steps: number
|
||
seed: number
|
||
cfg: number
|
||
lyricsStrength: number
|
||
instrumental: boolean
|
||
folderLocked?: boolean
|
||
engine?: MusicEngine
|
||
samplerName?: string
|
||
scheduler?: string
|
||
}
|
||
|
||
function sleep(ms: number) {
|
||
return new Promise(resolve => setTimeout(resolve, ms))
|
||
}
|
||
|
||
function classifyError(message: string) {
|
||
const lower = message.toLowerCase()
|
||
if (lower.includes('out of memory') || lower.includes('exceed allowed memory') || (lower.includes('cuda') && lower.includes('alloc')) || lower.includes('vram')) {
|
||
return 'ComfyUI ran out of VRAM on YuE. Sleep→Poke Comfy, don’t run video/image first, use one short [Chorus] only — or use ACE-Step.'
|
||
}
|
||
if (lower.includes('econnrefused') || lower.includes('unreachable') || lower.includes('fetch failed')) {
|
||
return 'ComfyUI host connection dropped. Confirm the desktop instance is running.'
|
||
}
|
||
if (lower.includes('timeout')) {
|
||
return 'Network timeout talking to ComfyUI. The job may still be running on the desktop.'
|
||
}
|
||
return message
|
||
}
|
||
|
||
async function settleMusicStudio(job: Job) {
|
||
try {
|
||
const { onLiveVideoSettled } = await import('~/server/utils/studioQueue')
|
||
await onLiveVideoSettled(job)
|
||
} catch {
|
||
/* queue kick is best-effort */
|
||
}
|
||
}
|
||
|
||
function nodeLabel(node: string) {
|
||
if (!node) return ''
|
||
return NODE_LABELS[node] || ''
|
||
}
|
||
|
||
async function waitForSocket(job: Job, ms = 4000) {
|
||
const started = Date.now()
|
||
while (Date.now() - started < ms) {
|
||
if (job.socketReady) return
|
||
await sleep(100)
|
||
}
|
||
}
|
||
|
||
function watchMusicJob(job: Job): Promise<void> {
|
||
job.socketReady = false
|
||
job.saving = false
|
||
const ws = new WebSocket(comfyWsUrl(job.clientId))
|
||
let settled = false
|
||
let finishing = false
|
||
let pollTimer: ReturnType<typeof setInterval> | null = null
|
||
let timeout: ReturnType<typeof setTimeout> | null = null
|
||
|
||
return new Promise((resolve) => {
|
||
let localProgress = 0
|
||
const startedAt = Date.now()
|
||
let lastActivity = Date.now()
|
||
const IDLE_MS = 30 * 60 * 1000
|
||
const ABSOLUTE_MS = 3 * 60 * 60 * 1000
|
||
|
||
const markActivity = () => { lastActivity = Date.now() }
|
||
|
||
const emitLocal = (event: Parameters<typeof emitJob>[1], samplePct?: number) => {
|
||
markActivity()
|
||
if (typeof event.progress === 'number') localProgress = event.progress
|
||
if (event.message && typeof samplePct === 'number') {
|
||
event.message = `${event.message} (${samplePct}%)`
|
||
}
|
||
emitJob(job, event)
|
||
}
|
||
|
||
const cleanup = () => {
|
||
if (timeout) clearTimeout(timeout)
|
||
if (pollTimer) clearInterval(pollTimer)
|
||
timeout = null
|
||
pollTimer = null
|
||
try { ws.close() } catch { /* ignore */ }
|
||
}
|
||
|
||
const armIdle = () => {
|
||
if (timeout) clearTimeout(timeout)
|
||
if (settled) return
|
||
const absLeft = ABSOLUTE_MS - (Date.now() - startedAt)
|
||
if (absLeft <= 0) {
|
||
void fail('Timed out waiting for ComfyUI (3 hours).')
|
||
return
|
||
}
|
||
const idleLeft = IDLE_MS - (Date.now() - lastActivity)
|
||
timeout = setTimeout(() => {
|
||
if (settled || finishing) return
|
||
if (Date.now() - lastActivity >= IDLE_MS) {
|
||
void fail('Timed out waiting for ComfyUI progress (30 minutes with no updates).')
|
||
return
|
||
}
|
||
armIdle()
|
||
}, Math.max(1000, Math.min(idleLeft, absLeft)))
|
||
}
|
||
|
||
const fail = async (error: string) => {
|
||
if (settled || finishing) return
|
||
settled = true
|
||
cleanup()
|
||
job.status = job.status === 'cancelled' ? 'cancelled' : 'error'
|
||
job.error = classifyError(error)
|
||
emitJob(job, { type: 'error', error: job.error, message: job.error })
|
||
await settleMusicStudio(job)
|
||
resolve()
|
||
}
|
||
|
||
const adoptPromptId = (id: unknown) => {
|
||
const promptId = String(id || '').trim()
|
||
if (promptId && !job.promptId) job.promptId = promptId
|
||
}
|
||
|
||
const loadJobAudio = async () => {
|
||
if (job.promptId) {
|
||
const scoped = await fetchHistory(job.promptId)
|
||
const fromPrompt = extractAudio(scoped, job.promptId)
|
||
if (fromPrompt) return fromPrompt
|
||
}
|
||
const found = findHistoryAudio(await fetchHistoryAll(), {
|
||
promptId: job.promptId,
|
||
tags: job.library?.tags,
|
||
since: job.startedAt - 60_000
|
||
})
|
||
if (found) {
|
||
job.promptId = found.promptId
|
||
return found.audio
|
||
}
|
||
return null
|
||
}
|
||
|
||
const succeed = async () => {
|
||
if (settled || finishing) return false
|
||
finishing = true
|
||
job.saving = true
|
||
try {
|
||
const audio = await loadJobAudio()
|
||
if (!audio) {
|
||
finishing = false
|
||
job.saving = false
|
||
return false
|
||
}
|
||
if (settled) return false
|
||
settled = true
|
||
cleanup()
|
||
try {
|
||
job.audio = audio
|
||
emitLocal({ type: 'status', message: 'Saving to library...', progress: 98 })
|
||
if (job.library) {
|
||
const buffer = await downloadComfyAudio(audio)
|
||
const ext = musicExtension(audio.filename)
|
||
const track = await saveTrack({
|
||
ownerKey: job.library.ownerKey,
|
||
folderId: job.library.folderId,
|
||
name: job.library.name,
|
||
tags: job.library.tags || job.library.prompt,
|
||
lyrics: job.library.lyrics || '',
|
||
duration: job.library.duration || 0,
|
||
steps: job.library.steps,
|
||
seed: job.library.seed,
|
||
cfg: job.library.cfg || 0,
|
||
instrumental: job.library.instrumental === true,
|
||
audio: buffer,
|
||
ext,
|
||
engine: job.library.engine,
|
||
samplerName: job.library.samplerName,
|
||
scheduler: job.library.scheduler,
|
||
lyricsStrength: job.library.lyricsStrength,
|
||
comfyFilename: audio.filename
|
||
})
|
||
job.trackId = track.id
|
||
job.library.audioExt = ext
|
||
await purgeComfyArtifacts({
|
||
audio,
|
||
promptId: job.promptId
|
||
})
|
||
}
|
||
job.status = 'complete'
|
||
emitJob(job, {
|
||
type: 'complete',
|
||
message: job.library?.folderLocked
|
||
? 'Saved to the locked folder. Unlock it to listen.'
|
||
: 'Track ready',
|
||
progress: 100,
|
||
filename: job.library?.folderLocked ? undefined : audio.filename,
|
||
subfolder: job.library?.folderLocked ? undefined : audio.subfolder,
|
||
mediaType: job.library?.folderLocked ? undefined : audio.type,
|
||
trackId: job.trackId,
|
||
folderLocked: job.library?.folderLocked
|
||
})
|
||
} catch (saveError) {
|
||
const message = saveError instanceof Error ? saveError.message : String(saveError)
|
||
job.status = 'error'
|
||
job.error = `Audio generated but library save failed: ${message}`
|
||
emitJob(job, { type: 'error', error: job.error, message: job.error })
|
||
}
|
||
await settleMusicStudio(job)
|
||
job.saving = false
|
||
resolve()
|
||
return true
|
||
} catch (error) {
|
||
finishing = false
|
||
job.saving = false
|
||
throw error
|
||
}
|
||
}
|
||
|
||
const pollHistory = async () => {
|
||
if (settled || finishing) return
|
||
if (job.promptId) {
|
||
try {
|
||
const response = await comfyFetch('/queue', { signal: AbortSignal.timeout(2500) })
|
||
if (response.ok) {
|
||
const queue = await response.json() as { queue_running?: unknown[][] }
|
||
job.musicActivity = {
|
||
checkedAt: Date.now(),
|
||
running: Boolean(queue.queue_running?.some(row => row[1] === job.promptId))
|
||
}
|
||
}
|
||
} catch { /* Keep the last confirmation timestamp so the UI can show stale checks. */ }
|
||
}
|
||
if (settled || finishing) return
|
||
try {
|
||
const history = await fetchHistoryAll()
|
||
const found = findHistoryAudio(history, {
|
||
promptId: job.promptId,
|
||
tags: job.library?.tags,
|
||
since: job.startedAt - 60_000
|
||
})
|
||
if (found) {
|
||
job.promptId = found.promptId
|
||
await succeed()
|
||
return
|
||
}
|
||
const entry = (job.promptId ? history?.[job.promptId] : undefined) as {
|
||
status?: {
|
||
status_str?: string
|
||
completed?: boolean
|
||
messages?: Array<[string, Record<string, unknown>]>
|
||
}
|
||
} | undefined
|
||
const status = entry?.status?.status_str
|
||
if (status === 'error') {
|
||
const err = entry?.status?.messages?.find(([type]) => type === 'execution_error')?.[1]
|
||
const detail = String(err?.exception_message || err?.message || '').trim()
|
||
await fail(detail || 'ComfyUI reported an execution error')
|
||
} else if (status === 'interrupted') await fail('Job interrupted.')
|
||
} catch {
|
||
// history can miss while the graph is still running
|
||
}
|
||
}
|
||
|
||
armIdle()
|
||
pollTimer = setInterval(() => { void pollHistory() }, 4000)
|
||
|
||
ws.addEventListener('open', () => {
|
||
job.socketReady = true
|
||
emitLocal({ type: 'status', message: 'Connected to ComfyUI', progress: Math.max(localProgress, 4) })
|
||
})
|
||
ws.addEventListener('error', () => {
|
||
job.socketReady = false
|
||
})
|
||
ws.addEventListener('message', (event) => {
|
||
let payload: { type?: string; data?: Record<string, any> }
|
||
try {
|
||
payload = JSON.parse(String(event.data))
|
||
} catch {
|
||
return
|
||
}
|
||
const type = String(payload.type || '')
|
||
const data = payload.data || {}
|
||
adoptPromptId(data.prompt_id)
|
||
if (type === 'status') markActivity()
|
||
if (type === 'progress') {
|
||
const value = Number(data.value || 0)
|
||
const max = Number(data.max || 0)
|
||
const pct = max > 0 ? Math.min(95, Math.round((value / max) * 90) + 8) : Math.max(localProgress, 12)
|
||
const node = String(data.node || '')
|
||
emitLocal({
|
||
type: 'progress',
|
||
message: nodeLabel(node) || 'Sampling',
|
||
progress: pct,
|
||
step: value,
|
||
maxStep: max,
|
||
node
|
||
}, max > 0 ? Math.round((value / max) * 100) : undefined)
|
||
}
|
||
if (type === 'executing') {
|
||
const node = data.node == null ? '' : String(data.node)
|
||
if (!node) void succeed()
|
||
else if (node) {
|
||
emitLocal({
|
||
type: 'status',
|
||
message: nodeLabel(node) || 'Running node',
|
||
progress: Math.max(localProgress, 10),
|
||
node
|
||
})
|
||
}
|
||
}
|
||
if (type === 'executed') {
|
||
void succeed()
|
||
}
|
||
if (type === 'execution_error') {
|
||
const message = String(data.exception_message || data.message || 'ComfyUI node execution failed')
|
||
void fail(message)
|
||
}
|
||
if (type === 'execution_interrupted') {
|
||
job.status = 'cancelled'
|
||
void fail('Job interrupted.')
|
||
}
|
||
})
|
||
})
|
||
}
|
||
|
||
export async function startMusicJob(params: MusicJobParams) {
|
||
const job = createJob('music')
|
||
job.library = {
|
||
ownerKey: params.ownerKey,
|
||
folderId: params.folderId,
|
||
hideThumbnail: false,
|
||
folderLocked: params.folderLocked === true,
|
||
name: params.name,
|
||
prompt: params.tags,
|
||
tags: params.tags,
|
||
lyrics: params.lyrics,
|
||
instrumental: params.instrumental,
|
||
aspect: 'audio',
|
||
width: 0,
|
||
height: 0,
|
||
steps: params.steps,
|
||
turbo: false,
|
||
seed: params.seed,
|
||
cfg: params.cfg,
|
||
duration: params.duration,
|
||
sound: true,
|
||
engine: params.engine || 'ace-step',
|
||
samplerName: params.samplerName || 'euler',
|
||
scheduler: params.scheduler || 'simple',
|
||
lyricsStrength: params.lyricsStrength
|
||
}
|
||
job.maxStep = params.steps
|
||
job.status = 'queued'
|
||
emitJob(job, { type: 'status', message: 'Checking ComfyUI...', progress: 1 })
|
||
|
||
void (async () => {
|
||
try {
|
||
const done = watchMusicJob(job)
|
||
await ensureComfyReady((status) => {
|
||
emitJob(job, {
|
||
type: status.state === 'busy' ? 'busy' : 'status',
|
||
message: status.message,
|
||
progress: status.state === 'online' ? 3 : 1,
|
||
busy: status.state === 'busy',
|
||
queueRunning: status.queueRunning,
|
||
queuePending: status.queuePending
|
||
})
|
||
})
|
||
job.status = 'running'
|
||
const engineLabel = params.engine === 'yue' ? 'YuE' : params.engine === 'ace-step-1.5' ? 'ACE-Step 1.5' : 'ACE-Step'
|
||
emitJob(job, { type: 'status', message: `Queueing ${engineLabel}...`, progress: 6 })
|
||
await waitForSocket(job, 4000)
|
||
await assertMusicEngineNodes(params.engine)
|
||
if (params.engine === 'yue') {
|
||
emitJob(job, { type: 'status', message: 'Freeing Comfy VRAM for YuE…', progress: 5 })
|
||
await freeComfyVram()
|
||
}
|
||
const graph = buildMusicWorkflow({
|
||
tags: params.tags,
|
||
lyrics: params.instrumental ? '' : params.lyrics,
|
||
duration: params.duration,
|
||
steps: params.steps,
|
||
seed: params.seed,
|
||
cfg: params.cfg,
|
||
lyricsStrength: params.instrumental ? 0 : params.lyricsStrength,
|
||
samplerName: params.samplerName,
|
||
scheduler: params.scheduler,
|
||
filenamePrefix: musicFilenamePrefix(),
|
||
engine: params.engine,
|
||
vocalMultiplier: params.instrumental ? 0.2 : 1
|
||
})
|
||
try {
|
||
const queued = await queuePrompt(graph, job.clientId, { timeoutMs: 25_000 })
|
||
if (queued?.prompt_id) job.promptId = queued.prompt_id
|
||
const runningMessage = params.engine === 'yue'
|
||
? 'YuE Stage A on Comfy — 16GB job, often 10–20+ minutes. Leave it alone.'
|
||
: 'Job queued on ComfyUI'
|
||
emitJob(job, { type: 'status', message: runningMessage, progress: 8 })
|
||
} catch (queueError: unknown) {
|
||
const statusCode = Number((queueError as { statusCode?: number; status?: number })?.statusCode
|
||
|| (queueError as { status?: number })?.status || 0)
|
||
if (statusCode >= 400) throw queueError
|
||
emitJob(job, {
|
||
type: 'status',
|
||
message: params.engine === 'yue'
|
||
? 'YuE submitted — waiting for Comfy to start Stage A…'
|
||
: 'Waiting for ComfyUI to pick up the job...',
|
||
progress: 8
|
||
})
|
||
}
|
||
await done
|
||
} catch (error) {
|
||
const message = error instanceof Error ? error.message : String(error)
|
||
if (job.status !== 'cancelled' && job.status !== 'complete') {
|
||
job.status = 'error'
|
||
job.error = classifyError(message)
|
||
emitJob(job, { type: 'error', error: job.error, message: job.error })
|
||
await settleMusicStudio(job)
|
||
}
|
||
}
|
||
})()
|
||
|
||
return job
|
||
}
|