From 62671d05daf679771b417c9b6c81e7f90f28ee23 Mon Sep 17 00:00:00 2001
From: Towsty <ianjmcarr@gmail.com>
Date: Mon, 14 Sep 2026 19:28:43 -0500
Subject: [PATCH] Allow multi-section YuE2 lyrics and force non-flash attention
 on Windows.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 pages/music.vue                   |  9 +++---
 scripts/yue2-host.mjs             |  4 +--
 scripts/yue2-worker.py            | 54 +++++++++++++++++++++++++------
 server/api/generate/music.post.ts |  5 ++-
 tests/test_yue2_worker.py         | 30 +++++++++++++----
 tests/yue2.test.mjs               |  8 +++--
 6 files changed, 85 insertions(+), 25 deletions(-)

diff --git a/pages/music.vue b/pages/music.vue
index e761503..0106405 100644
--- a/pages/music.vue
+++ b/pages/music.vue
@@ -34,7 +34,7 @@
             </select>
             <span class="block text-xs text-zinc-400">One lyric section. Compile off. Profile 3 is never selected automatically.</span>
           </label>
-          <p v-else-if="engineFamily === 'yue2'" class="text-xs text-zinc-400">YuE2-3B · one lyric section · plan → synthesize → decode · no score editor</p>
+          <p v-else-if="engineFamily === 'yue2'" class="text-xs text-zinc-400">YuE2-3B · lyrics with section tags · plan → synthesize → decode · no score editor</p>
           <div v-if="!advanced" class="space-y-4">
             <label v-if="engineFamily === 'ace'" class="block text-sm">Steps<input v-model.number="steps" type="number" :min="stepsMin" :max="stepsMax" class="mt-2 w-full rounded-xl border border-white/15 bg-zinc-900 p-3"></label>
             <label class="block text-sm">Seed<input v-model="seed" placeholder="random" class="mt-2 w-full rounded-xl border border-white/15 bg-zinc-900 p-3"></label>
@@ -109,7 +109,7 @@
               :max="durationMax"
               step="5"
             >
-            <span class="mt-1 block text-[11px] text-zinc-500">{{ engineFamily === 'yue' || engineFamily === 'yue2' ? 'Approximate length · one lyric section' : `${durationMin}–${durationMax} seconds` }}</span>
+            <span class="mt-1 block text-[11px] text-zinc-500">{{ engineFamily === 'yue' ? 'Approximate length · one lyric section' : engineFamily === 'yue2' ? 'Approximate length · multiple sections OK' : `${durationMin}–${durationMax} seconds` }}</span>
           </label>
           <label v-if="advanced && engineFamily === 'ace'" class="block text-sm">
             <span class="mb-1 block font-medium text-zinc-300">Steps · {{ steps }}</span>
@@ -333,7 +333,8 @@ const blockReason = computed(() => {
   if ((engineFamily.value === 'yue' || engineFamily.value === 'yue2') && instrumental.value) return 'YuE requires lyrics. Use ACE for instrumental music.'
   if ((engineFamily.value === 'yue' || engineFamily.value === 'yue2') && duration.value > 150) return 'YuE supports up to 150 seconds per section.'
   if (!tags.value.trim()) return 'Add genre and style tags.'
-  if (engineFamily.value === 'yue' || engineFamily.value === 'yue2') { const problem = yueLyricsProblem(lyrics.value); if (problem) return problem }
+  if (engineFamily.value === 'yue') { const problem = yueLyricsProblem(lyrics.value); if (problem) return problem }
+  if (engineFamily.value === 'yue2' && !lyrics.value.trim()) return 'Write lyrics for YuE2.'
   if (!instrumental.value && !lyrics.value.trim()) return 'Write lyrics, or turn on Instrumental.'
   return ''
 })
