Add optional per-image-profile Qwen prompt enhancers
This commit is contained in:
@@ -0,0 +1,32 @@
|
||||
import tempfile,unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock,patch
|
||||
from profiles import Profiles,PROMPT_ENHANCER_REPOS
|
||||
from prompt_enhancer import PromptEnhancer
|
||||
|
||||
class PromptEnhancerTests(unittest.TestCase):
|
||||
def test_profile_settings_are_persisted_and_validated(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path=Path(directory)/'profiles.json';catalog=Mock();catalog.entry.return_value={'id':'model','kind':'image','repo':'example/image','file':'model.safetensors','profile_eligible':True}
|
||||
profiles=Profiles(path,catalog)
|
||||
profiles.rows=[dict(id='p',revision=1,name='Image',kind='image',model_id='model',parameters=dict(width=512,height=512,steps=4,seed=1,guidance=1),components={})]
|
||||
config=dict(t2i=PROMPT_ENHANCER_REPOS['t2i'],i2i=None,device='auto')
|
||||
with self.assertRaises(ValueError):profiles.configure_prompt_enhancer(dict(id='p',revision=1,prompt_enhancer=dict(config,t2i='other/model')))
|
||||
self.assertFalse(path.exists())
|
||||
saved=profiles.configure_prompt_enhancer(dict(id='p',revision=1,prompt_enhancer=config))
|
||||
self.assertEqual(saved['revision'],2)
|
||||
self.assertEqual(Profiles(path,catalog).rows[0]['prompt_enhancer'],config)
|
||||
with self.assertRaises(ValueError):profiles.configure_prompt_enhancer(dict(id='p',revision=1,prompt_enhancer=config))
|
||||
def test_rewrite_off_never_starts_worker(self):
|
||||
with tempfile.TemporaryDirectory() as directory,patch('prompt_enhancer.subprocess.Popen') as spawn:
|
||||
enhancer=PromptEnhancer(directory,Mock(),Mock(),Mock())
|
||||
self.assertEqual(enhancer.rewrite({'prompt_enhancer':{'t2i':None}},'a red bicycle')['prompt'],'a red bicycle')
|
||||
spawn.assert_not_called()
|
||||
def test_device_selection_does_not_use_busy_gpu(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
enhancer=PromptEnhancer(directory,Mock(),Mock(),Mock())
|
||||
rows=[dict(name='RTX 5080',uuid='GPU-5080',processes=0,free_mib=12000),dict(name='RTX 3060',uuid='GPU-3060',processes=1,free_mib=11000)]
|
||||
with patch('prompt_enhancer.probe',return_value=rows):
|
||||
self.assertEqual(enhancer._devices('5080'),'GPU-5080')
|
||||
self.assertEqual(enhancer._devices('auto'),'GPU-5080')
|
||||
self.assertEqual(enhancer._devices('cpu'),'')
|
||||
Reference in New Issue
Block a user