towsty/aigen
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

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

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

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-09-03 07:42:54 -05:00
co-authored by Cursor
parent ce85e7f513
commit 78b158471b
17 changed files with 951 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/api/generate/music.post.ts
+14 -4
View File
@@ -5,6 +5,10 @@ import {
clampMusicCfg,
clampMusicDuration,
clampMusicSteps,
MUSIC_CFG_DEFAULT,
MUSIC_CFG_DEFAULT_15,
MUSIC_STEPS_DEFAULT,
MUSIC_STEPS_DEFAULT_15,
parseMusicEngine,
parseMusicSeed,
sanitizeMusicLyrics,
@@ -57,14 +61,20 @@ export default defineEventHandler(async (event) => {
const destFolder = library.folders.find(folder => folder.id === folderId)
const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
const engine = parseMusicEngine(body.engine)
const duration = clampMusicDuration(body.duration)
const steps = clampMusicSteps(body.steps)
const cfg = clampMusicCfg(body.cfg)
const steps = clampMusicSteps(
body.steps,
engine === 'ace-step-1.5' ? MUSIC_STEPS_DEFAULT_15 : MUSIC_STEPS_DEFAULT
)
const cfg = clampMusicCfg(
body.cfg,
engine === 'ace-step-1.5' ? MUSIC_CFG_DEFAULT_15 : MUSIC_CFG_DEFAULT
)
const lyricsStrength = clampLyricsStrength(body.lyricsStrength)
const seed = parseMusicSeed(body.seed)
const engine = parseMusicEngine(body.engine)
const clipName = allocateOutputName(ownerKey, folderId, String(body.name || '').trim(), {
engine: engine === 'ace-step-1.5' ? 'ace15' : 'ace'
engine: engine === 'ace-step-1.5' ? 'ace15' : engine === 'yue' ? 'yue' : 'ace'
})
const samplerName = String(body.sampler_name || 'euler')
server/api/generate/recover.post.ts
+73 -22
View File
@@ -16,11 +16,78 @@ function recentLibraryClip(event: Parameters<typeof assertLibraryOwner>[0]) {
}
}
function recentLibraryTrack(event: Parameters<typeof assertLibraryOwner>[0]) {
const { catalog } = assertLibraryOwner(event)
const newest = catalog.tracks?.[0]
if (!newest || Date.now() - newest.createdAt > RECENT_CLIP_MS) return null
return {
type: 'complete' as const,
status: 'complete' as const,
message: 'Track ready',
progress: 100,
trackId: newest.id,
filename: newest.comfyFilename,
folderLocked: false
}
}
async function recoverFinishedMedia(event: Parameters<typeof assertLibraryOwner>[0], opts?: { tags?: string }) {
const { owner } = assertLibraryOwner(event)
const library = publicLibrary(event)
const folderId = library.folders.find(folder => folder.unlocked)?.id
const tracks = folderId ? await importMissingComfyAudio(owner, folderId) : []
const want = String(opts?.tags || '').trim()
const { catalog } = assertLibraryOwner(event)
const match = [...tracks].reverse().find((track) => {
if (!want) return true
const tags = String(track.tags || '').trim()
return tags === want || tags.includes(want) || want.includes(tags)
}) || (want ? catalog.tracks.find((track) => {
const tags = String(track.tags || '').trim()
return Date.now() - track.createdAt < RECENT_CLIP_MS && (tags === want || tags.includes(want) || want.includes(tags))
}) : null)
if (match) {
return {
type: 'complete' as const,
status: 'complete' as const,
message: 'Recovered track from ComfyUI',
progress: 100,
trackId: match.id,
filename: match.comfyFilename,
folderLocked: false
}
}
if (opts?.tags) return null
const videos = folderId ? await importMissingComfyVideos(owner, folderId) : []
const newestVideo = videos[videos.length - 1]
if (newestVideo) {
return {
type: 'complete' as const,
status: 'complete' as const,
message: 'Recovered video from ComfyUI',
progress: 100,
clipId: newestVideo.id,
filename: newestVideo.comfyFilename,
hideThumbnail: newestVideo.hideThumbnail
}
}
return recentLibraryTrack(event) || recentLibraryClip(event)
}
export default defineEventHandler(async (event) => {
const body = await readBody<{ jobId?: string }>(event).catch(() => ({} as { jobId?: string }))
const jobId = String(body?.jobId || '')
const live = jobId ? getJob(jobId) : undefined
if (live) return jobSnapshot(live)
if (live) {
if (live.status === 'complete' || live.status === 'error' || live.trackId || live.clipId) {
return jobSnapshot(live)
}
if (live.kind === 'music' && Date.now() - live.startedAt > 15_000) {
const recovered = await recoverFinishedMedia(event, { tags: live.library?.tags })
if (recovered?.trackId) return recovered
}
return jobSnapshot(live)
}
if (jobId) {
const pending = readPendingJob(jobId)
@@ -29,28 +96,12 @@ export default defineEventHandler(async (event) => {
if (done) return done
return jobSnapshot(ensurePendingWatch(pending))
}
const saved = recentLibraryClip(event)
if (saved) return saved
const recovered = await recoverFinishedMedia(event)
if (recovered) return recovered
throw createError({ statusCode: 404, statusMessage: 'Job not found' })
}
const { owner } = assertLibraryOwner(event)
const library = publicLibrary(event)
const folderId = library.folders.find(folder => folder.unlocked)?.id
const imported = folderId ? await importMissingComfyVideos(owner, folderId) : []
const newest = imported[imported.length - 1]
if (newest) {
return {
type: 'complete',
status: 'complete',
message: 'Recovered video from ComfyUI',
progress: 100,
clipId: newest.id,
filename: newest.comfyFilename,
hideThumbnail: newest.hideThumbnail
}
}
const saved = recentLibraryClip(event)
if (saved) return saved
throw createError({ statusCode: 404, statusMessage: 'No finished ComfyUI video to recover' })
const recovered = await recoverFinishedMedia(event)
if (recovered) return recovered
throw createError({ statusCode: 404, statusMessage: 'No finished ComfyUI output to recover' })
})
server/assets/workflow_ace_step_15.json
+101
View File
@@ -0,0 +1,101 @@
{
"14": {
"inputs": {
"ckpt_name": "ace_step_1.5_turbo_aio.safetensors"
},
"class_type": "CheckpointLoaderSimple",
"_meta": {
"title": "Load ACE-Step 1.5"
}
},
"40": {
"inputs": {
"tags": "modern synthpop, 120 bpm, female vocals",
"lyrics": "[Verse 1]\n\n[Chorus]\n",
"seed": 0,
"bpm": 120,
"duration": 60,
"timesignature": "4",
"language": "en",
"keyscale": "C major",
"generate_audio_codes": true,
"cfg_scale": 2,
"temperature": 0.85,
"top_p": 0.9,
"top_k": 0,
"min_p": 0,
"clip": ["14", 1]
},
"class_type": "TextEncodeAceStepAudio1.5",
"_meta": {
"title": "Tags & Lyrics 1.5"
}
},
"47": {
"inputs": {
"conditioning": ["40", 0]
},
"class_type": "ConditioningZeroOut",
"_meta": {
"title": "Negative (zeroed)"
}
},
"17": {
"inputs": {
"seconds": 60,
"batch_size": 1
},
"class_type": "EmptyAceStep1.5LatentAudio",
"_meta": {
"title": "Audio Length 1.5"
}
},
"45": {
"inputs": {
"shift": 3,
"model": ["14", 0]
},
"class_type": "ModelSamplingAuraFlow",
"_meta": {
"title": "Model Sampling"
}
},
"3": {
"inputs": {
"seed": 0,
"steps": 8,
"cfg": 1,
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1,
"model": ["45", 0],
"positive": ["40", 0],
"negative": ["47", 0],
"latent_image": ["17", 0]
},
"class_type": "KSampler",
"_meta": {
"title": "Audio Sampler"
}
},
"16": {
"inputs": {
"samples": ["3", 0],
"vae": ["14", 2]
},
"class_type": "VAEDecodeAudio",
"_meta": {
"title": "VAE Decode Audio"
}
},
"19": {
"inputs": {
"filename_prefix": "audio/AIGen",
"audio": ["16", 0]
},
"class_type": "SaveAudio",
"_meta": {
"title": "Save Audio"
}
}
}
server/assets/workflow_yue.json
+75
View File
@@ -0,0 +1,75 @@
{
"15": {
"inputs": {
"stage_A_repo": "m-a-p/YuE-s1-7B-anneal-en-cot",
"xcodec_ckpt": "ckpt_00360000.pth",
"quantization_model": "fp16",
"use_mmgp": true,
"stage1_cache_size": 16384,
"exllamav2_cache_mode": "FP16",
"mmgp_profile": 2
},
"class_type": "YUE_Stage_A_Loader",
"_meta": {
"title": "YuE Stage A Loader"
}
},
"8": {
"inputs": {
"genres_prompt": "inspiring female uplifting pop airy vocal electronic bright",
"lyrics_prompt": "[verse]\n\n[chorus]\n",
"seed": 0,
"run_n_segment": 2,
"repetition_penalty": 1.1,
"prompt_start_time": 0,
"prompt_end_time": 30,
"max_new_tokens": 3000,
"use_dual_tracks_prompt": false,
"use_audio_prompt": false,
"offload_model": true,
"stage1_no_guidance": true,
"model": ["15", 0]
},
"class_type": "YUE_Stage_A_Sampler",
"_meta": {
"title": "YuE Stage A"
}
},
"16": {
"inputs": {
"stage_B_repo": "m-a-p/YuE-s2-1B-general",
"stage2_cache_size": 8192,
"stage2_batch_size": 2,
"exllamav2_cache_mode": "FP16",
"use_mmgp": true,
"info": ["8", 1]
},
"class_type": "YUE_Stage_B_Loader",
"_meta": {
"title": "YuE Stage B Loader"
}
},
"9": {
"inputs": {
"vocal_decoder_ckpt": "decoder_131000.pth",
"inst_decoder_ckpt": "decoder_151000.pth",
"rescale": true,
"stage1_set": ["8", 0],
"model": ["16", 0]
},
"class_type": "YUE_Stage_B_Sampler",
"_meta": {
"title": "YuE Stage B"
}
},
"19": {
"inputs": {
"filename_prefix": "audio/AIGen",
"audio": ["9", 0]
},
"class_type": "SaveAudio",
"_meta": {
"title": "Save Audio"
}
}
}
server/utils/comfy.ts
+108 -2
View File
@@ -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
+3 -1
View File
@@ -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
+47
View File
@@ -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
+57 -16
View File
@@ -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
+117 -13
View File
@@ -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}`])
])