Add ACE-Step music generation to the studio and library.
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
@@ -0,0 +1,352 @@
|
||||
import { createJob, emitJob, type Job } from '~/server/utils/jobs'
|
||||
import { extractAudio, fetchHistory, queuePrompt } from '~/server/utils/comfy'
|
||||
import { comfyWsUrl } from '~/server/utils/comfy'
|
||||
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
|
||||
import { downloadComfyAudio, saveTrack } from '~/server/utils/library'
|
||||
import { buildMusicWorkflow, musicFilenamePrefix } 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('cuda') && lower.includes('alloc')) || lower.includes('vram')) {
|
||||
return 'ComfyUI VRAM allocation failed. Try fewer steps or a shorter duration.'
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
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 })
|
||||
resolve()
|
||||
}
|
||||
|
||||
const succeed = async () => {
|
||||
if (settled || finishing || !job.promptId) return false
|
||||
finishing = true
|
||||
job.saving = true
|
||||
try {
|
||||
const history = await fetchHistory(job.promptId)
|
||||
const audio = extractAudio(history, job.promptId)
|
||||
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
|
||||
}
|
||||
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 })
|
||||
}
|
||||
job.saving = false
|
||||
resolve()
|
||||
return true
|
||||
} catch (error) {
|
||||
finishing = false
|
||||
job.saving = false
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
const pollHistory = async () => {
|
||||
if (settled || finishing || !job.promptId) return
|
||||
try {
|
||||
const history = await fetchHistory(job.promptId)
|
||||
if (extractAudio(history, job.promptId)) {
|
||||
await succeed()
|
||||
return
|
||||
}
|
||||
const entry = (history?.[job.promptId] || history) as {
|
||||
status?: { status_str?: string; completed?: boolean }
|
||||
} | undefined
|
||||
const status = entry?.status?.status_str
|
||||
if (status === 'error') await fail('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 || {}
|
||||
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 && job.promptId) 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 = 'uploading'
|
||||
emitJob(job, { type: 'status', message: 'Queueing ACE-Step...', progress: 6 })
|
||||
await waitForSocket(job, 4000)
|
||||
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
|
||||
})
|
||||
const queued = await queuePrompt(graph, job.clientId)
|
||||
job.promptId = queued.prompt_id
|
||||
job.status = 'running'
|
||||
emitJob(job, { type: 'status', message: 'Job queued on ComfyUI', 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 })
|
||||
}
|
||||
}
|
||||
})()
|
||||
|
||||
return job
|
||||
}
|
||||
Reference in New Issue
Block a user