2023-01-29 01:16:47 +00:00
|
|
|
import pytest
|
|
|
|
|
|
|
|
|
|
from imaginairy.prompt_schedules import parse_schedule_str
|
|
|
|
|
from imaginairy.utils import frange
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize(
|
2023-09-29 08:13:50 +00:00
|
|
|
("schedule_str", "expected"),
|
2023-01-29 01:16:47 +00:00
|
|
|
[
|
|
|
|
|
("prompt_strength[2:40:1]", ("prompt_strength", list(range(2, 40)))),
|
|
|
|
|
("prompt_strength[2:40:0.5]", ("prompt_strength", list(frange(2, 40, 0.5)))),
|
|
|
|
|
("prompt_strength[2,5,10,15]", ("prompt_strength", [2, 5, 10, 15])),
|
|
|
|
|
(
|
|
|
|
|
"prompt_strength[red,blue,10,15]",
|
|
|
|
|
("prompt_strength", ["red", "blue", 10, 15]),
|
|
|
|
|
),
|
|
|
|
|
],
|
|
|
|
|
)
|
|
|
|
|
def test_parse_schedule_str(schedule_str, expected):
|
|
|
|
|
cleaned_schedules = parse_schedule_str(schedule_str)
|
|
|
|
|
assert cleaned_schedules == expected
|