@@ -341,7 +342,7 @@ const blockReason = computed(() => {
 function selectEngineFamily(family: 'ace' | 'yue' | 'yue2') {
   engineFamily.value = family
   if (family === 'yue' || family === 'yue2') { instrumental.value = false; duration.value = Math.min(150, duration.value) }
-  if ((family === 'yue' || family === 'yue2') && lyrics.value === DEFAULT_MUSIC_LYRICS) lyrics.value = '[Verse 1]\n'
+  if (family === 'yue' && lyrics.value === DEFAULT_MUSIC_LYRICS) lyrics.value = '[Verse 1]\n'
   if (family === 'ace' && ace15.value && steps.value === MUSIC_STEPS_DEFAULT) {
     steps.value = MUSIC_STEPS_DEFAULT_15
   }
diff --git a/scripts/yue2-host.mjs b/scripts/yue2-host.mjs
index d7debf6..622aaaf 100644
--- a/scripts/yue2-host.mjs
+++ b/scripts/yue2-host.mjs
@@ -10,9 +10,7 @@ export function validateYue2Request(body) {
   if (!/^[a-zA-Z0-9-]{12,80}$/.test(body.id || '')) throw new Error('Invalid job ID.')
   const tags = String(body.tags || '').trim()
   const lyrics = String(body.lyrics || '').trim()
-  if (!tags || tags.length > 2000 || !lyrics || lyrics.length > 8000) throw new Error('Genre tags and one lyric section are required.')
-  const sections = [...lyrics.matchAll(/\[([^\]]+)\]\s*([^\[]*)/gs)].filter(m => m[2].trim())
-  if (sections.length !== 1) throw new Error('YuE2 requires one non-empty lyric section.')
+  if (!tags || tags.length > 2000 || !lyrics || lyrics.length > 8000) throw new Error('Genre tags and non-empty lyrics are required.')
   return { id: body.id, duration, seed: body.seed, tags, lyrics }
 }
 
diff --git a/scripts/yue2-worker.py b/scripts/yue2-worker.py
index 43195cd..2c7be97 100644
--- a/scripts/yue2-worker.py
+++ b/scripts/yue2-worker.py
@@ -18,14 +18,47 @@ def emit(**event):
 
 
 def normalize_lyrics(text):
-    # Same one-section rule as YuEGP field validation.
+    # YuE2 accepts multiple [Verse]/[Chorus] headings. Only wrap when none exist.
     import re
-    sections = re.findall(r'\[([^\]]+)\]\s*([^\[]*)', text.strip(), re.S)
-    sections = [(name, words.strip()) for name, words in sections if words.strip()]
-    if len(sections) != 1:
-        raise ValueError('YuE2 requires one non-empty lyric section. Combine the lyrics under one heading.')
-    name, words = sections[0]
-    return '[' + re.sub(r'\W+', '', name).lower() + ']\n' + words + '\n\n'
+    text = str(text or '').strip()
+    if not text:
+        raise ValueError('YuE2 requires non-empty lyrics.')
+    if re.search(r'\[[^\]]+\]', text):
+        return text if text.endswith('\n') else text + '\n'
+    return '[song]\n' + text + '\n\n'
+
+
+def resolve_attention_backend(torch_mod):
+    """Pick a non-flash attention path for the Windows CUDA wheel on the 5080."""
+    flash = False
+    try:
+        check = getattr(torch_mod.backends.cuda, 'is_flash_attention_available', None)
+        flash = bool(check()) if callable(check) else False
+    except Exception:
+        flash = False
+    # Cognito/YuE2-Windows: when flash is unavailable, fall back to cudnn.
+    # This host never enables flash: the torch wheel exposes the op without USE_FLASH_ATTENTION.
+    if sys.platform == 'win32' or not flash:
+        if torch_mod.cuda.is_available() and torch_mod.backends.cudnn.is_available():
+            return 'cudnn', 'torch'
+        return 'sdpa', 'torch-eager'
+    if torch_mod.cuda.is_available() and torch_mod.backends.cudnn.is_available():
+        return 'cudnn', 'torch'
+    return 'sdpa', 'torch-eager'
+
+
+def patch_graph_attention(attention_backend):
+    from yue2.cuda_graph import GraphAR
+    original = GraphAR.__init__
+
+    def init(self, model, prefixes, max_tokens, *, capture=True, attention_backend='auto', fuse_projections=False):
+        if attention_backend in ('auto', 'flash'):
+            attention_backend = patch_graph_attention.forced
+        return original(self, model, prefixes, max_tokens, capture=capture,
+                        attention_backend=attention_backend, fuse_projections=fuse_projections)
+
+    patch_graph_attention.forced = attention_backend
+    GraphAR.__init__ = init
 
 
 def main():
@@ -56,16 +89,19 @@ def main():
     os.chdir(root)
     if str(root) not in sys.path:
         sys.path.insert(0, str(root))
-    emit(stage='loading', progress=1, message='Loading YuE2', model=model, vae=vae, duration=duration)
     import torch
     import soundfile as sf
+    attention_backend, pipeline_backend = resolve_attention_backend(torch)
+    emit(stage='loading', progress=1, message='Loading YuE2', model=model, vae=vae,
+         duration=duration, attention_backend=attention_backend, pipeline_backend=pipeline_backend)
     from yue2 import YuE2Pipeline
     if not torch.cuda.is_available():
         raise RuntimeError('YuE2 requires a CUDA GPU; CPU fallback is disabled.')
+    patch_graph_attention(attention_backend)
     pipe_kwargs = dict(style=style, lyrics=lyrics, cot='full', seed=seed)
     # Duration is kept for library metadata and validation. Upstream one-shot
     # requests do not take a seconds field; song length follows the plan.
-    with YuE2Pipeline.from_pretrained(model, vae=vae, device='cuda') as pipe:
+    with YuE2Pipeline.from_pretrained(model, vae=vae, device='cuda', backend=pipeline_backend) as pipe:
         emit(stage='plan', message='Planning melody and chords', progress=5)
         plan = pipe.plan(**pipe_kwargs)
         emit(stage='semantic', message='Generating semantic tokens', progress=25)
diff --git a/server/api/generate/music.post.ts b/server/api/generate/music.post.ts
index 06a7aa8..650ae02 100644
--- a/server/api/generate/music.post.ts
+++ b/server/api/generate/music.post.ts
@@ -71,10 +71,13 @@ export default defineEventHandler(async (event) => {
   if ((engine === 'yue' || engine === 'yue2') && instrumental) throw createError({ statusCode: 400, statusMessage: 'YuE currently requires lyrics. Use ACE for instrumental music.' })
   const yueProfile = body.yueProfile ?? 1
   if (engine === 'yue' && yueProfile !== 1 && yueProfile !== 3) throw createError({ statusCode: 400, statusMessage: 'Choose YuEGP profile 1 or manual fallback 3.' })
-  if ((engine === 'yue' || engine === 'yue2') && !instrumental) {
+  if (engine === 'yue' && !instrumental) {
     const problem = yueLyricsProblem(lyrics)
     if (problem) throw createError({ statusCode: 400, statusMessage: problem })
   }
+  if (engine === 'yue2' && !instrumental && !lyrics) {
+    throw createError({ statusCode: 400, statusMessage: 'YuE2 needs lyrics.' })
+  }
   const duration = clampMusicDuration(body.duration)
   if ((engine === 'yue' || engine === 'yue2') && duration > 150) throw createError({ statusCode: 400, statusMessage: 'YuE supports up to 150 seconds per section.' })
   const steps = clampMusicSteps(
diff --git a/tests/test_yue2_worker.py b/tests/test_yue2_worker.py
index 92e4165..539bba2 100644
--- a/tests/test_yue2_worker.py
+++ b/tests/test_yue2_worker.py
@@ -2,6 +2,7 @@
 import importlib.util
 from pathlib import Path
 import unittest
+from types import SimpleNamespace
 
 spec = importlib.util.spec_from_file_location('yue2_worker', Path(__file__).parents[1] / 'scripts/yue2-worker.py')
 worker = importlib.util.module_from_spec(spec)
@@ -9,12 +10,29 @@ spec.loader.exec_module(worker)
 
 
 class WorkerTests(unittest.TestCase):
-    def test_lyrics_preserve_words_and_normalize_ui_headings(self):
-        self.assertEqual(worker.normalize_lyrics('[Pre-Chorus]\nEvery word stays\n[Outro]\n'), '[prechorus]\nEvery word stays\n\n')
-        self.assertEqual(worker.normalize_lyrics('[Verse 1]\nHello'), '[verse1]\nHello\n\n')
-        for lyrics in ['', 'No heading', '[Verse]\nA\n[Chorus]\nB']:
-            with self.assertRaises(ValueError):
-                worker.normalize_lyrics(lyrics)
+    def test_lyrics_allow_multiple_sections_and_wrap_plain_text(self):
+        multi = '[Verse 1]\nHello\n\n[Chorus]\nSing it'
+        self.assertEqual(worker.normalize_lyrics(multi), multi + '\n')
+        self.assertEqual(worker.normalize_lyrics('[Verse 1]\nHello'), '[Verse 1]\nHello\n')
+        self.assertEqual(worker.normalize_lyrics('No heading yet'), '[song]\nNo heading yet\n\n')
+        with self.assertRaises(ValueError):
+            worker.normalize_lyrics('')
+        with self.assertRaises(ValueError):
+            worker.normalize_lyrics('   ')
+
+    def test_attention_backend_never_selects_flash(self):
+        torch_mod = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=lambda: True),
+            backends=SimpleNamespace(cudnn=SimpleNamespace(is_available=lambda: True)),
+        )
+        attention, pipeline = worker.resolve_attention_backend(torch_mod)
+        self.assertEqual(attention, 'cudnn')
+        self.assertEqual(pipeline, 'torch')
+        self.assertNotEqual(attention, 'flash')
+        torch_mod.backends.cudnn.is_available = lambda: False
+        attention, pipeline = worker.resolve_attention_backend(torch_mod)
+        self.assertEqual(attention, 'sdpa')
+        self.assertEqual(pipeline, 'torch-eager')
 
 
 if __name__ == '__main__':
diff --git a/tests/yue2.test.mjs b/tests/yue2.test.mjs
index 8de5bc5..1ccd587 100644
--- a/tests/yue2.test.mjs
+++ b/tests/yue2.test.mjs
@@ -11,9 +11,13 @@ import { createYueGpHost, validateYueGpRequest } from '../scripts/yuegp-host.mjs
 
 const request = { id: 'test-song-123456', tags: 'pop, warm vocals', lyrics: '[Verse 1]\nA quiet morning, a new day', seed: 0 }
 
-test('YuE2 defaults are 60 seconds and keep Yue lyric validation', () => {
+test('YuE2 defaults are 60 seconds and allow multiple lyric sections', () => {
   assert.deepEqual(validateYue2Request(request), { ...request, duration: 60 })
-  assert.throws(() => validateYue2Request({ ...request, lyrics: '[Verse]\nA\n[Chorus]\nB' }))
+  assert.deepEqual(
+    validateYue2Request({ ...request, lyrics: '[Verse]\nA\n[Chorus]\nB' }).lyrics,
+    '[Verse]\nA\n[Chorus]\nB'
+  )
+  assert.throws(() => validateYue2Request({ ...request, lyrics: '' }))
   assert.throws(() => validateYue2Request({ ...request, duration: 20 }))
   assert.throws(() => validateYue2Request({ ...request, id: '../escape' }))
 })