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

Allow multi-section YuE2 lyrics and force non-flash attention on Windows.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-09-14 19:28:43 -05:00
co-authored by Cursor
parent fae3b23e37
commit 62671d05da
6 changed files with 85 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_yue2_worker.py
+24 -6
View File
@@ -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__':