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>
2026-09-03 07:42:54 -05:00
parent ce85e7f513
commit 78b158471b
17 changed files with 951 additions and 110 deletions
@@ -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}`])
+])