48 lines
3.6 KiB
Python
48 lines
3.6 KiB
Python
import tempfile,unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock
|
|
from profiles import Profiles,YUE2_REPO,YUE2_COMPONENTS,music_recipe
|
|
from catalog import file_role
|
|
class MusicComponentsTests(unittest.TestCase):
|
|
def test_download_recipe_assignment_and_remaining_worker_block(self):
|
|
model=dict(id='main',kind='music',repo=YUE2_REPO,file='yue2-3b-bf16.gguf',revision='a'*40)
|
|
entries=[dict(id=role,kind='music',repo=YUE2_REPO,file=info['files'][0],revision='a'*40,size=100) for role,info in YUE2_COMPONENTS.items()]
|
|
lookup={e['id']:e for e in [model]+entries}
|
|
c=Mock();c.entry.side_effect=lambda id:lookup[id] if id in lookup else (_ for _ in ()).throw(ValueError('missing'))
|
|
c.status.return_value=dict(entries=entries)
|
|
c.files.return_value=dict(repo=YUE2_REPO,revision='a'*40,files=[dict(name=e['file'],size=100) for e in entries],gated=False,license='apache-2.0')
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p=Profiles(Path(d)/'profiles.json',c);p.music_runtime_ready=lambda:True
|
|
p.rows=[dict(id='profile',kind='music',revision=1,name='YuE2',model_id='main',parameters={})]
|
|
r=p.components('main');self.assertEqual(len(r['requirements']),5);self.assertTrue(r['runtime_installed']);c.files.assert_called_once_with(YUE2_REPO,'a'*40)
|
|
self.assertTrue(all(len(x['available'])==1 for x in r['requirements']))
|
|
p.assign(dict(id='profile',revision=1,components={e['id']:e['id'] for e in entries}))
|
|
result=p.status()['profiles'][0];self.assertFalse(result['runnable']);self.assertEqual(len(result['blockers']),1)
|
|
entries[0]['revision']='b'*40
|
|
with self.assertRaises(ValueError):p._component(model,'vae','vae')
|
|
def test_components_cannot_be_main_models(self):
|
|
for info in YUE2_COMPONENTS.values():
|
|
for filename in info['files']:self.assertFalse(file_role(filename)['profile_eligible'])
|
|
self.assertIsNone(music_recipe(dict(kind='music',repo='other/Yue2',file='yue2-3b-bf16.gguf')))
|
|
def test_manual_setup_available_in_every_category_and_preserves_recipe(self):
|
|
for kind in ('chat','image','audio','stt','music','voice','video'):
|
|
model=dict(id='main',kind=kind,repo='unknown/model',file='weights.gguf',revision='a'*40)
|
|
component=dict(id='aux',kind=kind,repo='external/repo',file='vae.safetensors',revision='b'*40,size=1)
|
|
c=Mock();c.entry.side_effect=lambda id: {'main':model,'aux':component}[id]
|
|
c.status.return_value=dict(entries=[component]);c.files.return_value=dict(repo='unknown/model',revision='a'*40,files=[],gated=False,url='https://huggingface.co/unknown/model',license=None)
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p=Profiles(Path(d)/'profiles.json',c);p.rows=[dict(id='p',kind=kind,revision=1,name='Test',model_id='main',parameters={})]
|
|
self.assertTrue(p.components('main')['manual'])
|
|
p.assign(dict(id='p',revision=1,components={'manual_vae':'aux'}))
|
|
self.assertEqual(p.rows[0]['components']['manual_vae'],'aux')
|
|
self.assertTrue(p.references('aux'))
|
|
with self.assertRaises(ValueError):p.assign(dict(id='p',revision=2,components={'arbitrary':'aux'}))
|
|
def test_audio_runtime_package_definition_matches_exact_repo_and_model(self):
|
|
import json
|
|
from profiles import audio_package_recipe
|
|
with tempfile.TemporaryDirectory() as d:
|
|
Path(d,'model_specs').mkdir();Path(d,'model_specs/test.json').write_text(json.dumps(dict(packages=[dict(target_directory='Example-GGUF',files=['model.gguf','config.json'])])))
|
|
model=dict(repo='audio-cpp/Example-GGUF',file='model.gguf')
|
|
self.assertEqual(audio_package_recipe(model,d)['runtime_file_0']['files'],['config.json'])
|
|
self.assertIsNone(audio_package_recipe(dict(model,repo='other/Example-GGUF'),d))
|