Add ACE-Step 1.5 and YuE on Music, and put Video/Image/Music cards back.

Recover finished Comfy audio when queueing hangs so a completed track is not stuck on Queueing.

Co-authored-by: Cursor <cursoragent@cursor.com>

server/utils/comfy.ts

@@ -171,11 +171,12 @@ function comfyQueueMessage(payload: unknown, status: number) {
   return message || `Queue failed (${status})`
 }
 
-export async function queuePrompt(graph: unknown, clientId: string) {
+export async function queuePrompt(graph: unknown, clientId: string, opts?: { timeoutMs?: number }) {
   const res = await comfyFetch('/prompt', {
     method: 'POST',
     headers: { 'Content-Type': 'application/json' },
-    body: JSON.stringify({ prompt: graph, client_id: clientId })
+    body: JSON.stringify({ prompt: graph, client_id: clientId }),
+    ...(opts?.timeoutMs ? { signal: AbortSignal.timeout(opts.timeoutMs) } : {})
   })
   const payload = await res.json().catch(() => ({}))
   if (!res.ok) {
@@ -295,6 +296,111 @@ export function extractAudio(history: Record<string, unknown> | null, promptId:
   return findAudio(history)
 }
 
+export function comfyAudioPrefix() {
+  return String(useRuntimeConfig().comfyAudioPrefix || process.env.COMFY_AUDIO_PREFIX || 'audio/AIGen')
+}
+
+export function isOurComfyAudio(audio: { filename: string; subfolder?: string }) {
+  const prefix = comfyAudioPrefix().replace(/\/$/, '')
+  const parts = prefix.split('/')
+  const namePrefix = parts[parts.length - 1] || 'AIGen'
+  const sub = parts.length > 1 ? parts.slice(0, -1).join('/') : 'audio'
+  const nameOk = String(audio.filename || '').startsWith(namePrefix)
+  const subOk = !audio.subfolder || audio.subfolder === sub
+  return nameOk && subOk
+}
+
+export function extractAceMetaFromHistory(entry: unknown) {
+  const prompt = (entry as { prompt?: unknown[] })?.prompt
+  const graph = Array.isArray(prompt) ? prompt[2] : null
+  const meta = {
+    tags: '',
+    lyrics: '',
+    duration: 0,
+    steps: 0,
+    seed: 0,
+    cfg: 0,
+    lyricsStrength: 0.9,
+    samplerName: '',
+    scheduler: '',
+    engine: 'ace-step'
+  }
+  if (!graph || typeof graph !== 'object') return meta
+  for (const node of Object.values(graph as Record<string, { class_type?: string; inputs?: Record<string, unknown> }>)) {
+    const inputs = node?.inputs || {}
+    if (node?.class_type === 'TextEncodeAceStepAudio' || node?.class_type === 'TextEncodeAceStepAudio1.5') {
+      meta.tags = String(inputs.tags || '')
+      meta.lyrics = String(inputs.lyrics || '')
+      meta.lyricsStrength = Number(inputs.lyrics_strength || meta.lyricsStrength)
+    }
+    if (node?.class_type === 'YUE_Stage_A_Sampler') {
+      meta.tags = String(inputs.genres_prompt || '')
+      meta.lyrics = String(inputs.lyrics_prompt || '')
+      meta.engine = 'yue'
+      meta.duration = Number(inputs.prompt_end_time || meta.duration)
+      meta.seed = Number(inputs.seed || meta.seed)
+    }
+    if (node?.class_type === 'EmptyAceStepLatentAudio' || node?.class_type === 'EmptyAceStep1.5LatentAudio') {
+      meta.duration = Number(inputs.seconds || 0)
+    }
+    if (node?.class_type === 'KSampler') {
+      meta.steps = Number(inputs.steps || 0)
+      meta.seed = Number(inputs.seed || 0)
+      meta.cfg = Number(inputs.cfg || 0)
+      meta.samplerName = String(inputs.sampler_name || '')
+      meta.scheduler = String(inputs.scheduler || '')
+    }
+    if (node?.class_type === 'CheckpointLoaderSimple') {
+      const ckpt = String(inputs.ckpt_name || '')
+      if (/1\.5|turbo_aio/i.test(ckpt)) meta.engine = 'ace-step-1.5'
+    }
+  }
+  return meta
+}
+
+function historyCompletedAt(entry: unknown) {
+  const messages = (entry as { status?: { messages?: Array<[string, { timestamp?: number }]> } })?.status?.messages
+  const last = messages?.find(item => item[0] === 'execution_success')
+  return Number(last?.[1]?.timestamp || 0)
+}
+
+export function findHistoryAudio(
+  history: Record<string, unknown> | null,
+  opts?: { promptId?: string; tags?: string; since?: number }
+) {
+  if (!history) return null
+  if (opts?.promptId) {
+    const audio = extractAudio(history, opts.promptId)
+    if (audio) return { promptId: opts.promptId, audio }
+  }
+  const want = String(opts?.tags || '').trim()
+  const matches: {
+    promptId: string
+    audio: NonNullable<ReturnType<typeof extractAudio>>
+    ts: number
+    tags: string
+  }[] = []
+  for (const [id, entry] of Object.entries(history)) {
+    const audio = extractAudio({ [id]: entry as Record<string, unknown> }, id)
+    if (!audio || !isOurComfyAudio(audio)) continue
+    const ts = historyCompletedAt(entry)
+    if (opts?.since && ts && ts < opts.since) continue
+    const meta = extractAceMetaFromHistory(entry)
+    matches.push({ promptId: id, audio, ts: historyCompletedAt(entry), tags: meta.tags })
+  }
+  if (!matches.length) return null
+  const tagged = want
+    ? matches.filter((item) => {
+      const tags = item.tags.trim()
+      return tags === want || tags.includes(want) || want.includes(tags)
+    })
+    : matches
+  const pool = want ? tagged : matches
+  if (!pool.length) return null
+  pool.sort((a, b) => b.ts - a.ts)
+  return { promptId: pool[0].promptId, audio: pool[0].audio }
+}
+
 export function extractVideo(history: Record<string, unknown> | null, promptId: string) {
   if (!history) return null
   const wrapped = history[promptId] as { outputs?: Record<string, unknown> } | undefined

server/utils/generationPresets.ts

@@ -40,6 +40,7 @@ import {
   clampMusicCfg,
   clampMusicDuration,
   clampMusicSteps,
+  parseMusicEngine,
   sanitizeMusicLyrics,
   sanitizeMusicTags
 } from '~/utils/music'
@@ -173,7 +174,8 @@ function sanitizeMusicSettings(raw: unknown): MusicPresetSettings {
     duration: clampMusicDuration(rec.duration),
     steps: clampMusicSteps(rec.steps),
     seed: typeof rec.seed === 'number' && Number.isFinite(rec.seed) ? Math.max(0, Math.floor(rec.seed)) : undefined,
-    cfg: rec.cfg == null ? undefined : clampMusicCfg(rec.cfg)
+    cfg: rec.cfg == null ? undefined : clampMusicCfg(rec.cfg),
+    engine: parseMusicEngine(rec.engine)
   }
 }
 

server/utils/library.ts

@@ -19,6 +19,7 @@ import {
   type KeepPromptLists
 } from '~/utils/keepPrompts'
 import { familyBaseName, nextClipPartName, clipPartLabel } from '~/utils/libraryBrowse'
+import { musicExtension } from '~/utils/music'
 
 export { familyBaseName, nextClipPartName, clipPartLabel }
 
@@ -1170,6 +1171,7 @@ export async function maintainLibrary(owner: string, folderId?: string) {
   const busy = Boolean(live && live.running > 0)
   if (!busy) {
     await importMissingComfyVideos(owner, folderId).catch(() => [])
+    await importMissingComfyAudio(owner, folderId).catch(() => [])
   }
   await ensureClipThumbnails(owner, { skipHistory: busy })
   if (!busy) await sweepOwnerExtensionStills(owner)
@@ -2043,3 +2045,48 @@ export async function importMissingComfyVideos(owner: string, folderId?: string)
   }
   return imported
 }
+
+export async function importMissingComfyAudio(owner: string, folderId?: string) {
+  const catalog = readCatalog(owner)
+  const folder = catalog.folders.find(item => item.id === folderId) || catalog.folders[0]
+  if (!folder) return []
+  const known = new Set((catalog.tracks || []).map(track => track.comfyFilename).filter(Boolean) as string[])
+  const history = await fetchHistoryAll()
+  const found: {
+    promptId: string
+    audio: { filename: string; subfolder: string; type: string }
+    meta: ReturnType<typeof extractAceMetaFromHistory>
+  }[] = []
+  for (const [promptId, entry] of Object.entries(history)) {
+    const audio = extractAudio({ [promptId]: entry as Record<string, unknown> }, promptId)
+    if (!audio || known.has(audio.filename) || !isOurComfyAudio(audio)) continue
+    found.push({ promptId, audio, meta: extractAceMetaFromHistory(entry) })
+  }
+  found.sort((a, b) => a.audio.filename.localeCompare(b.audio.filename, undefined, { numeric: true }))
+  const imported = []
+  for (const item of found) {
+    const buffer = await downloadComfyAudio(item.audio)
+    const ext = musicExtension(item.audio.filename)
+    const track = await saveTrack({
+      ownerKey: owner,
+      folderId: folder.id,
+      name: clipTitle({ name: '', prompt: item.meta.tags || 'Recovered track' }),
+      tags: item.meta.tags || 'Recovered from ComfyUI',
+      lyrics: item.meta.lyrics || '',
+      duration: item.meta.duration,
+      steps: item.meta.steps,
+      seed: item.meta.seed,
+      cfg: item.meta.cfg,
+      instrumental: !String(item.meta.lyrics || '').trim(),
+      audio: buffer,
+      ext,
+      engine: item.meta.engine,
+      samplerName: item.meta.samplerName,
+      scheduler: item.meta.scheduler,
+      lyricsStrength: item.meta.lyricsStrength,
+      comfyFilename: item.audio.filename
+    })
+    imported.push(track)
+  }
+  return imported
+}

server/utils/musicChain.ts

@@ -1,9 +1,9 @@
 import { createJob, emitJob, type Job } from '~/server/utils/jobs'
-import { extractAudio, fetchHistory, queuePrompt } from '~/server/utils/comfy'
+import { extractAudio, fetchHistory, fetchHistoryAll, findHistoryAudio, 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 { 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'
@@ -131,13 +131,35 @@ function watchMusicJob(job: Job): Promise<void> {
       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 || !job.promptId) return false
+      if (settled || finishing) return false
       finishing = true
       job.saving = true
       try {
-        const history = await fetchHistory(job.promptId)
-        const audio = extractAudio(history, job.promptId)
+        const audio = await loadJobAudio()
         if (!audio) {
           finishing = false
           job.saving = false
@@ -205,14 +227,20 @@ function watchMusicJob(job: Job): Promise<void> {
     }
 
     const pollHistory = async () => {
-      if (settled || finishing || !job.promptId) return
+      if (settled || finishing) return
       try {
-        const history = await fetchHistory(job.promptId)
-        if (extractAudio(history, job.promptId)) {
+        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 = (history?.[job.promptId] || history) as {
+        const entry = (job.promptId ? history?.[job.promptId] : undefined) as {
           status?: { status_str?: string; completed?: boolean }
         } | undefined
         const status = entry?.status?.status_str
@@ -242,6 +270,7 @@ function watchMusicJob(job: Job): Promise<void> {
       }
       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)
@@ -259,7 +288,7 @@ function watchMusicJob(job: Job): Promise<void> {
       }
       if (type === 'executing') {
         const node = data.node == null ? '' : String(data.node)
-        if (!node && job.promptId) void succeed()
+        if (!node) void succeed()
         else if (node) {
           emitLocal({
             type: 'status',
@@ -327,9 +356,11 @@ export async function startMusicJob(params: MusicJobParams) {
           queuePending: status.queuePending
         })
       })
-      job.status = 'uploading'
-      emitJob(job, { type: 'status', message: 'Queueing ACE-Step...', progress: 6 })
+      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)
       const graph = buildMusicWorkflow({
         tags: params.tags,
         lyrics: params.instrumental ? '' : params.lyrics,
@@ -344,10 +375,20 @@ export async function startMusicJob(params: MusicJobParams) {
         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 })
+      try {
+        const queued = await queuePrompt(graph, job.clientId, { timeoutMs: 25_000 })
+        if (queued?.prompt_id) job.promptId = queued.prompt_id
+        emitJob(job, { type: 'status', message: 'Job queued on ComfyUI', 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: 'Waiting for ComfyUI to pick up the job...',
+          progress: 8
+        })
+      }
       await done
     } catch (error) {
       const message = error instanceof Error ? error.message : String(error)

server/utils/musicWorkflow.ts

@@ -1,5 +1,8 @@
 import workflowAceStep from '../assets/workflow_ace_step.json'
+import workflowAceStep15 from '../assets/workflow_ace_step_15.json'
+import workflowYue from '../assets/workflow_yue.json'
 import type { MusicEngine } from '~/utils/music'
+import { parseMusicBpm } from '~/utils/music'
 
 export type MusicWorkflowParams = {
   tags: string
@@ -15,6 +18,7 @@ export type MusicWorkflowParams = {
   checkpoint?: string
   engine?: MusicEngine
   vocalMultiplier?: number
+  bpm?: number
 }
 
 type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
@@ -22,23 +26,38 @@ type WorkflowGraph = Record<string, WorkflowNode>
 
 const ACE_V1_CKPT = 'ace_step_v1_3.5b.safetensors'
 const ACE_15_CKPT = 'ace_step_1.5_turbo_aio.safetensors'
+const YUE_S1_DEFAULT = 'C:\\Users\\ianjm\\AppData\\Local\\Comfy-Desktop\\ComfyUI-Shared\\models\\YuE\\YuE-s1-7B-anneal-en-cot'
+const YUE_S2_DEFAULT = 'C:\\Users\\ianjm\\AppData\\Local\\Comfy-Desktop\\ComfyUI-Shared\\models\\YuE\\YuE-s2-1B-general'
+const YUE_NODE = 'YUE_Stage_A_Loader'
 
 export function musicFilenamePrefix() {
   return String(useRuntimeConfig().comfyAudioPrefix || process.env.COMFY_AUDIO_PREFIX || 'audio/AIGen')
 }
 
-export function buildMusicWorkflow(params: MusicWorkflowParams): WorkflowGraph {
+function yueStageARepo() {
+  return String(process.env.YUE_STAGE_A || useRuntimeConfig().yueStageA || YUE_S1_DEFAULT)
+}
+
+function yueStageBRepo() {
+  return String(process.env.YUE_STAGE_B || useRuntimeConfig().yueStageB || YUE_S2_DEFAULT)
+}
+
+function applyPrefix(graph: WorkflowGraph, prefix: string) {
+  for (const node of Object.values(graph)) {
+    if ((node.class_type === 'SaveAudio' || node.class_type === 'SaveAudioMP3') && 'filename_prefix' in node.inputs) {
+      node.inputs.filename_prefix = prefix
+    }
+  }
+}
+
+function buildAceV1Workflow(params: MusicWorkflowParams): WorkflowGraph {
   const graph = structuredClone(workflowAceStep as WorkflowGraph)
-  const engine = params.engine || 'ace-step'
-  const checkpoint = params.checkpoint
-    || (engine === 'ace-step-1.5' ? ACE_15_CKPT : ACE_V1_CKPT)
   const prefix = params.filenamePrefix || musicFilenamePrefix()
   const sampler = params.samplerName || 'euler'
   const scheduler = params.scheduler || 'simple'
-
   for (const node of Object.values(graph)) {
     if (node.class_type === 'CheckpointLoaderSimple' && 'ckpt_name' in node.inputs) {
-      node.inputs.ckpt_name = checkpoint
+      node.inputs.ckpt_name = params.checkpoint || ACE_V1_CKPT
     }
     if (node.class_type === 'TextEncodeAceStepAudio') {
       node.inputs.tags = params.tags
@@ -55,14 +74,99 @@ export function buildMusicWorkflow(params: MusicWorkflowParams): WorkflowGraph {
       node.inputs.sampler_name = sampler
       node.inputs.scheduler = scheduler
     }
-    if (node.class_type === 'SaveAudio' && 'filename_prefix' in node.inputs) {
-      node.inputs.filename_prefix = prefix
-    }
   }
-
+  applyPrefix(graph, prefix)
   return graph
 }
 
-export const MUSIC_NODE_LABELS: Record<string, string> = Object.fromEntries(
-  Object.entries(workflowAceStep as WorkflowGraph).map(([id, node]) => [id, node._meta?.title || `Node ${id}`])
-)
+function buildAce15Workflow(params: MusicWorkflowParams): WorkflowGraph {
+  const graph = structuredClone(workflowAceStep15 as WorkflowGraph)
+  const prefix = params.filenamePrefix || musicFilenamePrefix()
+  const sampler = params.samplerName || 'euler'
+  const scheduler = params.scheduler || 'simple'
+  const bpm = params.bpm || parseMusicBpm(params.tags)
+  for (const node of Object.values(graph)) {
+    if (node.class_type === 'CheckpointLoaderSimple' && 'ckpt_name' in node.inputs) {
+      node.inputs.ckpt_name = params.checkpoint || ACE_15_CKPT
+    }
+    if (node.class_type === 'TextEncodeAceStepAudio1.5') {
+      node.inputs.tags = params.tags
+      node.inputs.lyrics = params.lyrics
+      node.inputs.seed = params.seed
+      node.inputs.bpm = bpm
+      node.inputs.duration = params.duration
+    }
+    if (node.class_type === 'EmptyAceStep1.5LatentAudio' && 'seconds' in node.inputs) {
+      node.inputs.seconds = params.duration
+    }
+    if (node.class_type === 'KSampler') {
+      node.inputs.seed = params.seed
+      node.inputs.steps = params.steps
+      node.inputs.cfg = params.cfg
+      node.inputs.sampler_name = sampler
+      node.inputs.scheduler = scheduler
+    }
+  }
+  applyPrefix(graph, prefix)
+  return graph
+}
+
+function buildYueWorkflow(params: MusicWorkflowParams): WorkflowGraph {
+  const graph = structuredClone(workflowYue as WorkflowGraph)
+  const prefix = params.filenamePrefix || musicFilenamePrefix()
+  const tags = params.tags
+  const lyrics = params.lyrics || '[instrumental]\n'
+  const segments = Math.max(1, Math.min(10, lyrics.split(/\[[^\]]+\]/g).filter(part => part.trim()).length || 2))
+  for (const node of Object.values(graph)) {
+    if (node.class_type === 'YUE_Stage_A_Loader') {
+      node.inputs.stage_A_repo = yueStageARepo()
+    }
+    if (node.class_type === 'YUE_Stage_B_Loader') {
+      node.inputs.stage_B_repo = yueStageBRepo()
+    }
+    if (node.class_type === 'YUE_Stage_A_Sampler') {
+      node.inputs.genres_prompt = tags.replace(/,/g, ' ').replace(/\s+/g, ' ').trim()
+      node.inputs.lyrics_prompt = lyrics
+      node.inputs.seed = params.seed
+      node.inputs.prompt_end_time = params.duration
+      node.inputs.run_n_segment = segments
+    }
+  }
+  applyPrefix(graph, prefix)
+  return graph
+}
+
+export function buildMusicWorkflow(params: MusicWorkflowParams): WorkflowGraph {
+  const engine = params.engine || 'ace-step'
+  if (engine === 'yue') return buildYueWorkflow(params)
+  if (engine === 'ace-step-1.5') return buildAce15Workflow(params)
+  return buildAceV1Workflow(params)
+}
+
+export async function assertMusicEngineNodes(engine: MusicEngine | undefined) {
+  const { comfyHasClassType } = await import('~/server/utils/comfy')
+  if (engine === 'yue') {
+    const present = await comfyHasClassType(YUE_NODE)
+    if (present === false) {
+      throw createError({
+        statusCode: 400,
+        statusMessage: 'YuE is not installed on Beast. Clone ComfyUI_YuE into custom_nodes and restart Comfy, then try again.'
+      })
+    }
+  }
+  if (engine === 'ace-step-1.5') {
+    const present = await comfyHasClassType('TextEncodeAceStepAudio1.5')
+    if (present === false) {
+      throw createError({
+        statusCode: 400,
+        statusMessage: 'This ComfyUI build does not have ACE-Step 1.5 nodes. Update Comfy Desktop, or turn the 1.5 toggle off.'
+      })
+    }
+  }
+}
+
+export const MUSIC_NODE_LABELS: Record<string, string> = Object.fromEntries([
+  ...Object.entries(workflowAceStep as WorkflowGraph).map(([id, node]) => [id, node._meta?.title || `Node ${id}`]),
+  ...Object.entries(workflowAceStep15 as WorkflowGraph).map(([id, node]) => [id, node._meta?.title || `Node ${id}`]),
+  ...Object.entries(workflowYue as WorkflowGraph).map(([id, node]) => [id, node._meta?.title || `Node ${id}`])
+])