105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from workflows import WorkflowBuilder, WorkflowError, fit_latent_size
|
|
|
|
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["25"]["inputs"]["vae"] == ["19", 2]
|
|
assert workflow["10"]["inputs"]["text"] == "bad"
|
|
assert workflow["3"]["inputs"]["noise_seed"] == 7
|
|
assert workflow["3"]["inputs"]["steps"] == 12
|
|
assert workflow["3"]["inputs"]["sampler_name"] == "dpmpp_2m"
|
|
assert workflow["3"]["inputs"]["scheduler"] == "karras"
|
|
assert workflow["20"]["class_type"] == "LoadImage"
|
|
assert workflow["20"]["inputs"]["image"] == "input.png"
|
|
assert workflow["15"]["inputs"]["width"] == 1024
|
|
assert workflow["15"]["inputs"]["height"] == 1024
|
|
|
|
|
|
def test_build_img2img_custom_size() -> None:
|
|
builder = make_builder()
|
|
workflow = builder.build_img2img(
|
|
prompt="make blue",
|
|
negative="bad",
|
|
seed=7,
|
|
image_name="input.png",
|
|
width=832,
|
|
height=1216,
|
|
unet_name="Krea2-Q6.gguf",
|
|
)
|
|
|
|
assert workflow["15"]["inputs"]["width"] == 832
|
|
assert workflow["15"]["inputs"]["height"] == 1216
|
|
assert workflow["19"]["inputs"]["unet_name"] == "Krea2-Q6.gguf"
|
|
|
|
|
|
def test_fit_latent_size_preserves_aspect_ratio() -> None:
|
|
assert fit_latent_size(1024, 1024) == (1024, 1024)
|
|
|
|
width, height = fit_latent_size(1080, 1920)
|
|
assert width % 16 == 0
|
|
assert height % 16 == 0
|
|
assert height > width
|
|
assert max(width, height) <= 1536
|
|
|
|
width, height = fit_latent_size(4000, 3000)
|
|
assert width % 16 == 0
|
|
assert height % 16 == 0
|
|
assert width > height
|
|
|
|
|
|
def test_fit_latent_size_fallback_for_invalid_size() -> None:
|
|
assert fit_latent_size(0, 100) == (1024, 1024)
|
|
|
|
|
|
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)
|