63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from workflows import WorkflowBuilder, WorkflowError
|
|
|
|
BASE = Path(__file__).resolve().parents[1] / "workflows"
|
|
|
|
|
|
def make_builder() -> WorkflowBuilder:
|
|
return WorkflowBuilder(
|
|
BASE / "txt2img.json",
|
|
BASE / "img2img.json",
|
|
BASE / "mapping.json",
|
|
)
|
|
|
|
|
|
def test_build_txt2img() -> None:
|
|
builder = make_builder()
|
|
workflow = builder.build_txt2img(prompt="cat", negative="bad", seed=123)
|
|
|
|
assert workflow["67"]["inputs"]["text"] == "cat"
|
|
assert workflow["71"]["inputs"]["text"] == "bad"
|
|
assert workflow["70"]["inputs"]["seed"] == 123
|
|
assert workflow["81"]["class_type"] == "EmptySD3LatentImage"
|
|
assert workflow["70"]["inputs"]["denoise"] == 1.0
|
|
|
|
|
|
def test_build_img2img() -> None:
|
|
builder = make_builder()
|
|
workflow = builder.build_img2img(
|
|
prompt="make blue",
|
|
negative="bad",
|
|
seed=7,
|
|
image_name="input.png",
|
|
)
|
|
|
|
assert workflow["25"]["inputs"]["prompt"] == "make blue"
|
|
assert workflow["10"]["inputs"]["text"] == "bad"
|
|
assert workflow["3"]["inputs"]["noise_seed"] == 7
|
|
assert workflow["20"]["class_type"] == "LoadImage"
|
|
assert workflow["20"]["inputs"]["image"] == "input.png"
|
|
|
|
|
|
def test_templates_are_not_mutated() -> None:
|
|
builder = make_builder()
|
|
builder.build_txt2img(prompt="first", negative="bad", seed=1)
|
|
second = builder.build_txt2img(prompt="second", negative="bad", seed=2)
|
|
assert second["67"]["inputs"]["text"] == "second"
|
|
|
|
|
|
def test_invalid_mapping(tmp_path: Path) -> None:
|
|
mapping = json.loads((BASE / "mapping.json").read_text(encoding="utf-8"))
|
|
mapping["txt2img"]["positive"] = "999"
|
|
path = tmp_path / "mapping.json"
|
|
path.write_text(json.dumps(mapping), encoding="utf-8")
|
|
|
|
with pytest.raises(WorkflowError):
|
|
WorkflowBuilder(BASE / "txt2img.json", BASE / "img2img.json", path)
|