44 lines
1.9 KiB
Python
44 lines
1.9 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_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()
|