Houmeres/LLeMbas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Draw a picture, on a ComfyUI you are running

Browse Source 
The last unbuilt capability, and built the way CLAUDE.md said it had to be: a
ToolDef reaching resolve_tools plus a permission and a capability flag, not a new
code path. The only genuinely new UI is one branch in the transcript.

services/images/ is three modules. comfy.py speaks HTTP -- submit, poll /history,
fetch the PNG, /free, and an /object_info discovery for the admin page only.
Polled and not socketed, because holding a connection open for the length of a
generation is the live-connection state the whole ssh.py design forbids, and the
thing being waited for takes tens of seconds anyway. The base URL is exempt from
the SSRF guard by construction, exactly as Connection.base_url and the audio
endpoints are -- said out loud in the docstring, because a default of
127.0.0.1:8188 is precisely the shape that guard exists to refuse and therefore
reads as a hole rather than a decision.

workflow.py fills a template, and the one thing that matters is that it walks the
parsed JSON rather than the text of it. A value that is exactly "{{steps}}"
becomes the number 20; ComfyUI validates types and refuses the string. A
placeholder inside a longer string is still text, which is what makes
"{{prompt}}, masterpiece" work -- and text substitution would additionally mean a
prompt containing a quotation mark produced a document that no longer parses, on
the one input guaranteed to hold arbitrary text. Which node holds the prompt is
the administrator's statement rather than a guess from node types: sniffing for
the first CLIPTextEncode works on the shipped workflow and on nothing else, and
swaps positive for negative the first time somebody reorders them. seed has no
fixed default, because one would make every unspecified generation identical and
make the retry loop redraw the same rejected picture four times.

tool.py is one call, one finished image. Returning every attempt to the
conversation would cost a round each, make the ceiling advisory rather than
enforced, and walk the reader past every reject -- so the reviewer lives inside
the tool and is asked about *bytes*: an attempt about to be discarded should not
leave an Attachment behind, so it sees a downscaled preview built in memory and
only the kept image is written. Anything that goes wrong in review is a keep;
losing a picture because a judging request timed out would be the check
destroying the thing it was checking. The last attempt is kept whatever the
verdict, so a request always produces something. Rejects are recorded, not
stored.

Preserve VRAM unloads the chat's own connection and nothing else, because the
memory being freed belongs to one machine: local llama-swap answers GET /unload,
and a box on the network has no reason to be unloaded when ComfyUI wants memory
here. The swap goes round the review rather than round the tool, which costs two
model loads per retry -- so the two settings are independent and the page warns
when both are on. Nothing loads the LLM back: the reply's next request does, and
that step exists in the description and not in the code, so the code says so.

Two rules elsewhere had to be drawn for the first time. message_payload sends
images only on user turns -- no assistant message had ever carried one, and the
moment one does the multimodal list form on an assistant turn is rejected by
OpenAI and most local runners, breaking every later turn in the chat. And
files.store gained keep_original, because _process_image turns anything without
alpha into JPEG q85 at 1400px: right for a phone photo, a visible loss on the one
output this feature exists to produce.

/image sends the ordinary message with force_tool, which becomes tool_choice for
the first round only -- left in place the reply would draw a picture, be asked
again, and draw another. FORCEABLE_TOOLS is an allow list because the name is
read off a form.

ToolContext gained chat_id, and that fixed a tool nobody had ever successfully
run: _run_scratch_write read context.chat_id on a dataclass with no such field,
so every call raised AttributeError, swallowed by run_tool's blanket except into
"the scratch_write tool failed" -- indistinguishable from a model calling it
wrongly. The test that existed asserted the family and the risk, which are
properties of the declaration rather than of the code.

