Tighten YuE lyrics validation and Comfy music workflow wiring.

Require section tags before queueing YuE and keep the ACE/YuE path changes that were already in progress.

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

pages/music.vue

@@ -390,6 +390,35 @@ async function tryRecover() {
   }
 }
 
+async function attachLiveMusic(opts: {
+  studioId?: string
+  liveId: string
+  message?: string
+  progress?: number
+}) {
+  const liveId = String(opts.liveId || '').trim()
+  if (!liveId) return
+  if (opts.studioId) studioJobId.value = opts.studioId
+  liveMusicId.value = opts.studioId || liveMusicId.value || liveId
+  queued.value = false
+  busy.value = true
+  settled = false
+  if (typeof opts.progress === 'number') progress.value = opts.progress
+  if (opts.message) status.value = opts.message
+  else if (!status.value || /waiting for a generate|waiting in the job queue|waiting for gpu/i.test(status.value)) {
+    status.value = selectedEngine.value === 'yue'
+      ? 'YuE running on Comfy — Stage A can take 10–20+ minutes'
+      : 'Generating…'
+  }
+  if (jobId.value !== liveId) {
+    jobId.value = liveId
+    listen(liveId)
+  }
+  if (!recoverPoll) {
+    recoverPoll = setInterval(() => { void tryRecover() }, 8000)
+  }
+}
+
 async function refreshStudioQueue() {
   const data = await $fetch<{
     jobs?: Array<{
@@ -405,6 +434,23 @@ async function refreshStudioQueue() {
   const musicRow = rows.find(job => job.kind === 'music' && (job.status === 'waiting' || job.status === 'running' || job.status === 'held'))
   liveMusicId.value = musicRow?.id || ''
   queueCount.value = data.waitingCount || rows.filter(job => job.status === 'waiting').length
+
+  if (musicRow?.status === 'running' && musicRow.liveJobId) {
+    await attachLiveMusic({
+      studioId: musicRow.id,
+      liveId: musicRow.liveJobId
+    })
+  } else if (musicRow && (musicRow.status === 'waiting' || (musicRow.status === 'held' && !musicRow.liveJobId))) {
+    studioJobId.value = musicRow.id
+    queued.value = true
+    busy.value = true
+    status.value = musicRow.status === 'held'
+      ? 'Paused in the job queue'
+      : (selectedEngine.value === 'yue'
+        ? 'Waiting for GPU… YuE is heavy; do not spam Generate'
+        : 'Waiting in the job queue…')
+  }
+
   if (!studioJobId.value) return
   const row = rows.find(job => job.id === studioJobId.value)
   if (!row) {
@@ -429,16 +475,7 @@ async function refreshStudioQueue() {
     return
   }
   if (row.status === 'running' && row.liveJobId) {
-    queued.value = false
-    if (!jobId.value || jobId.value !== row.liveJobId) {
-      jobId.value = row.liveJobId
-      listen(row.liveJobId)
-    }
-    if (/waiting/i.test(status.value)) {
-      status.value = selectedEngine.value === 'yue'
-        ? 'YuE running on Comfy — Stage A can take 10–20+ minutes'
-        : 'Generating…'
-    }
+    await attachLiveMusic({ studioId: row.id, liveId: row.liveJobId })
   }
   if (row.status === 'error') {
     queued.value = false
@@ -449,13 +486,24 @@ async function refreshStudioQueue() {
     return
   }
   if (row.liveJobId && row.liveJobId !== jobId.value) {
-    queued.value = false
-    jobId.value = row.liveJobId
-    status.value = status.value || 'Generating…'
-    listen(row.liveJobId)
+    await attachLiveMusic({ studioId: row.id, liveId: row.liveJobId })
   }
 }
 
+async function resumeActiveMusic() {
+  const active = await $fetch<{
+    music?: { jobId?: string; message?: string; progress?: number; status?: string }
+  } | null>('/api/generate/active').catch(() => null)
+  const liveId = active?.music?.jobId
+  if (!liveId) return
+  if (active?.music?.status === 'error' || active?.music?.status === 'complete') return
+  await attachLiveMusic({
+    liveId,
+    message: active?.music?.message,
+    progress: typeof active?.music?.progress === 'number' ? active.music.progress : undefined
+  })
+}
+
 function applyEvent(payload: Record<string, any>) {
   if (payload.message) status.value = payload.message
   if (typeof payload.progress === 'number') progress.value = payload.progress
@@ -660,6 +708,7 @@ onMounted(async () => {
   await loadLibrary()
   await loadMusicPresets()
   applyIncomingMusic()
+  await resumeActiveMusic()
   await refreshStudioQueue()
   stopQueuePoll()
   queuePoll = setInterval(() => { void refreshStudioQueue() }, 8000)
@@ -670,7 +719,7 @@ onMounted(async () => {
     status.value = 'Loaded from library'
     return
   }
-  if (!audioUrl.value) {
+  if (!audioUrl.value && !busy.value) {
     try {
       const recovered = await $fetch<Record<string, any>>('/api/generate/recover', { method: 'POST', body: {} })
       if (recovered?.trackId && /recovered/i.test(String(recovered.message || ''))) {

server/api/generate/active.get.ts

@@ -72,7 +72,7 @@ export default defineEventHandler(async (event) => {
   }
 
   return {
-    jobId: video?.jobId || edit?.jobId || music?.jobId || null,
+    jobId: video?.jobId || edit?.jobId || null,
     video,
     edit,
     music

server/api/generate/music.post.ts

@@ -12,7 +12,8 @@ import {
   parseMusicEngine,
   parseMusicSeed,
   sanitizeMusicLyrics,
-  sanitizeMusicTags
+  sanitizeMusicTags,
+  yueLyricsProblem
 } from '~/utils/music'
 
 export default defineEventHandler(async (event) => {
@@ -62,6 +63,10 @@ export default defineEventHandler(async (event) => {
   const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
 
   const engine = parseMusicEngine(body.engine)
+  if (engine === 'yue' && !instrumental) {
+    const problem = yueLyricsProblem(lyrics)
+    if (problem) throw createError({ statusCode: 400, statusMessage: problem })
+  }
   const duration = clampMusicDuration(body.duration)
   const steps = clampMusicSteps(
     body.steps,

server/assets/workflow_yue.json

@@ -19,10 +19,10 @@
       "genres_prompt": "inspiring female uplifting pop airy vocal electronic bright",
       "lyrics_prompt": "[verse]\n\n[chorus]\n",
       "seed": 0,
-      "run_n_segment": 2,
+      "run_n_segment": 1,
       "repetition_penalty": 1.1,
       "prompt_start_time": 0,
-      "prompt_end_time": 45,
+      "prompt_end_time": 30,
       "max_new_tokens": 2944,
       "use_dual_tracks_prompt": false,
       "use_audio_prompt": false,

server/utils/comfy.ts

@@ -195,6 +195,16 @@ export async function interruptComfy() {
   return Boolean(res?.ok)
 }
 
+/** Ask Comfy to drop cached models so a heavy YuE job starts with a clean 16GB card. */
+export async function freeComfyVram() {
+  await comfyFetch('/free', {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({ unload_models: true, free_memory: true }),
+    signal: AbortSignal.timeout(15000)
+  }).catch(() => null)
+}
+
 export async function fetchHistory(promptId: string) {
   const res = await comfyFetch(`/history/${encodeURIComponent(promptId)}`)
   if (!res.ok) return null

server/utils/musicChain.ts

@@ -1,5 +1,5 @@
 import { createJob, emitJob, type Job } from '~/server/utils/jobs'
-import { extractAudio, fetchHistory, fetchHistoryAll, findHistoryAudio, purgeComfyArtifacts, queuePrompt } from '~/server/utils/comfy'
+import { extractAudio, fetchHistory, fetchHistoryAll, findHistoryAudio, freeComfyVram, purgeComfyArtifacts, queuePrompt } from '~/server/utils/comfy'
 import { comfyWsUrl } from '~/server/utils/comfy'
 import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
 import { downloadComfyAudio, saveTrack } from '~/server/utils/library'
@@ -33,7 +33,7 @@ function sleep(ms: number) {
 function classifyError(message: string) {
   const lower = message.toLowerCase()
   if (lower.includes('out of memory') || lower.includes('exceed allowed memory') || (lower.includes('cuda') && lower.includes('alloc')) || lower.includes('vram')) {
-    return 'ComfyUI ran out of VRAM. For YuE use shorter lyrics (fewer sections) or ACE-Step instead.'
+    return 'ComfyUI ran out of VRAM on YuE. Sleep→Poke Comfy, don’t run video/image first, use one short [Chorus] only — or use ACE-Step.'
   }
   if (lower.includes('econnrefused') || lower.includes('unreachable') || lower.includes('fetch failed')) {
     return 'ComfyUI host connection dropped. Confirm the desktop instance is running.'
@@ -372,6 +372,10 @@ export async function startMusicJob(params: MusicJobParams) {
       emitJob(job, { type: 'status', message: `Queueing ${engineLabel}...`, progress: 6 })
       await waitForSocket(job, 4000)
       await assertMusicEngineNodes(params.engine)
+      if (params.engine === 'yue') {
+        emitJob(job, { type: 'status', message: 'Freeing Comfy VRAM for YuE…', progress: 5 })
+        await freeComfyVram()
+      }
       const graph = buildMusicWorkflow({
         tags: params.tags,
         lyrics: params.instrumental ? '' : params.lyrics,

server/utils/musicWorkflow.ts

@@ -26,8 +26,8 @@ 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_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() {
@@ -115,14 +115,14 @@ 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 lyrics = params.lyrics || '[Instrumental]\n'
   // Beast 5080 16GB recipe (keep in sync with workflow_yue.json):
   // - mmgp profile 2 only (3–5 enable quanto and crash under Comfy InferenceMode)
-  // - max 2 Stage-A segments (3+ OOMs)
+  // - exactly 1 Stage-A lyric section (2+ still OOMs beside xcodec)
   // - max_new_tokens floor is 2944 (node min); do not go lower
   // - Triton/torch.compile off on Windows
-  const sections = lyrics.split(/\[[^\]]+\]/g).filter(part => part.trim()).length || 2
-  const segments = Math.max(1, Math.min(2, sections))
+  // - xcodec stays on CPU during Stage A (node patch)
+  const segments = 1
   for (const node of Object.values(graph)) {
     if (node.class_type === 'YUE_Stage_A_Loader') {
       node.inputs.stage_A_repo = yueStageARepo()
@@ -138,10 +138,11 @@ function buildYueWorkflow(params: MusicWorkflowParams): WorkflowGraph {
       node.inputs.stage2_batch_size = 1
     }
     if (node.class_type === 'YUE_Stage_A_Sampler') {
-      node.inputs.genres_prompt = tags.replace(/,/g, ' ').replace(/\s+/g, ' ').trim()
+      // Genres go to the model prompt only. Stage-A npy filenames are short (Windows MAX_PATH).
+      node.inputs.genres_prompt = tags.replace(/,/g, ' ').replace(/\s+/g, ' ').trim().slice(0, 400)
       node.inputs.lyrics_prompt = lyrics
       node.inputs.seed = params.seed
-      node.inputs.prompt_end_time = Math.min(params.duration, 45)
+      node.inputs.prompt_end_time = Math.min(params.duration, 30)
       node.inputs.run_n_segment = segments
       node.inputs.max_new_tokens = 2944
       node.inputs.use_dual_tracks_prompt = false

utils/music.ts

@@ -94,6 +94,22 @@ export function sanitizeMusicLyrics(raw: unknown, instrumental = false) {
   return String(raw || '').replace(/\r\n/g, '\n').trim().slice(0, MUSIC_LYRICS_MAX)
 }
 
+/** YuE splits on [Section] headers. Returns how many sections are present. */
+export function countYueLyricSections(lyrics: unknown) {
+  const text = String(lyrics || '')
+  const matches = text.match(/\[[^\]]+\]/g)
+  return matches ? matches.length : 0
+}
+
+export function yueLyricsProblem(lyrics: unknown): string | null {
+  const text = String(lyrics || '').trim()
+  if (!text) return 'YuE needs lyrics with section tags like [Verse 1] and [Chorus].'
+  if (countYueLyricSections(text) < 1) {
+    return 'YuE lyrics must include at least one [Section] tag (use the chips under Lyrics).'
+  }
+  return null
+}
+
 export function parseMusicSeed(raw: unknown) {
   const text = String(raw ?? '').trim().toLowerCase()
   if (!text || text === 'random') return Math.floor(Math.random() * 2_147_483_647)