imaginAIry/tests/test_prompt_schedules.py
Bryce Drennan 9ee09ac842
feature: add compilation animations (#224)
- add generation/compare gifs
2023-01-28 17:16:47 -08:00

22 lines
700 B
Python

import pytest
from imaginairy.prompt_schedules import parse_schedule_str
from imaginairy.utils import frange
@pytest.mark.parametrize(
"schedule_str,expected",
[
("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