Доработка управления моделями

This commit is contained in:
Shaman
2026-10-05 02:19:46 +03:00
parent f5ec608fc2
commit 0c5c0b423b
16 changed files with 629 additions and 252 deletions
+42 -11
View File
@@ -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()
+5
View File
@@ -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
View File
@@ -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"})
+101
View File
@@ -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
View File
@@ -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))