Verified against the real ComfyUI 0.27.0 on this machine rather than against
documentation: every endpoint shape here was read off it, a generation ran end to
end through the client, the reviewer was shown a matching and a mismatched prompt
and answered KEEP and RETRY correctly, and the unload hook fired for the local
llama-swap and not for the remote box.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
2026-08-05 14:13:19 +02:00
parent 9f5ff72e32
commit 47d1ddbc3c
38 changed files with 3958 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_admin_images.py
+272
View File
@@ -0,0 +1,272 @@
"""The image generation admin page, and the workflows behind it."""
from __future__ import annotations
import json
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select
import lembas
from lembas.db.models import ImageWorkflow
from lembas.services import settings_store
from lembas.services.crypto import UNCHANGED_SENTINEL, decrypt
BASE_TEXT = (Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text(
encoding="utf-8"
)
BASE = json.loads(BASE_TEXT)
@pytest.fixture
def admin(client: TestClient, db, registered):
from lembas.db.models import User
db.query(User).update({"role": "admin"})
db.commit()
return client
def _settings(client, **overrides):
data = {
"enabled": "true",
"base_url": "http://comfy.test:8188",
"api_key": "",
"timeout": "600",
"checkpoints": "sd.safetensors\nother.safetensors",
"default_workflow_id": "",
"review_enabled": "true",
"review_model_id": "",
"max_tries": "4",
"preserve_vram": "",
"instructions": "",
}
data.update(overrides)
return client.post("/admin/images", data=data, follow_redirects=False)
# --- The page ------------------------------------------------------------------
def test_the_page_renders_and_is_in_the_nav(admin):
page = admin.get("/admin/images").text
assert 'href="/admin/images"' in page
assert "ComfyUI" in page
def test_settings_are_saved(admin, db):
assert _settings(admin).status_code == 303
values = settings_store.images(db)
assert values["enabled"] is True
assert values["base_url"] == "http://comfy.test:8188"
assert values["checkpoints"] == ["sd.safetensors", "other.safetensors"]
assert values["review_enabled"] is True
def test_a_trailing_slash_is_stripped_from_the_url(admin, db):
_settings(admin, base_url="http://comfy.test:8188/")
assert settings_store.images(db)["base_url"] == "http://comfy.test:8188"
def test_an_absent_checkbox_is_off(admin, db):
"""An unticked box is simply missing from a form post -- that absence *is*
the off signal."""
data = {"base_url": "http://x", "timeout": "600", "max_tries": "4"}
admin.post("/admin/images", data=data, follow_redirects=False)
values = settings_store.images(db)
assert values["enabled"] is False
assert values["review_enabled"] is False
assert values["preserve_vram"] is False
def test_the_numbers_are_clamped(admin, db):
_settings(admin, max_tries="500", timeout="1")
values = settings_store.images(db)
assert values["max_tries"] == 10
assert values["timeout"] == 10.0
def test_max_tries_never_reaches_zero(admin, db):
"""Zero would mean the tool generates nothing and reports success. Unlike
the zeroes elsewhere in the settings, there is no reading of it anybody
wants."""
_settings(admin, max_tries="0")
assert settings_store.images(db)["max_tries"] == 1
def test_the_key_survives_a_save_that_did_not_touch_it(admin, db):
_settings(admin, api_key="secret-key")
assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key"
_settings(admin, api_key=UNCHANGED_SENTINEL, instructions="something else")
assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key"
def test_an_emptied_key_is_removed(admin, db):
_settings(admin, api_key="secret-key")
_settings(admin, api_key="")
assert settings_store.images(db)["api_key_encrypted"] == ""
def test_the_key_is_never_rendered(admin, db):
_settings(admin, api_key="secret-key")
assert "secret-key" not in admin.get("/admin/images").text
def test_testing_without_a_url_says_so(admin, db):
_settings(admin, base_url="")
assert "Set a base URL first" in admin.post("/admin/images/test").text
# --- Workflows -----------------------------------------------------------------
def test_new_is_not_read_as_an_id(admin):
"""FastAPI matches in registration order, so `/workflows/new` has to be
registered before `/workflows/{id}` or "new" is an id and 404s. This has
been a bug twice in this codebase."""
response = admin.get("/admin/images/workflows/new")
assert response.status_code == 200
assert "Export (API)" in response.text
def test_the_default_workflow_is_offered_to_start_from(admin):
"""Somebody setting this up for the first time should not have to find a
working API-format document before they can try anything."""
page = admin.get("/admin/images/workflows/new").text
assert "CheckpointLoaderSimple" in page
assert "{{prompt}}" in page
def test_a_workflow_is_created(admin, db):
response = admin.post(
"/admin/images/workflows",
data={
"slug": "sdxl",
"name": "SDXL",
"description": "photo",
"workflow": BASE_TEXT,
"enabled": "true",
},
follow_redirects=False,
)
assert response.status_code == 303
row = db.scalar(select(ImageWorkflow).where(ImageWorkflow.slug == "sdxl"))
assert row is not None
assert row.workflow_json["4"]["inputs"]["ckpt_name"] == "{{model}}"
def test_broken_json_is_refused_with_a_sentence(admin, db):
"""A sentence and not a 422: losing forty lines of pasted JSON to a
validation error is not a thing to do to somebody, so the form comes back
with what was typed still in it."""
response = admin.post(
"/admin/images/workflows",
data={"slug": "x", "name": "X", "workflow": "{not json"},
)
assert response.status_code == 200
assert "not valid JSON" in response.text
assert "{not json" in response.text, "and what was typed is still there"
assert db.scalar(select(ImageWorkflow)) is None
def test_a_workflow_with_no_prompt_placeholder_is_refused(admin, db):
"""It would draw the same picture whatever anybody typed, and would look
like a broken model rather than an unparameterised template."""
without = json.loads(BASE_TEXT)
without["6"]["inputs"]["text"] = "a fixed prompt"
response = admin.post(
"/admin/images/workflows",
data={"slug": "x", "name": "X", "workflow": json.dumps(without)},
)
assert "{{prompt}}" in response.text
assert "every image would be the same" in response.text
assert db.scalar(select(ImageWorkflow)) is None
def test_an_unknown_placeholder_is_refused(admin, db):
response = admin.post(
"/admin/images/workflows",
data={
"slug": "x",
"name": "X",
"workflow": json.dumps({"1": {"a": "{{prompt}} {{lora}}"}}),
},
)
assert "lora" in response.text
assert db.scalar(select(ImageWorkflow)) is None
def test_a_ui_format_export_is_refused_helpfully(admin, db):
"""The two ComfyUI formats look alike enough that this is the mistake
everybody makes first, and "invalid" would not tell them which button to
press instead."""
response = admin.post(
"/admin/images/workflows",
data={"slug": "x", "name": "X", "workflow": "[]"},
)
assert "Export (API)" in response.text
def test_a_duplicate_slug_is_refused(admin, db):
for _ in range(2):
response = admin.post(
"/admin/images/workflows",
data={"slug": "dup", "name": "Dup", "workflow": BASE_TEXT},
)
assert "already a workflow" in response.text
assert len(list(db.scalars(select(ImageWorkflow)))) == 1
def test_a_rejected_edit_leaves_the_stored_row_alone(admin, db):
admin.post(
"/admin/images/workflows",
data={"slug": "keep", "name": "Keep", "workflow": BASE_TEXT},
)
row = db.scalar(select(ImageWorkflow))
admin.post(
f"/admin/images/workflows/{row.id}",
data={"slug": "keep", "name": "Renamed", "workflow": "{broken"},
)
db.expire_all()
row = db.scalar(select(ImageWorkflow))
assert row.name == "Keep", "validated against a draft, so nothing was written"
def test_a_workflow_is_deleted(admin, db):
admin.post(
"/admin/images/workflows",
data={"slug": "gone", "name": "Gone", "workflow": BASE_TEXT},
)
row = db.scalar(select(ImageWorkflow))
admin.post(f"/admin/images/workflows/{row.id}/delete", follow_redirects=False)
db.expire_all()
assert db.scalar(select(ImageWorkflow)) is None
def test_only_an_administrator_can_reach_it(client: TestClient, db, registered):
"""The first account registered is an administrator, so this needs a second
one -- the `test_admin_tools` fixture's reason for existing."""
from lembas.db.models import User
client.post("/auth/logout")
client.post(
"/auth/register",
data={"name": "Sam", "email": "sam@shire.test", "password": "correct horse battery"},
follow_redirects=False,
)
user = db.scalar(select(User).where(User.email == "sam@shire.test"))
user.role = "user"
user.active = True
db.commit()
client.post(
"/auth/login",
data={"email": "sam@shire.test", "password": "correct horse battery"},
follow_redirects=False,
)
assert client.get("/admin/images").status_code in (302, 303, 403, 404)
assert client.post("/admin/images/workflows", data={}).status_code in (302, 303, 403, 404)
tests/test_canvas_stream.py
+35
View File
@@ -200,6 +200,41 @@ def test_scratch_write_opens_its_tab(db, user_id):
assert tool.risk == "read"
async def test_scratch_write_actually_runs(db, user_id):
"""It did not, for the whole life of the feature.
The runner read `context.chat_id` off a `ToolContext` that had no such
field, so every call raised `AttributeError` -- swallowed by `run_tool`'s
blanket except into "the scratch_write tool failed", which is
indistinguishable from a model calling it wrongly. Nothing exercised the
runner: the test above asserts the family and the risk, which are
attributes of the declaration rather than of the code.
"""
from lembas.db.models import Chat
from lembas.services.tools import REGISTRY, ToolContext
chat = Chat(user_id=user_id, model_id="m")
db.add(chat)
db.commit()
outcome = await REGISTRY["scratch_write"].run(
ToolContext(owner_id=user_id, chat_id=chat.id), {"text": "hello", "mode": "replace"}
)
assert outcome.event["status"] == "ok"
assert outcome.event.get("canvas", {}).get("source") == "scratch"
async def test_scratch_write_without_a_chat_says_so(db, user_id):
"""Which is what the `chat_id` check was always for."""
from lembas.services.tools import REGISTRY, ToolContext
outcome = await REGISTRY["scratch_write"].run(
ToolContext(owner_id=user_id), {"text": "hello"}
)
assert outcome.event["status"] == "error"
def test_the_event_survives_into_the_stored_transcript():
"""Harmless and mildly useful: `_tool_activity.html` reads named keys."""
event = {"name": "file_read", "canvas": {"key": "agent:/a.py"}}
tests/test_images_chat.py
+297
View File
@@ -0,0 +1,297 @@
"""Image generation where it meets the conversation.
Three seams, each of which fails silently if it is got wrong: how a generated
image is bound to the reply that made it, why it is never replayed on an
assistant turn, and what `/image` actually sends.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select
import lembas
from lembas.db.models import ROLE_ASSISTANT, ROLE_USER, Attachment, Chat, Connection, Message, Model
from lembas.services import chat as chat_service
from lembas.services import files as files_service
from lembas.services import generation as generation_service
from lembas.services.crypto import encrypt
PNG = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00"
b"\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178U\x00\x00\x00"
b"\x00IEND\xaeB`\x82"
)
@pytest.fixture
def vision_chat(db, user_id):
connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt(""))
db.add(connection)
db.commit()
db.add(
Model(
connection_id=connection.id,
model_id="m",
capabilities_json={"vision": True, "tools": True},
)
)
chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
db.add(chat)
db.commit()
return chat
# --- Storing what was drawn ----------------------------------------------------
def test_a_generated_image_is_kept_as_it_arrived(db, user_id, vision_chat):
"""`_process_image` transcodes to JPEG q85 and downscales to 1400px, which
is right for a phone photo and a visible loss on the one output this feature
exists to produce."""
attachment = files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="out.png",
keep_original=True,
)
assert attachment.media_type == "image/png"
assert files_service.stored_path(attachment.stored_name).read_bytes() == PNG
def test_an_ordinary_upload_is_still_processed(db, user_id, vision_chat):
"""The flag is opt-in, and the protection it skips still applies to
everything that arrives from outside."""
attachment = files_service.store(
db, user_id=user_id, chat_id=vision_chat.id, payload=PNG, filename="photo.png"
)
assert attachment.media_type == "image/jpeg"
def test_a_corrupt_image_is_still_refused(db, user_id, vision_chat):
"""What `keep_original` skips is the resize and the transcode, not the
check that this is an image at all."""
with pytest.raises(files_service.FileError):
files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=b"\x89PNG\r\n\x1a\n" + b"rubbish",
filename="broken.png",
keep_original=True,
)
def test_the_loop_binds_the_image_to_the_reply(db, user_id, vision_chat):
"""A runner cannot write the message row -- `_persist` is the single writer
-- so the runner makes the attachment and the loop says which turn owns it,
exactly as it already does for a canvas tab."""
message = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "here it is")
attachment = files_service.store(
db, user_id=user_id, chat_id=vision_chat.id, payload=PNG, filename="a.png"
)
generation_service._bind_attachments(db, vision_chat, message, [attachment.id])
db.commit()
db.refresh(attachment)
assert attachment.message_id == message.id
def test_binding_refuses_a_row_from_another_chat(db, user_id, vision_chat):
"""The ids arrive on a tool event, which is a dict a runner built. Without
the narrowing a forged one would pull somebody else's file into this
conversation -- the reason `files.claim` is scoped the same way."""
other = Chat(user_id=user_id, model_id="m")
db.add(other)
db.commit()
elsewhere = files_service.store(
db, user_id=user_id, chat_id=other.id, payload=PNG, filename="a.png"
)
message = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "x")
generation_service._bind_attachments(db, vision_chat, message, [elsewhere.id])
db.commit()
db.refresh(elsewhere)
assert elsewhere.message_id is None
def test_binding_refuses_a_row_already_bound(db, user_id, vision_chat):
first = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "one")
second = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "two")
attachment = files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=first.id,
)
generation_service._bind_attachments(db, vision_chat, second, [attachment.id])
db.commit()
db.refresh(attachment)
assert attachment.message_id == first.id
# --- What reaches the model ----------------------------------------------------
def test_an_image_on_an_assistant_turn_is_never_replayed(db, user_id, vision_chat):
"""The multimodal list form on an `assistant` turn is rejected outright by
OpenAI and by most local runners -- and it would break not that turn but
every later one in the chat."""
reply = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "here it is")
files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=reply.id,
)
db.refresh(reply)
payload = chat_service.message_payload(reply, vision=True)
assert isinstance(payload["content"], str), "no image_url parts on an assistant turn"
assert payload["content"] == "here it is"
def test_an_image_a_person_sent_still_reaches_the_model(db, user_id, vision_chat):
"""The rule narrows assistant turns and nothing else."""
turn = chat_service.create_message(db, vision_chat, ROLE_USER, "what is this?")
files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=turn.id,
)
db.refresh(turn)
payload = chat_service.message_payload(turn, vision=True)
assert isinstance(payload["content"], list)
assert any(part["type"] == "image_url" for part in payload["content"])
def test_a_generated_image_still_renders_in_the_bubble(
client: TestClient, db, user_id, vision_chat
):
"""It is an attachment on the assistant message, and `_message.html` renders
attachments for either role -- so the reader sees it without the template
learning anything new."""
reply = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "here it is")
attachment = files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=reply.id,
)
page = client.get(f"/chat/{vision_chat.id}").text
assert f"/api/files/{attachment.id}/content" in page
# --- Forcing the tool ----------------------------------------------------------
def test_the_forced_tool_reaches_the_request(db, vision_chat):
tools = [{"type": "function", "function": {"name": "image_generate", "parameters": {}}}]
body = chat_service.build_request(db, vision_chat, tools=tools, force_tool="image_generate")
assert body["tool_choice"] == {
"type": "function",
"function": {"name": "image_generate"},
}
def test_nothing_is_forced_by_default(db, vision_chat):
"""A provider strict about unknown parameters must see exactly the request
it always did until somebody types a slash command."""
tools = [{"type": "function", "function": {"name": "image_generate", "parameters": {}}}]
assert "tool_choice" not in chat_service.build_request(db, vision_chat, tools=tools)
def test_a_tool_that_was_not_offered_cannot_be_forced(db, vision_chat):
"""`resolve_tools` still decides what exists. Forcing something absent from
the array is a request most endpoints reject outright."""
tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}]
body = chat_service.build_request(db, vision_chat, tools=tools, force_tool="image_generate")
assert "tool_choice" not in body
def test_forcing_needs_a_tools_array_at_all(db, vision_chat):
assert "tool_choice" not in chat_service.build_request(
db, vision_chat, force_tool="image_generate"
)
def test_the_endpoint_only_accepts_names_from_the_allow_list(
client: TestClient, db, registered, vision_chat, monkeypatch
):
"""This becomes `tool_choice`, so a name read straight off a form would let
anyone who can send a message decide what the model must do next."""
from lembas.api import chats as chats_api
seen: dict = {}
monkeypatch.setattr(
chats_api.generation_service,
"ensure",
lambda chat_id, message_id, *, force_tool="": seen.update(force_tool=force_tool),
)
client.post(
f"/api/chats/{vision_chat.id}/messages",
data={"content": "hello", "force_tool": "shell_run"},
)
assert seen["force_tool"] == "", "not on the list, so not forced"
# That first turn left an unfinished assistant row behind, and a second
# message while one is in flight is *queued* rather than sent -- so it would
# never reach `ensure` at all. Finish it first.
db.query(Message).filter(Message.complete.is_(False)).update({"complete": True})
db.commit()
client.post(
f"/api/chats/{vision_chat.id}/messages",
data={"content": "a bicycle", "force_tool": "image_generate"},
)
assert seen["force_tool"] == "image_generate"
def test_the_image_command_is_in_the_table():
"""`/help` reads this list, so a command missing from it is one nobody can
discover -- the direction this actually rots."""
source = (Path(lembas.__file__).parent / "web/static/js/commands.js").read_text(
encoding="utf-8"
)
assert 'name: "image"' in source
assert 'body.append("force_tool", "image_generate")' in source
assert "htmx.process" in source, "or the reply's sse-connect is inert markup"
def test_a_forced_round_does_not_force_the_next_one(db, vision_chat):
"""Leaving `tool_choice` in place would make every round call the tool
again: draw a picture, be asked again, draw another."""
source = (Path(lembas.__file__).parent / "services/generation.py").read_text(encoding="utf-8")
assert 'payload.pop("tool_choice", None)' in source
def test_a_queued_turn_is_not_lost_when_it_was_forced(db, user_id, vision_chat):
"""`/image` typed while a reply is streaming queues like anything else. The
forcing is on the generation, so a queued turn simply arrives unforced --
which is right: by then the model has the words and the context."""
db.add(Message(chat_id=vision_chat.id, role=ROLE_ASSISTANT, content="", complete=False))
db.commit()
from lembas.api import chats as chats_api
assert chats_api._reply_in_flight(db, vision_chat) is True
assert db.scalar(select(Attachment)) is None
tests/test_images_comfy.py
+281
View File
@@ -0,0 +1,281 @@
"""The ComfyUI client, against recorded shapes.
Every response body here was read off a real ComfyUI 0.27.0 while this was being
written, rather than reconstructed from documentation -- including the one that
matters most, `/history` answering with an empty object while the job is still
queued.
"""
from __future__ import annotations
import httpx
import pytest
from lembas.services.images import comfy
CONFIG = comfy.Config(base_url="http://comfy.test:8188", timeout=5.0)
HISTORY_DONE = {
"p1": {
"status": {"status_str": "success", "completed": True, "messages": []},
"outputs": {
"9": {"images": [{"filename": "LLeMbas_00001_.png", "subfolder": "", "type": "output"}]}
},
}
}
def _handler(routes):
"""Answer by path, and record what was asked."""
def handle(request: httpx.Request) -> httpx.Response:
for path, responder in routes.items():
if request.url.path.startswith(path):
return responder(request)
return httpx.Response(404, json={"error": "no route"})
return handle
# --- Submitting ----------------------------------------------------------------
async def test_a_workflow_is_queued_and_its_id_returned(mock_http):
seen: dict = {}
def prompt(request):
seen["body"] = request.read().decode()
return httpx.Response(200, json={"prompt_id": "p1", "number": 3, "node_errors": {}})
mock_http(_handler({"/prompt": prompt}))
assert await comfy.submit(CONFIG, {"3": {"class_type": "KSampler"}}) == "p1"
assert "KSampler" in seen["body"]
assert "client_id" in seen["body"], "ComfyUI keys its progress socket on this"
async def test_a_graph_error_names_the_node(mock_http):
""" "Invalid prompt" against a twelve-node document says nothing. The
commonest cause by far is a checkpoint that does not exist on that machine,
and the node id is what points at it."""
mock_http(
_handler(
{
"/prompt": lambda r: httpx.Response(
400,
json={
"error": {"type": "prompt_outputs_failed_validation"},
"node_errors": {
"4": {"errors": [{"message": "value not in list: ckpt_name"}]}
},
},
)
}
)
)
with pytest.raises(comfy.ComfyError) as caught:
await comfy.submit(CONFIG, {})
assert "node 4" in caught.value.message
assert "ckpt_name" in caught.value.message
async def test_an_accepted_workflow_with_no_id_is_refused(mock_http):
mock_http(_handler({"/prompt": lambda r: httpx.Response(200, json={"number": 1})}))
with pytest.raises(comfy.ComfyError):
await comfy.submit(CONFIG, {})
async def test_an_unreachable_comfyui_says_so(mock_http):
def refuse(request):
raise httpx.ConnectError("nope", request=request)
mock_http(_handler({"/prompt": refuse}))
with pytest.raises(comfy.ComfyError) as caught:
await comfy.submit(CONFIG, {})
assert "Could not reach ComfyUI" in caught.value.message
# --- Waiting -------------------------------------------------------------------
async def test_an_empty_history_means_not_yet(mock_http, monkeypatch):
"""`/history/{id}` is empty while the job is queued and gains the whole
record when it ends -- so empty is "not yet", not "nothing". Reading it the
other way makes every generation fail the instant it is submitted."""
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
calls = {"n": 0}
def history(request):
calls["n"] += 1
return httpx.Response(200, json=HISTORY_DONE if calls["n"] >= 3 else {})
mock_http(_handler({"/history": history}))
refs = await comfy.await_images(CONFIG, "p1")
assert calls["n"] == 3, "it kept asking"
assert [ref.filename for ref in refs] == ["LLeMbas_00001_.png"]
async def test_waiting_gives_up_eventually(mock_http, monkeypatch):
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
mock_http(_handler({"/history": lambda r: httpx.Response(200, json={})}))
with pytest.raises(comfy.ComfyError) as caught:
await comfy.await_images(comfy.Config(base_url=CONFIG.base_url, timeout=0.05), "p1")
assert "did not finish" in caught.value.message
async def test_a_failed_workflow_is_not_an_empty_one(mock_http, monkeypatch):
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
mock_http(
_handler(
{
"/history": lambda r: httpx.Response(
200,
json={
"p1": {"status": {"status_str": "error", "completed": True}, "outputs": {}}
},
)
}
)
)
with pytest.raises(comfy.ComfyError):
await comfy.await_images(CONFIG, "p1")
async def test_images_are_read_from_every_node(mock_http, monkeypatch):
"""A template is somebody else's document and may save from a node called
anything, or from two of them. Looking for `SaveImage` by name works on the
shipped workflow and on nothing else."""
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
mock_http(
_handler(
{
"/history": lambda r: httpx.Response(
200,
json={
"p1": {
"status": {"status_str": "success", "completed": True},
"outputs": {
"12": {
"images": [
{"filename": "a.png", "subfolder": "s", "type": "output"}
]
},
"40": {"images": [{"filename": "b.png"}]},
},
}
},
)
}
)
)
refs = await comfy.await_images(CONFIG, "p1")
assert [r.filename for r in refs] == ["a.png", "b.png"]
assert refs[0].subfolder == "s"
# --- Fetching ------------------------------------------------------------------
async def test_the_image_comes_back_as_bytes(mock_http):
seen: dict = {}
def view(request):
seen["params"] = dict(request.url.params)
return httpx.Response(200, content=b"\x89PNG\r\n\x1a\n" + b"x" * 100)
mock_http(_handler({"/view": view}))
payload = await comfy.fetch_image(CONFIG, comfy.Ref("a.png", "sub", "output"))
assert payload.startswith(b"\x89PNG")
assert seen["params"] == {"filename": "a.png", "subfolder": "sub", "type": "output"}
async def test_an_oversized_image_is_refused(mock_http, monkeypatch):
"""The one place an external service hands back raw bytes that are written
to disk. Neither `audio.speak` nor `openai_client` has a cap to copy."""
monkeypatch.setattr(comfy, "MAX_IMAGE_BYTES", 64)
mock_http(_handler({"/view": lambda r: httpx.Response(200, content=b"x" * 4096)}))
with pytest.raises(comfy.ComfyError):
await comfy.fetch_image(CONFIG, comfy.Ref("a.png"))
async def test_an_empty_file_is_refused(mock_http):
mock_http(_handler({"/view": lambda r: httpx.Response(200, content=b"")}))
with pytest.raises(comfy.ComfyError):
await comfy.fetch_image(CONFIG, comfy.Ref("a.png"))
# --- Freeing -------------------------------------------------------------------
async def test_freeing_never_raises(mock_http):
"""It runs on the way out of a generation that has already produced its
image. Failing the whole tool because a memory hint was refused would be
turning a tidy-up into an error."""
def boom(request):
raise httpx.ConnectError("gone", request=request)
mock_http(_handler({"/free": boom}))
await comfy.free(CONFIG) # must not raise
async def test_freeing_asks_for_both(mock_http):
seen: dict = {}
def free(request):
seen["body"] = request.read().decode()
return httpx.Response(200, json={})
mock_http(_handler({"/free": free}))
await comfy.free(CONFIG)
assert "unload_models" in seen["body"] and "free_memory" in seen["body"]
# --- Discovery -----------------------------------------------------------------
async def test_discovery_reads_the_option_lists(mock_http):
"""The shape is `{node: {input: {required: {field: [[...values], {meta}]}}}}`
-- a list whose first element is the options."""
mock_http(
_handler(
{
"/object_info/CheckpointLoaderSimple": lambda r: httpx.Response(
200,
json={
"CheckpointLoaderSimple": {
"input": {
"required": {"ckpt_name": [["a.safetensors", "b.safetensors"]]}
}
}
},
),
"/object_info/KSampler": lambda r: httpx.Response(
200,
json={
"KSampler": {
"input": {
"required": {
"sampler_name": [["euler", "dpmpp_2m"]],
"scheduler": [["normal", "karras"]],
}
}
}
},
),
}
)
)
checkpoints, samplers, schedulers = await comfy.discover(CONFIG)
assert checkpoints == ["a.safetensors", "b.safetensors"]
assert samplers == ["euler", "dpmpp_2m"]
assert schedulers == ["normal", "karras"]
async def test_a_custom_node_pack_that_changes_the_shape_is_survived(mock_http):
"""It is somebody else's schema. An empty list is a page that says "none
found"; an exception is a page that says nothing at all."""
mock_http(
_handler(
{
"/object_info": lambda r: httpx.Response(200, json={"Whatever": {"input": {}}}),
}
)
)
assert await comfy.discover(CONFIG) == ([], [], [])
tests/test_images_tool.py
+395
View File
@@ -0,0 +1,395 @@
"""Generating an image: what is offered, and what the loop does.
The gate tests matter more than they look. Image generation is the first tool
whose availability depends on an instance being *configured* as well as
permitted, and the branch it needs in `_family_allowed` is easy to leave out --
without it the family falls through to the last line and silently requires
`library.use`, which has nothing to do with drawing a picture.
"""
from __future__ import annotations
import json
from pathlib import Path
import pytest
import lembas
from lembas.db.models import Chat, Connection, ImageWorkflow, Model, User
from lembas.services import settings_store
from lembas.services import tools as tools_service
from lembas.services.crypto import encrypt
from lembas.services.images import comfy
from lembas.services.images import tool as image_tool
from lembas.services.tools import ToolContext
BASE = json.loads(
(Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text(
encoding="utf-8"
)
)
PNG = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00"
b"\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178U\x00\x00\x00"
b"\x00IEND\xaeB`\x82"
)
@pytest.fixture
def configured(db, user_id):
"""An instance that can draw: a ComfyUI, a checkpoint and a workflow."""
connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt(""))
db.add(connection)
db.commit()
model = Model(
connection_id=connection.id,
model_id="m",
capabilities_json={"tools": True, "tool_image": True},
)
db.add(model)
db.add(ImageWorkflow(slug="base", name="Base", description="general", workflow_json=BASE))
db.commit()
chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
db.add(chat)
settings_store.update(
db,
{
"enabled": True,
"base_url": "http://comfy.test:8188",
"checkpoints": ["sd.safetensors", "other.safetensors"],
"review_enabled": False,
},
key=settings_store.IMAGES,
)
db.commit()
return chat
def _context(db, user_id, chat) -> ToolContext:
return tools_service.context_for(db, db.get(User, user_id), chat)
# --- Whether it is offered at all ----------------------------------------------
def test_it_is_offered_once_the_instance_can_draw(db, user_id, configured):
offered = tools_service.resolve_tools(db, configured, db.get(User, user_id))
assert "image_generate" in offered.by_name
def test_it_is_withheld_without_a_comfyui(db, user_id, configured):
settings_store.update(db, {"base_url": ""}, key=settings_store.IMAGES)
assert (
"image_generate"
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
)
def test_it_is_withheld_with_no_checkpoints(db, user_id, configured):
"""A model naming a checkpoint that does not exist gets a refusal from
ComfyUI and spends a round finding out, so an instance that has listed none
should not offer the tool -- the same rule `skill_get` follows for an empty
library."""
settings_store.update(db, {"checkpoints": []}, key=settings_store.IMAGES)
assert (
"image_generate"
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
)
def test_it_is_withheld_when_the_feature_is_off(db, user_id, configured):
settings_store.update(db, {"enabled": False}, key=settings_store.IMAGES)
assert (
"image_generate"
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
)
def test_it_is_withheld_without_the_permission(db, user_id, configured):
user = db.get(User, user_id)
user.role = "user" # administrators are given every permission
settings_store.update(db, {"default_permissions": {"tools.image": False}})
db.commit()
assert "image_generate" not in tools_service.resolve_tools(db, configured, user).by_name
def test_it_does_not_need_the_library_permission(db, user_id, configured):
"""The trap the explicit branch in `_family_allowed` exists to avoid: with
no branch the family falls to the last line and requires `library.use`,
which is a coincidence of naming rather than a rule."""
user = db.get(User, user_id)
user.role = "user"
settings_store.update(db, {"default_permissions": {"library.use": False}})
db.commit()
assert "image_generate" in tools_service.resolve_tools(db, configured, user).by_name
def test_the_model_capability_can_withhold_it(db, user_id, configured):
model = db.scalar(Model.__table__.select().where(Model.model_id == "m"))
db.query(Model).filter(Model.model_id == "m").update(
{"capabilities_json": {"tools": True, "tool_image": False}}
)
db.commit()
assert model is not None
assert (
"image_generate"
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
)
def test_a_chat_can_switch_it_off(db, user_id, configured):
configured.scope_json = {"families": {"image": False}}
db.commit()
assert (
"image_generate"
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
)
def test_the_registry_knows_it_so_the_guidance_applies(db):
"""`harness._families` maps an offered name back to a family through
`registry(db)`. A tool missing from there resolves to no family, and its
fragment never appears -- the omission that cost custom tools their guidance
once already."""
assert tools_service.registry(db)["image_generate"].family == "image"
assert "image" in tools_service.families(db)
# --- The schema ----------------------------------------------------------------
def test_the_schema_offers_this_instance_s_own_choices(db, user_id, configured):
schema = image_tool.schema_for(db, settings_store.images(db))
assert schema["properties"]["model"]["enum"] == ["sd.safetensors", "other.safetensors"]
assert schema["properties"]["template"]["enum"] == ["base"]
assert "general" in schema["properties"]["template"]["description"]
def test_only_the_prompt_is_required_and_it_is_first(db, user_id, configured):
"""`tools.parse_arguments` puts the raw string into the first *required*
parameter when a model emits arguments that are not valid JSON -- common
with small models. This way that degrades into a prompt rather than a seed.
"""
schema = image_tool.schema_for(db, settings_store.images(db))
assert schema["required"] == ["prompt"]
assert next(iter(schema["properties"])) == "prompt"
def test_the_samplers_are_not_an_enum(db, user_id, configured):
"""Forty-four of them, on every request, forever, to prevent a mistake worth
one sentence of correction."""
schema = image_tool.schema_for(db, settings_store.images(db))
assert "enum" not in schema["properties"]["sampler"]
# --- The loop ------------------------------------------------------------------
class FakeComfy:
"""Stands in for the far side, counting what it was asked to do."""
def __init__(self):
self.submits: list[dict] = []
self.frees = 0
async def submit(self, config, wf):
self.submits.append(wf)
return f"p{len(self.submits)}"
async def await_images(self, config, prompt_id):
return [comfy.Ref(f"{prompt_id}.png")]
async def fetch_image(self, config, ref):
return PNG
async def free(self, config):
self.frees += 1
@pytest.fixture
def fake(monkeypatch):
stub = FakeComfy()
for name in ("submit", "await_images", "fetch_image", "free"):
monkeypatch.setattr(comfy, name, getattr(stub, name))
return stub
async def test_one_call_produces_one_stored_image(db, user_id, configured, fake):
from lembas.db.models import Attachment
outcome = await image_tool.run(
_context(db, user_id, configured), {"prompt": "a bicycle", "model": "sd.safetensors"}
)
assert outcome.event["status"] == "ok"
attachment = db.get(Attachment, outcome.event["attachment_id"])
assert attachment is not None
assert attachment.media_type == "image/png", "kept as ComfyUI made it, not transcoded"
assert attachment.message_id is None, "the loop binds it, not the runner"
assert attachment.source_label == "Image generation"
assert len(fake.submits) == 1
async def test_the_model_is_told_it_is_already_on_screen(db, user_id, configured, fake):
"""Without this the commonest next thing a model does is offer to show you
the image, which it has no way of doing and which has already happened."""
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "a bicycle"})
assert "already" in outcome.content or "is shown to them" in outcome.content
async def test_a_prompt_is_required(db, user_id, configured, fake):
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": " "})
assert outcome.event["status"] == "error"
assert not fake.submits
async def test_an_unknown_checkpoint_falls_back_rather_than_failing(db, user_id, configured, fake):
"""It would reach ComfyUI, be refused, and cost a round to discover -- and
the model was shown the list it may choose from."""
await image_tool.run(
_context(db, user_id, configured), {"prompt": "x", "model": "invented.safetensors"}
)
assert fake.submits[0]["4"]["inputs"]["ckpt_name"] == "sd.safetensors"
async def test_the_chat_s_own_preference_is_used_when_the_model_names_none(
db, user_id, configured, fake
):
configured.image_checkpoint = "other.safetensors"
db.commit()
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert fake.submits[0]["4"]["inputs"]["ckpt_name"] == "other.safetensors"
async def test_a_failure_out_there_is_reported_not_raised(
db, user_id, configured, fake, monkeypatch
):
async def refuse(config, wf):
raise comfy.ComfyError("ComfyUI refused the workflow — node 4")
monkeypatch.setattr(comfy, "submit", refuse)
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert outcome.event["status"] == "error"
assert "node 4" in outcome.event["error"]
# --- Reviewing -----------------------------------------------------------------
def _with_review(db, chat, *, tries=3, vision=True):
db.query(Model).filter(Model.model_id == "m").update(
{"capabilities_json": {"tools": True, "tool_image": True, "vision": vision}}
)
settings_store.update(
db, {"review_enabled": True, "max_tries": tries}, key=settings_store.IMAGES
)
db.commit()
async def test_a_kept_image_stops_the_loop(db, user_id, configured, fake, monkeypatch):
_with_review(db, configured)
monkeypatch.setattr(image_tool, "_review", _verdicts([(True, "")]))
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert len(fake.submits) == 1
async def test_a_rejected_image_is_tried_again(db, user_id, configured, fake, monkeypatch):
_with_review(db, configured)
monkeypatch.setattr(image_tool, "_review", _verdicts([(False, "no bicycle"), (True, "")]))
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert len(fake.submits) == 2
assert "no bicycle" in outcome.event["text"]
assert "2 attempts" in outcome.content
async def test_the_ceiling_is_enforced_and_the_last_one_is_kept(
db, user_id, configured, fake, monkeypatch
):
"""A request always produces a picture. Returning nothing after four
rejections would be spending a minute of somebody's GPU to say no."""
_with_review(db, configured, tries=3)
monkeypatch.setattr(image_tool, "_review", _verdicts([(False, "nope")] * 5))
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert len(fake.submits) == 3
assert outcome.event["status"] == "ok"
assert outcome.event["attachment_id"]
async def test_each_attempt_gets_a_new_seed(db, user_id, configured, fake, monkeypatch):
"""Retrying with the same seed regenerates the same rejected image."""
_with_review(db, configured, tries=3)
monkeypatch.setattr(
image_tool, "_review", _verdicts([(False, "no"), (False, "no"), (True, "")])
)
await image_tool.run(_context(db, user_id, configured), {"prompt": "x", "seed": 99})
seeds = [wf["3"]["inputs"]["seed"] for wf in fake.submits]
assert seeds[0] == 99, "the seed that was asked for is honoured"
assert len(set(seeds)) == 3, "and the retries are not the same picture again"
async def test_nothing_is_reviewed_without_a_vision_model(db, user_id, configured, fake):
""" "You asked for a picture and got an error about vision" is a worse answer
than a picture."""
_with_review(db, configured, vision=False)
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert len(fake.submits) == 1
assert outcome.event["status"] == "ok"
def _verdicts(sequence):
calls = {"n": 0}
async def review(context, endpoint, model_id, prompt, payload):
index = min(calls["n"], len(sequence) - 1)
calls["n"] += 1
return sequence[index]
return review
# --- Preserve VRAM -------------------------------------------------------------
async def test_nothing_is_unloaded_by_default(db, user_id, configured, fake, monkeypatch):
calls = {"n": 0}
async def unload(context):
calls["n"] += 1
return True
monkeypatch.setattr(image_tool, "_unload_llm", unload)
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert calls["n"] == 0
assert fake.frees == 0
async def test_preserve_vram_unloads_and_frees(db, user_id, configured, fake, monkeypatch):
settings_store.update(db, {"preserve_vram": True}, key=settings_store.IMAGES)
db.commit()
calls = {"n": 0}
async def unload(context):
calls["n"] += 1
return True
monkeypatch.setattr(image_tool, "_unload_llm", unload)
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert calls["n"] == 1, "once, before generating"
assert fake.frees >= 1, "and ComfyUI is asked to let go afterwards"
async def test_a_connection_with_no_unload_url_is_never_called(db, user_id, configured):
"""The whole of the per-connection design: a chat talking to a box on the
network must not have that box unloaded, because its memory is not the
memory ComfyUI wants."""
context = _context(db, user_id, configured)
assert await image_tool._unload_llm(context) is False
async def test_an_unload_that_fails_does_not_stop_the_generation(
db, user_id, configured, fake, monkeypatch
):
"""It is a hint before a slow operation. A machine that will not answer it is
one where the picture should still be drawn."""
connection = db.scalar(Connection.__table__.select())
db.query(Connection).update({"unload_url": "http://127.0.0.1:9/unload"})
settings_store.update(db, {"preserve_vram": True}, key=settings_store.IMAGES)
db.commit()
assert connection is not None
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
assert outcome.event["status"] == "ok"
tests/test_images_workflow.py
+122
View File
@@ -0,0 +1,122 @@
"""Filling a ComfyUI template.
The thing worth pinning here is types. ComfyUI validates its inputs, so a
workflow whose `steps` arrives as the string "20" is refused -- and the refusal
happens on somebody's machine a minute later rather than in a test.
"""
from __future__ import annotations
import json
from pathlib import Path
import pytest
import lembas
from lembas.services.images import workflow as wf
BASE = json.loads(
(Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text(
encoding="utf-8"
)
)
def _filled(**given):
return wf.fill(BASE, wf.resolve({"prompt": "a bicycle", **given}))
# --- Substitution --------------------------------------------------------------
def test_a_whole_placeholder_keeps_its_type():
"""`"steps": "{{steps}}"` has to become the number 20, not the text "20".
Substituting textually is the obvious implementation and produces a document
ComfyUI refuses."""
sampler = _filled(steps=25, cfg=7.5, seed=42)["3"]["inputs"]
assert sampler["steps"] == 25
assert isinstance(sampler["steps"], int)
assert sampler["cfg"] == 7.5
assert isinstance(sampler["cfg"], float)
assert sampler["seed"] == 42
assert isinstance(sampler["seed"], int)
def test_a_placeholder_inside_a_string_is_text():
"""Which is what makes `"{{prompt}}, masterpiece"` a usable template."""
filled = wf.fill(
{"n": {"inputs": {"text": "{{prompt}}, masterpiece"}}}, wf.resolve({"prompt": "a cat"})
)
assert filled["n"]["inputs"]["text"] == "a cat, masterpiece"
def test_the_prompt_and_the_negative_go_to_different_nodes():
filled = _filled(negative="blurry")
assert filled["6"]["inputs"]["text"] == "a bicycle"
assert filled["7"]["inputs"]["text"] == "blurry"
def test_the_links_between_nodes_survive():
"""A workflow is a graph, and `["4", 0]` is an edge rather than a value. A
filler that walked only dicts would flatten every one of them."""
filled = _filled()
assert filled["3"]["inputs"]["model"] == ["4", 0]
assert filled["8"]["inputs"]["vae"] == ["4", 2]
assert filled["5"]["inputs"]["batch_size"] == 1
def test_an_unknown_placeholder_is_left_alone():
"""`prompts.substitute`'s rule. A literal `{{x}}` is not a feature, but
silently deleting one is worse than leaving it where somebody can see it."""
filled = wf.fill({"n": {"inputs": {"lora": "{{lora_name}}"}}}, wf.resolve({"prompt": "x"}))
assert filled["n"]["inputs"]["lora"] == "{{lora_name}}"
def test_the_shipped_template_uses_every_placeholder_it_should():
assert wf.placeholders_in(BASE) == set(wf.PLACEHOLDERS)
# --- Resolution ----------------------------------------------------------------
def test_what_is_not_asked_for_takes_its_default():
values = wf.resolve({"prompt": "a bicycle"})
assert values["steps"] == 20
assert values["cfg"] == 8.0
assert values["width"] == 512
assert values["negative"] == "text, watermark"
assert values["sampler"] == "euler"
def test_the_seed_is_random_when_nobody_names_one():
"""A fixed default would make every unspecified generation identical -- and
would make the retry loop produce the same rejected image four times."""
seeds = {wf.resolve({"prompt": "x"})["seed"] for _ in range(8)}
assert len(seeds) == 8
def test_a_named_seed_is_kept():
assert wf.resolve({"prompt": "x", "seed": 1234})["seed"] == 1234
@pytest.mark.parametrize("empty", [None, ""])
def test_null_and_empty_mean_no_opinion(empty):
"""A model that emits `"steps": null` rather than omitting the key is common
enough that reading it as a request for zero steps would be a bug nobody
could see."""
assert wf.resolve({"prompt": "x", "steps": empty})["steps"] == 20
def test_out_of_range_numbers_are_clamped_rather_than_refused():
"""A model asking for 300 steps has misjudged rather than misbehaved, and a
clarifying round costs more than doing the sensible thing."""
values = wf.resolve({"prompt": "x", "steps": 5000, "cfg": -4, "denoise": 12})
assert values["steps"] == 150
assert values["cfg"] == 0.0
assert values["denoise"] == 1.0
def test_nonsense_falls_back_instead_of_raising():
"""Arguments come from a model, so "twenty" is a thing that will arrive."""
assert wf.resolve({"prompt": "x", "steps": "twenty"})["steps"] == 20
def test_a_seed_larger_than_comfyui_allows_is_wrapped():
assert 0 <= wf.resolve({"prompt": "x", "seed": wf.MAX_SEED + 5})["seed"] <= wf.MAX_SEED