Add ACE-Step music generation to the studio and library.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-02 23:58:18 -05:00
parent c79cca39ae
commit 5555155884
25 changed files with 2157 additions and 44 deletions
@@ -0,0 +1,68 @@
+import workflowAceStep from '../assets/workflow_ace_step.json'
+import type { MusicEngine } from '~/utils/music'
+
+export type MusicWorkflowParams = {
+  tags: string
+  lyrics: string
+  duration: number
+  steps: number
+  seed: number
+  cfg: number
+  lyricsStrength: number
+  samplerName?: string
+  scheduler?: string
+  filenamePrefix?: string
+  checkpoint?: string
+  engine?: MusicEngine
+  vocalMultiplier?: number
+}
+
+type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
+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'
+
+export function musicFilenamePrefix() {
+  return String(useRuntimeConfig().comfyAudioPrefix || process.env.COMFY_AUDIO_PREFIX || 'audio/AIGen')
+}
+
+export function buildMusicWorkflow(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
+    }
+    if (node.class_type === 'TextEncodeAceStepAudio') {
+      node.inputs.tags = params.tags
+      node.inputs.lyrics = params.lyrics
+      node.inputs.lyrics_strength = params.lyricsStrength
+    }
+    if (node.class_type === 'EmptyAceStepLatentAudio' && '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
+    }
+    if (node.class_type === 'SaveAudio' && 'filename_prefix' in node.inputs) {
+      node.inputs.filename_prefix = 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}`])
+)