Доработка управления моделями
This commit is contained in:
+42
-11
@@ -1,4 +1,7 @@
|
||||
from app.config import DEFAULT_MAPPING, Settings
|
||||
import pytest
|
||||
|
||||
from app.config import Settings
|
||||
from app.errors import ConfigError
|
||||
|
||||
|
||||
def test_defaults():
|
||||
@@ -6,7 +9,7 @@ def test_defaults():
|
||||
assert s.comfy_base_url == "http://192.168.0.3:8188"
|
||||
assert s.execution_timeout == 300.0
|
||||
assert s.poll_interval == 1.5
|
||||
assert s.mapping() == DEFAULT_MAPPING
|
||||
assert s.models_file.name == "models.json"
|
||||
|
||||
|
||||
def test_env_override(monkeypatch):
|
||||
@@ -21,12 +24,40 @@ def test_env_override(monkeypatch):
|
||||
assert s.poll_interval == 0.5
|
||||
|
||||
|
||||
def test_mapping_file_override(tmp_path):
|
||||
f = tmp_path / "mapping.json"
|
||||
f.write_text('{"create": {"prompt": {"node": "1", "field": "t"}}}', encoding="utf-8")
|
||||
s = Settings(_env_file=None, mapping_file=f)
|
||||
m = s.mapping()
|
||||
assert m["create"]["prompt"] == {"node": "1", "field": "t"}
|
||||
# other sections/keys keep defaults
|
||||
assert m["edit"] == DEFAULT_MAPPING["edit"]
|
||||
assert m["create"]["seed"] == DEFAULT_MAPPING["create"]["seed"]
|
||||
def test_registry_loads_shipped_in_order():
|
||||
s = Settings(_env_file=None)
|
||||
registry = s.load_registry()
|
||||
# order preserved (first = preferred for default selection)
|
||||
assert list(registry) == ["z-image-turbo", "qwen-image-edit"]
|
||||
assert registry["z-image-turbo"].workflow.name == "z_image_turbo.json"
|
||||
assert registry["z-image-turbo"].owned_by == "comfy-adapter"
|
||||
|
||||
|
||||
def test_registry_missing_file_fails_fast(tmp_path):
|
||||
s = Settings(_env_file=None, models_file=tmp_path / "nope.json")
|
||||
with pytest.raises(ConfigError, match="not found"):
|
||||
s.load_registry()
|
||||
|
||||
|
||||
def test_registry_malformed_json_fails_fast(tmp_path):
|
||||
f = tmp_path / "models.json"
|
||||
f.write_text("{not valid", encoding="utf-8")
|
||||
s = Settings(_env_file=None, models_file=f)
|
||||
with pytest.raises(ConfigError, match="not valid JSON"):
|
||||
s.load_registry()
|
||||
|
||||
|
||||
def test_registry_empty_fails_fast(tmp_path):
|
||||
f = tmp_path / "models.json"
|
||||
f.write_text("{}", encoding="utf-8")
|
||||
s = Settings(_env_file=None, models_file=f)
|
||||
with pytest.raises(ConfigError, match="non-empty"):
|
||||
s.load_registry()
|
||||
|
||||
|
||||
def test_registry_bad_entry_fails_fast(tmp_path):
|
||||
f = tmp_path / "models.json"
|
||||
f.write_text('{"m": {"not_workflow": "x.json"}}', encoding="utf-8")
|
||||
s = Settings(_env_file=None, models_file=f)
|
||||
with pytest.raises(ConfigError, match="invalid entry"):
|
||||
s.load_registry()
|
||||
|
||||
@@ -89,6 +89,11 @@ async def test_unknown_model_400(client):
|
||||
await _assert_400(await client.post(GEN, json={"prompt": "x", "model": "dall-e-3"}))
|
||||
|
||||
|
||||
async def test_incompatible_model_400(client):
|
||||
# qwen-image-edit declares {image}, which generations cannot supply
|
||||
await _assert_400(await client.post(GEN, json={"prompt": "x", "model": "qwen-image-edit"}))
|
||||
|
||||
|
||||
async def test_bad_json_400(client):
|
||||
await _assert_400(await client.post(GEN, content=b"{not json", headers={"Content-Type": "application/json"}))
|
||||
|
||||
|
||||
+40
-25
@@ -1,15 +1,14 @@
|
||||
import pytest
|
||||
|
||||
from app.config import DEFAULT_MAPPING
|
||||
from app.config import Settings
|
||||
from app.errors import InvalidRequest
|
||||
from app.injector import build_workflow, new_seed, parse_size
|
||||
from app.templates import load_all
|
||||
from app.config import Settings
|
||||
from app.templates import load_registry
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def templates():
|
||||
return load_all(Settings(_env_file=None))
|
||||
def registry():
|
||||
return load_registry(Settings(_env_file=None))
|
||||
|
||||
|
||||
def test_parse_size():
|
||||
@@ -33,41 +32,57 @@ def test_seed_fixed_and_range():
|
||||
new_seed(-1)
|
||||
|
||||
|
||||
def test_create_injection(templates):
|
||||
create, _ = templates
|
||||
mapping = DEFAULT_MAPPING["create"]
|
||||
graph = build_workflow(create, mapping, prompt="a cat", seed=77, size="832x1216")
|
||||
def test_create_injection_and_template_untouched(registry):
|
||||
create = registry["z-image-turbo"].graph
|
||||
graph = build_workflow(create, {"prompt": "a cat", "seed": 77, "negative": "", "width": 832, "height": 1216})
|
||||
assert graph["83:27"]["inputs"]["text"] == "a cat"
|
||||
assert graph["83:13"]["inputs"]["width"] == 832
|
||||
assert graph["83:13"]["inputs"]["height"] == 1216
|
||||
assert graph["83:3"]["inputs"]["seed"] == 77
|
||||
# template untouched
|
||||
assert create["83:27"]["inputs"]["text"] == "\n"
|
||||
assert create["83:13"]["inputs"]["width"] == 1280
|
||||
# template still holds placeholders (deep copy, no mutation)
|
||||
assert create["83:27"]["inputs"]["text"] == "{prompt}"
|
||||
assert create["83:13"]["inputs"]["width"] == "{width}"
|
||||
assert create["83:3"]["inputs"]["seed"] == "{seed}"
|
||||
|
||||
|
||||
def test_edit_injection(templates):
|
||||
_, edit = templates
|
||||
mapping = DEFAULT_MAPPING["edit"]
|
||||
graph = build_workflow(edit, mapping, prompt="make it night", seed=5, image_name="adapter_x.png")
|
||||
def test_injection_is_typed(registry):
|
||||
graph = build_workflow(registry["z-image-turbo"].graph, {"prompt": "x", "seed": 7, "negative": "", "width": 64, "height": 128})
|
||||
assert isinstance(graph["83:3"]["inputs"]["seed"], int) # numeric token -> int, not "{seed}"
|
||||
assert isinstance(graph["83:13"]["inputs"]["width"], int)
|
||||
assert isinstance(graph["83:27"]["inputs"]["text"], str)
|
||||
|
||||
|
||||
def test_embedded_token_is_string_interpolated():
|
||||
graph = {"1": {"class_type": "X", "inputs": {"text": "prefix {prompt} suffix"}}}
|
||||
out = build_workflow(graph, {"prompt": "hi"})
|
||||
assert out["1"]["inputs"]["text"] == "prefix hi suffix"
|
||||
|
||||
|
||||
def test_edit_injection(registry):
|
||||
edit = registry["qwen-image-edit"].graph
|
||||
graph = build_workflow(edit, {"prompt": "make it night", "seed": 5, "negative": "", "image": "adapter_x.png"})
|
||||
assert graph["115:111"]["inputs"]["prompt"] == "make it night"
|
||||
assert graph["115:110"]["inputs"]["prompt"] == ""
|
||||
assert graph["78"]["inputs"]["image"] == "adapter_x.png"
|
||||
assert graph["115:3"]["inputs"]["seed"] == 5
|
||||
|
||||
|
||||
def test_edit_seed_varies_per_call(templates):
|
||||
_, edit = templates
|
||||
mapping = DEFAULT_MAPPING["edit"]
|
||||
def test_edit_seed_varies_per_call(registry):
|
||||
edit = registry["qwen-image-edit"].graph
|
||||
seeds = set()
|
||||
for _ in range(3):
|
||||
g = build_workflow(edit, mapping, prompt="p", seed=new_seed(None))
|
||||
g = build_workflow(edit, {"prompt": "p", "seed": new_seed(None), "negative": "", "image": "a.png"})
|
||||
seeds.add(g["115:3"]["inputs"]["seed"])
|
||||
assert len(seeds) == 3
|
||||
|
||||
|
||||
def test_unknown_mapping_node_raises(templates):
|
||||
create, _ = templates
|
||||
bad_mapping = {"prompt": {"node": "999", "field": "text"}}
|
||||
with pytest.raises(InvalidRequest, match="999"):
|
||||
build_workflow(create, bad_mapping, prompt="p", seed=1)
|
||||
def test_unrecognized_token_rejected():
|
||||
graph = {"1": {"class_type": "X", "inputs": {"text": "{foo}"}}}
|
||||
with pytest.raises(InvalidRequest, match="unrecognized token"):
|
||||
build_workflow(graph, {"prompt": "p"})
|
||||
|
||||
|
||||
def test_unfilled_token_rejected():
|
||||
graph = {"1": {"class_type": "X", "inputs": {"seed": "{seed}"}}}
|
||||
with pytest.raises(InvalidRequest, match="no value supplied"):
|
||||
build_workflow(graph, {"prompt": "p"})
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""A third model added purely by editing models.json (no code change): it shows up
|
||||
in /v1/models, is dry-inject validated at startup, and is selectable by id."""
|
||||
import contextlib
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
from app.config import REPO_ROOT, Settings
|
||||
from app.main import create_app
|
||||
from tests.conftest import COMFY, FAKE_KEY
|
||||
|
||||
PROMPT_ID = "p-int"
|
||||
|
||||
|
||||
def _extra_generation_workflow() -> dict:
|
||||
# declared = {prompt, negative, seed, width, height} -> generations-compatible, no {image}
|
||||
return {
|
||||
"1": {
|
||||
"class_type": "X",
|
||||
"inputs": {
|
||||
"text": "{prompt}",
|
||||
"neg": "{negative}",
|
||||
"sd": "{seed}",
|
||||
"w": "{width}",
|
||||
"h": "{height}",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _models_file(tmp_path) -> Settings:
|
||||
(tmp_path / "extra.json").write_text(json.dumps(_extra_generation_workflow()), encoding="utf-8")
|
||||
registry = {
|
||||
"z-image-turbo": {"workflow": str(REPO_ROOT / "workflows" / "z_image_turbo.json")},
|
||||
"extra-gen": {"workflow": "extra.json"}, # relative to models.json dir
|
||||
"qwen-image-edit": {"workflow": str(REPO_ROOT / "workflows" / "qwen_image_edit.json")},
|
||||
}
|
||||
models = tmp_path / "models.json"
|
||||
models.write_text(json.dumps(registry), encoding="utf-8")
|
||||
return Settings(
|
||||
_env_file=None,
|
||||
comfy_base_url=COMFY,
|
||||
adapter_api_key=FAKE_KEY,
|
||||
models_file=models,
|
||||
validate_templates_on_server=True,
|
||||
poll_interval=0.0,
|
||||
execution_timeout=5.0,
|
||||
)
|
||||
|
||||
|
||||
async def _with_lifespan(app):
|
||||
stack = contextlib.AsyncExitStack()
|
||||
await stack.enter_async_context(app.router.lifespan_context(app))
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
ac = await stack.enter_async_context(httpx.AsyncClient(transport=transport, base_url="http://adapter.test"))
|
||||
ac.headers["Authorization"] = f"Bearer {FAKE_KEY}"
|
||||
return stack, ac
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_third_model_listed_and_validated_without_code_change(tmp_path):
|
||||
validate_route = respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
|
||||
app = create_app(_models_file(tmp_path))
|
||||
stack, ac = await _with_lifespan(app)
|
||||
try:
|
||||
ids = [m["id"] for m in (await ac.get("/v1/models")).json()["data"]]
|
||||
assert ids == ["z-image-turbo", "extra-gen", "qwen-image-edit"]
|
||||
# every model dry-injected and validated once at startup
|
||||
assert validate_route.call_count == 3
|
||||
finally:
|
||||
await stack.aclose()
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_extra_model_selectable_and_injects(tmp_path):
|
||||
prompt_route = respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
|
||||
respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={PROMPT_ID: {"outputs": {"1": {"images": [{"filename": "o.png", "subfolder": "", "type": "output"}]}},
|
||||
"status": {"status_str": "success", "completed": True}}},
|
||||
)
|
||||
)
|
||||
respx.get(f"{COMFY}/view").mock(return_value=httpx.Response(200, content=b"\x89PNG\r\n\x1a\n"))
|
||||
|
||||
settings = _models_file(tmp_path).model_copy(update={"validate_templates_on_server": False})
|
||||
app = create_app(settings)
|
||||
stack, ac = await _with_lifespan(app)
|
||||
try:
|
||||
r = await ac.post(
|
||||
"http://adapter.test/v1/images/generations",
|
||||
json={"model": "extra-gen", "prompt": "hello", "size": "512x768", "seed": 9},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]
|
||||
assert sent["1"]["inputs"]["text"] == "hello"
|
||||
assert sent["1"]["inputs"]["sd"] == 9
|
||||
assert sent["1"]["inputs"]["w"] == 512 and sent["1"]["inputs"]["h"] == 768
|
||||
finally:
|
||||
await stack.aclose()
|
||||
+80
-24
@@ -1,7 +1,15 @@
|
||||
import pytest
|
||||
|
||||
from app.config import Settings
|
||||
from app.templates import TemplateError, check_mapping, find_broken_links, load_all
|
||||
from app.errors import InvalidRequest
|
||||
from app.injector import extract_tokens
|
||||
from app.templates import (
|
||||
TemplateError,
|
||||
find_broken_links,
|
||||
is_compatible,
|
||||
load_registry,
|
||||
select_model,
|
||||
)
|
||||
|
||||
|
||||
def test_find_broken_links_detects_missing_node():
|
||||
@@ -22,36 +30,84 @@ def test_find_broken_links_ok_on_valid_graph():
|
||||
assert find_broken_links(graph) == []
|
||||
|
||||
|
||||
def test_check_mapping_missing_node():
|
||||
graph = {"1": {"class_type": "X", "inputs": {"text": ""}}}
|
||||
mapping = {"create": {"prompt": {"node": "404", "field": "text"}}}
|
||||
errors = check_mapping(graph, mapping["create"], "create")
|
||||
assert any("not found" in e for e in errors)
|
||||
def test_extract_tokens_reports_known_and_unknown():
|
||||
graph = {
|
||||
"1": {"class_type": "X", "inputs": {"text": "{prompt}", "seed": "{seed}"}},
|
||||
"2": {"class_type": "Y", "inputs": {"note": "static {notatoken} text"}},
|
||||
}
|
||||
declared, unknown = extract_tokens(graph)
|
||||
assert declared == {"prompt", "seed"}
|
||||
assert len(unknown) == 1 and "notatoken" in unknown[0]
|
||||
|
||||
|
||||
def test_check_mapping_non_scalar_field():
|
||||
graph = {"1": {"class_type": "X", "inputs": {"text": ["2", 0]}}}
|
||||
mapping = {"create": {"prompt": {"node": "1", "field": "text"}}}
|
||||
errors = check_mapping(graph, mapping["create"], "create")
|
||||
assert any("no scalar input" in e for e in errors)
|
||||
def test_shipped_registry_loads_with_expected_tokens():
|
||||
registry = load_registry(Settings(_env_file=None))
|
||||
assert set(registry) == {"z-image-turbo", "qwen-image-edit"}
|
||||
assert registry["z-image-turbo"].declared == frozenset({"prompt", "seed", "width", "height"})
|
||||
assert registry["qwen-image-edit"].declared == frozenset({"prompt", "negative", "seed", "image"})
|
||||
|
||||
|
||||
def test_shipped_templates_load_cleanly():
|
||||
s = Settings(_env_file=None)
|
||||
create, edit = load_all(s)
|
||||
assert create["83:28"]["class_type"] == "UNETLoader"
|
||||
assert edit["115:89"]["class_type"] == "LoraLoaderModelOnly"
|
||||
def test_generations_graph_incompatible_with_edits():
|
||||
assert not is_compatible(frozenset({"prompt", "seed", "width", "height"}), "edits")
|
||||
|
||||
|
||||
def test_broken_template_blocks_startup(tmp_path):
|
||||
def test_edits_graph_incompatible_with_generations():
|
||||
assert not is_compatible(frozenset({"prompt", "negative", "seed", "image"}), "generations")
|
||||
|
||||
|
||||
def test_select_model_defaults(registry):
|
||||
assert select_model(registry, None, "generations").id == "z-image-turbo"
|
||||
assert select_model(registry, None, "edits").id == "qwen-image-edit"
|
||||
|
||||
|
||||
def test_select_model_explicit(registry):
|
||||
assert select_model(registry, "qwen-image-edit", "edits").id == "qwen-image-edit"
|
||||
|
||||
|
||||
def test_select_model_unknown_400(registry):
|
||||
with pytest.raises(InvalidRequest, match="unknown model"):
|
||||
select_model(registry, "dall-e-3", "generations")
|
||||
|
||||
|
||||
def test_select_model_incompatible_400(registry):
|
||||
with pytest.raises(InvalidRequest, match="not compatible"):
|
||||
select_model(registry, "qwen-image-edit", "generations")
|
||||
with pytest.raises(InvalidRequest, match="not compatible"):
|
||||
select_model(registry, "z-image-turbo", "edits")
|
||||
|
||||
|
||||
def test_select_model_no_suitable_400(registry):
|
||||
gen_only = {"z-image-turbo": registry["z-image-turbo"]}
|
||||
with pytest.raises(InvalidRequest, match="no model"):
|
||||
select_model(gen_only, None, "edits")
|
||||
|
||||
|
||||
def test_broken_workflow_blocks_startup(tmp_path):
|
||||
bad = tmp_path / "bad.json"
|
||||
bad.write_text('{"1": {"class_type": "X", "inputs": {"m": ["9", 0]}}}', encoding="utf-8")
|
||||
good = tmp_path / "good.json"
|
||||
good.write_text(
|
||||
'{"1": {"class_type": "X", "inputs": {"m": ["2", 0]}}, "2": {"class_type": "Y", "inputs": {}}}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
s = Settings(_env_file=None, create_template_path=bad, edit_template_path=good)
|
||||
models = tmp_path / "models.json"
|
||||
models.write_text('{"m1": {"workflow": "bad.json"}}', encoding="utf-8")
|
||||
with pytest.raises(TemplateError) as exc:
|
||||
load_all(s)
|
||||
load_registry(Settings(_env_file=None, models_file=models))
|
||||
assert "missing node '9'" in str(exc.value)
|
||||
|
||||
|
||||
def test_unknown_token_blocks_startup(tmp_path):
|
||||
wf = tmp_path / "wf.json"
|
||||
wf.write_text('{"1": {"class_type": "X", "inputs": {"text": "{bogus}"}}}', encoding="utf-8")
|
||||
models = tmp_path / "models.json"
|
||||
models.write_text('{"m1": {"workflow": "wf.json"}}', encoding="utf-8")
|
||||
with pytest.raises(TemplateError, match="unknown token"):
|
||||
load_registry(Settings(_env_file=None, models_file=models))
|
||||
|
||||
|
||||
def test_missing_workflow_file_blocks_startup(tmp_path):
|
||||
models = tmp_path / "models.json"
|
||||
models.write_text('{"m1": {"workflow": "absent.json"}}', encoding="utf-8")
|
||||
with pytest.raises(TemplateError, match="not found"):
|
||||
load_registry(Settings(_env_file=None, models_file=models))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def registry():
|
||||
return load_registry(Settings(_env_file=None))
|
||||
|
||||
Reference in New Issue
Block a user