54 lines
2.5 KiB
Python
54 lines
2.5 KiB
Python
"""CPU-only adapter regression tests; never import torch or run inference."""
|
|
import importlib.util
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import Mock
|
|
|
|
spec = importlib.util.spec_from_file_location('worker', Path(__file__).parents[1] / 'scripts/yuegp-worker.py')
|
|
worker = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(worker)
|
|
|
|
|
|
class WorkerTests(unittest.TestCase):
|
|
def test_windows_memory_fix_preserves_profile_and_precision(self):
|
|
options = worker.memory_profile_options(1, False, 'win32', 32 * 1024 ** 3)
|
|
self.assertFalse(options['pinnedMemory'])
|
|
self.assertEqual(options['profile_no'], 1)
|
|
self.assertFalse(options['quantizeTransformer'])
|
|
self.assertFalse(options['compile'])
|
|
self.assertNotIn('budgets', options)
|
|
self.assertNotIn('pinnedMemory', worker.memory_profile_options(1, False, 'linux', 32 * 1024 ** 3))
|
|
self.assertTrue(worker.memory_profile_options(3, False, 'win32', 32 * 1024 ** 3)['quantizeTransformer'])
|
|
|
|
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_codec_construction_and_weight_loading_are_explicitly_cpu(self):
|
|
active = []
|
|
class Device:
|
|
def __enter__(self): active.append('cpu')
|
|
def __exit__(self, *args): active.pop()
|
|
codec = Mock()
|
|
def construct(**kwargs):
|
|
self.assertEqual(active, ['cpu'])
|
|
return codec
|
|
def load(path, **kwargs):
|
|
self.assertEqual(active, ['cpu'])
|
|
self.assertEqual(kwargs['map_location'], 'cpu')
|
|
return {'codec_model': 'weights'}
|
|
torch = SimpleNamespace(device=lambda name: Device() if name == 'cpu' else self.fail('GPU context'), load=load)
|
|
ns = {'SoundStream': construct, 'OmegaConf': SimpleNamespace(load=lambda _: SimpleNamespace(generator=SimpleNamespace(config={}))) }
|
|
self.assertIs(worker.load_codec_on_cpu(ns, torch), codec)
|
|
codec.load_state_dict.assert_called_once_with('weights')
|
|
codec.eval.assert_called_once()
|
|
self.assertEqual(active, [])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|