49 lines
1.9 KiB
Python
49 lines
1.9 KiB
Python
import unittest
|
|
|
|
from modules import util
|
|
|
|
|
|
class TestUtils(unittest.TestCase):
|
|
def test_can_parse_tokens_with_lora(self):
|
|
test_cases = [
|
|
{
|
|
"input": ("some prompt, very cool, <lora:hey-lora:0.4>, cool <lora:you-lora:0.2>", [], 5),
|
|
"output": [("hey-lora.safetensors", 0.4), ("you-lora.safetensors", 0.2)],
|
|
},
|
|
# Test can not exceed limit
|
|
{
|
|
"input": ("some prompt, very cool, <lora:hey-lora:0.4>, cool <lora:you-lora:0.2>", [], 1),
|
|
"output": [("hey-lora.safetensors", 0.4)],
|
|
},
|
|
# test Loras from UI take precedence over prompt
|
|
{
|
|
"input": (
|
|
"some prompt, very cool, <lora:l1:0.4>, <lora:l2:-0.2>, <lora:l3:0.3>, <lora:l4:0.5>, <lora:l6:0.24>, <lora:l7:0.1>",
|
|
[("hey-lora.safetensors", 0.4)],
|
|
5,
|
|
),
|
|
"output": [
|
|
("hey-lora.safetensors", 0.4),
|
|
("l1.safetensors", 0.4),
|
|
("l2.safetensors", -0.2),
|
|
("l3.safetensors", 0.3),
|
|
("l4.safetensors", 0.5),
|
|
],
|
|
},
|
|
# Test lora specification not separated by comma are ignored, only latest specified is used
|
|
{
|
|
"input": ("some prompt, very cool, <lora:hey-lora:0.4><lora:you-lora:0.2>", [], 3),
|
|
"output": [("you-lora.safetensors", 0.2)],
|
|
},
|
|
{
|
|
"input": ("<lora:foo:1..2>, <lora:bar:.>, <lora:baz:+> and <lora:quux:>", [], 6),
|
|
"output": []
|
|
}
|
|
]
|
|
|
|
for test in test_cases:
|
|
prompt, loras, loras_limit = test["input"]
|
|
expected = test["output"]
|
|
actual = util.parse_lora_references_from_prompt(prompt, loras, loras_limit)
|
|
self.assertEqual(expected, actual)
|