Files
aigen/server/utils/musicWorkflow.ts
T

69 lines
2.3 KiB
TypeScript

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}`])
)