Reuse component downloads and assignments for YuE2 music profiles

2026-09-30 21:21:04 +02:00
parent d700c68394
commit b731811482
5 changed files with 53 additions and 6 deletions
@@ -0,0 +1,26 @@
+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')))