Jaroslav Beneš
ca7eb6cedb
Omitting the seed was already random. Passing -1 was not: it went through the uint64 wrap and arrived as 18446744073709551615, which is a perfectly valid *fixed* seed -- so "give me something new" returned the identical picture every time, silently, and the retry loop would have redrawn the same rejected image until it ran out of attempts. -1 is what ComfyUI's own interface uses for random, and A1111, and everything else that has ever asked somebody for a seed. A model that has read any of them will write it, so the one reading that had to work was the one that did not. Any negative value, not only -1, because the sentinel is the *idea* rather than the number and a model that writes -2 means the same thing. Zero stays a real seed: it is the boundary this change could easily have swallowed, and it is one somebody deliberately picks. Confirmed against the real ComfyUI: -1 now sends a random uint64 that it accepts and draws from. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
530 lines
21 KiB
Python
530 lines
21 KiB
Python
"""The tool that makes a picture, and the loop that decides to keep it.
|
|
|
|
One call is one finished image. The alternative -- return every attempt to the
|
|
conversation and let the model decide whether to call again -- costs a full
|
|
round per retry, makes the ceiling advisory rather than enforced, and shows the
|
|
reader every reject on the way past. So the retrying happens here, and what
|
|
comes back is the image that was kept.
|
|
|
|
**Three things are ordered rather than incidental.**
|
|
|
|
*The reviewer is asked about bytes, not about a row.* An attempt that is going
|
|
to be thrown away should not leave an `Attachment` behind, so the judge is shown
|
|
a downscaled preview built in memory and only the kept image is ever written.
|
|
|
|
*Preserve VRAM swaps around the review, not around the tool.* The sequence is
|
|
unload the LLM, generate, free ComfyUI, ask the reviewer (which loads the LLM
|
|
again), and round once more if it said no. Two model loads per retry, which is
|
|
why the two settings are independent and the admin page says so.
|
|
|
|
*Nothing loads the LLM back at the end.* The reply's next request does it, and
|
|
llama-swap -- or Ollama, or anything else worth pointing this at -- loads on
|
|
demand. A step that exists in the description and not in the code looks like an
|
|
omission, so it is said here instead.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from sqlalchemy import select
|
|
|
|
from lembas.services.images import comfy, workflow
|
|
from lembas.services.llm.openai_client import Endpoint, LLMError, complete
|
|
from lembas.services.tools import RISK_WRITE, ToolContext, ToolDef, ToolOutcome
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# What the reviewer is allowed to write back. It is one verdict and one line of
|
|
# reason, and a model that writes an essay about a picture is a model whose
|
|
# answer nobody reads.
|
|
MAX_VERDICT_TOKENS = 200
|
|
|
|
# How long to wait for a connection to admit it has unloaded. Short: this is a
|
|
# hint before a slow operation, and a machine that will not answer it is one
|
|
# where the generation should go ahead anyway rather than fail.
|
|
UNLOAD_TIMEOUT = 30.0
|
|
|
|
SCHEMA: dict[str, Any] = {
|
|
"type": "object",
|
|
"properties": {
|
|
# First, and the only required one, because `tools.parse_arguments`
|
|
# puts the whole raw string into the first required parameter when a
|
|
# model emits arguments that are not valid JSON. That failure is common
|
|
# with small models, and this way it degrades into a prompt rather than
|
|
# into a seed.
|
|
"prompt": {
|
|
"type": "string",
|
|
"description": "What to draw. Describe the subject, the setting and the style.",
|
|
},
|
|
"negative": {
|
|
"type": "string",
|
|
"description": "What to keep out of the picture. Defaults to 'text, watermark'.",
|
|
},
|
|
"template": {
|
|
"type": "string",
|
|
"description": "Which workflow to use. Omit to use this chat's usual one.",
|
|
},
|
|
"model": {
|
|
"type": "string",
|
|
"description": "Which checkpoint to draw with. Omit to use this chat's usual one.",
|
|
},
|
|
"seed": {
|
|
"type": "integer",
|
|
"description": (
|
|
"Omit it, or pass -1, for a new random image. Repeat a seed you were "
|
|
"told about to get the same image again."
|
|
),
|
|
},
|
|
"steps": {"type": "integer", "description": "Sampling steps. Default 20."},
|
|
"cfg": {"type": "number", "description": "Prompt adherence. Default 8."},
|
|
"width": {"type": "integer", "description": "Pixels. Default 512."},
|
|
"height": {"type": "integer", "description": "Pixels. Default 512."},
|
|
"sampler": {"type": "string", "description": "Sampler name. Default euler."},
|
|
"scheduler": {"type": "string", "description": "Scheduler name. Default normal."},
|
|
"denoise": {"type": "number", "description": "0 to 1. Default 1."},
|
|
},
|
|
"required": ["prompt"],
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Attempt:
|
|
"""One generated image and what was decided about it."""
|
|
|
|
number: int
|
|
seed: int
|
|
kept: bool
|
|
verdict: str = ""
|
|
|
|
|
|
def config_of(context: ToolContext) -> comfy.Config:
|
|
"""The client snapshot, with the key decrypted at the last moment."""
|
|
from lembas.services.crypto import decrypt
|
|
|
|
values = context.image_config or {}
|
|
return comfy.Config(
|
|
base_url=str(values.get("base_url") or ""),
|
|
api_key=decrypt(str(values.get("api_key_encrypted") or "")),
|
|
timeout=float(values.get("timeout") or comfy.DEFAULT_TIMEOUT),
|
|
)
|
|
|
|
|
|
def _choices(db, values: dict[str, Any]) -> tuple[list[Any], list[str]]:
|
|
"""The templates and checkpoints on offer, for the schema and the harness."""
|
|
from lembas.db.models import ImageWorkflow
|
|
|
|
rows = list(
|
|
db.scalars(
|
|
select(ImageWorkflow)
|
|
.where(ImageWorkflow.enabled.is_(True))
|
|
.order_by(ImageWorkflow.position, ImageWorkflow.slug)
|
|
)
|
|
)
|
|
return rows, [str(name) for name in (values.get("checkpoints") or [])]
|
|
|
|
|
|
def schema_for(db, values: dict[str, Any]) -> dict[str, Any]:
|
|
"""The parameter schema, with this instance's own choices in it.
|
|
|
|
`template` and `model` become enums because a name that does not exist is a
|
|
refusal from ComfyUI and a wasted round; `sampler` and `scheduler` stay
|
|
plain strings because there are forty-four and nine of them, and an enum
|
|
that size costs tokens on every request forever to prevent a mistake worth
|
|
one sentence of correction.
|
|
"""
|
|
rows, checkpoints = _choices(db, values)
|
|
schema = json.loads(json.dumps(SCHEMA))
|
|
if rows:
|
|
schema["properties"]["template"]["enum"] = [row.slug for row in rows]
|
|
schema["properties"]["template"]["description"] = "Which workflow to use. " + "; ".join(
|
|
f"{row.slug}: {row.description or row.name}" for row in rows[:12]
|
|
)
|
|
if checkpoints:
|
|
schema["properties"]["model"]["enum"] = checkpoints
|
|
return schema
|
|
|
|
|
|
def tool_def(db, values: dict[str, Any]) -> ToolDef:
|
|
return ToolDef(
|
|
name="image_generate",
|
|
family="image",
|
|
description=(
|
|
"Draw a picture from a description and show it to the person you are "
|
|
"talking to. Returns once the image has been made and is on screen."
|
|
),
|
|
parameters=schema_for(db, values),
|
|
run=run,
|
|
# Not RISK_READ: it spends somebody's GPU for a minute and puts a new
|
|
# artefact in the conversation. In an agent chat that means the mode
|
|
# decides whether to ask first, which is the right answer for a call
|
|
# that cannot be undone by reading something again.
|
|
risk=RISK_WRITE,
|
|
)
|
|
|
|
|
|
# --- Preserve VRAM -------------------------------------------------------------
|
|
async def _unload_llm(context: ToolContext) -> bool:
|
|
"""Ask this chat's own endpoint to drop its model. Best-effort.
|
|
|
|
*This chat's own* is the whole of the design. The unload hook is a column on
|
|
`Connection`, so a chat talking to a local llama-swap unloads that and a
|
|
chat talking to a box on the network unloads nothing -- its VRAM is not the
|
|
VRAM ComfyUI is about to want.
|
|
"""
|
|
from lembas.db.models import Connection
|
|
from lembas.db.session import session_scope
|
|
|
|
url = ""
|
|
method = "POST"
|
|
try:
|
|
with session_scope() as db:
|
|
connection = db.get(Connection, context.connection_id)
|
|
if connection is not None:
|
|
url = (connection.unload_url or "").strip()
|
|
method = (connection.unload_method or "POST").upper()
|
|
except Exception: # noqa: BLE001 - a hint that could not be looked up is not a failure
|
|
log.debug("could not read the unload hook", exc_info=True)
|
|
return False
|
|
|
|
if not url:
|
|
return False
|
|
try:
|
|
async with httpx.AsyncClient(timeout=UNLOAD_TIMEOUT) as client:
|
|
await client.request(method, url)
|
|
return True
|
|
except Exception: # noqa: BLE001 - see the module docstring: a hint, not a step
|
|
log.info("could not unload the model at %s", url, exc_info=True)
|
|
return False
|
|
|
|
|
|
# --- The reviewer --------------------------------------------------------------
|
|
def _reviewer(context: ToolContext) -> tuple[Endpoint, str] | None:
|
|
"""The model that judges an image, or None if there is nobody to ask.
|
|
|
|
The admin's choice first, then the chat's own model when it has vision. A
|
|
chat on a text-only model with no reviewer configured simply keeps the first
|
|
image, which is the behaviour with review switched off -- said here rather
|
|
than failing, because "you asked for a picture and got an error about
|
|
vision" is a worse answer than a picture.
|
|
"""
|
|
from lembas.db.models import Connection, Model
|
|
from lembas.db.session import session_scope
|
|
|
|
values = context.image_config or {}
|
|
if not values.get("review_enabled"):
|
|
return None
|
|
|
|
wanted = str(values.get("review_model_id") or "")
|
|
try:
|
|
with session_scope() as db:
|
|
model = None
|
|
if wanted:
|
|
model = db.get(Model, wanted)
|
|
if model is None and context.model_id:
|
|
model = db.scalar(
|
|
select(Model).where(
|
|
Model.model_id == context.model_id,
|
|
Model.connection_id == context.connection_id,
|
|
)
|
|
)
|
|
if model is None or not (model.capabilities_json or {}).get("vision"):
|
|
return None
|
|
connection = db.get(Connection, model.connection_id)
|
|
if connection is None or not connection.enabled:
|
|
return None
|
|
return Endpoint.from_connection(connection), model.model_id
|
|
except Exception: # noqa: BLE001 - no reviewer is a degraded mode, not an error
|
|
log.warning("could not resolve an image reviewer", exc_info=True)
|
|
return None
|
|
|
|
|
|
async def _review(
|
|
context: ToolContext, endpoint: Endpoint, model_id: str, prompt: str, payload: bytes
|
|
) -> tuple[bool, str]:
|
|
"""Show the reviewer the image and ask whether to keep it.
|
|
|
|
Answers `(keep, reason)`. **Anything that goes wrong is a keep**: the
|
|
reviewer is a second opinion on a picture that already exists, and losing an
|
|
image because a judging request timed out would be the check destroying the
|
|
thing it was checking.
|
|
"""
|
|
from lembas.db.session import session_scope
|
|
from lembas.services import files as files_service
|
|
from lembas.services import prompts as prompts_service
|
|
|
|
preview = files_service.preview_data_uri(payload, max_edge=768)
|
|
if preview is None:
|
|
return True, ""
|
|
|
|
with session_scope() as db:
|
|
instruction = prompts_service.resolve(db, "task.image_review")
|
|
# An administrator who cleared the fragment has switched reviewing off, the
|
|
# same way clearing `task.compact` switches compaction off. Nothing is asked
|
|
# of anyone and the image is kept.
|
|
if not instruction.strip():
|
|
return True, ""
|
|
|
|
body = {
|
|
"model": model_id,
|
|
"messages": [
|
|
{"role": "system", "content": instruction},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": f"The request was: {prompt}"},
|
|
{"type": "image_url", "image_url": {"url": preview}},
|
|
],
|
|
},
|
|
],
|
|
"max_tokens": MAX_VERDICT_TOKENS,
|
|
"temperature": 0,
|
|
}
|
|
try:
|
|
answer = (await complete(endpoint, body)).strip()
|
|
except LLMError as exc:
|
|
log.info("could not review a generated image: %s", exc.message)
|
|
return True, ""
|
|
|
|
verdict, _, reason = answer.partition("\n")
|
|
keep = not verdict.strip().upper().startswith("RETRY")
|
|
return keep, (reason or verdict).strip()[:300]
|
|
|
|
|
|
# --- The runner ----------------------------------------------------------------
|
|
async def run(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
|
"""Generate one image, review it if there is anybody to ask, and keep one."""
|
|
from lembas.db.session import session_scope
|
|
from lembas.services import files as files_service
|
|
|
|
event: dict[str, Any] = {
|
|
"name": "image_generate",
|
|
"kind": "image",
|
|
"query": str(args.get("prompt") or "")[:200],
|
|
"results": [],
|
|
}
|
|
|
|
prompt = str(args.get("prompt") or "").strip()
|
|
if not prompt:
|
|
return ToolOutcome(
|
|
"No prompt was given, so nothing was drawn. Say what the picture should show.",
|
|
{**event, "status": "error", "error": "No prompt."},
|
|
)
|
|
if not context.chat_id:
|
|
return ToolOutcome(
|
|
"Images can only be generated inside a chat.",
|
|
{**event, "status": "error", "error": "No chat."},
|
|
)
|
|
|
|
values = context.image_config or {}
|
|
config = config_of(context)
|
|
if not config.configured:
|
|
return ToolOutcome(
|
|
"No image generator is configured on this instance.",
|
|
{**event, "status": "error", "error": "No ComfyUI configured."},
|
|
)
|
|
|
|
# Resolve the template and the checkpoint: what the model asked for, then
|
|
# this chat's usual, then the instance default. Every rung is a preference
|
|
# and none of them is a constraint, which is what lets a model that only
|
|
# wrote a prompt still get a picture.
|
|
try:
|
|
with session_scope() as db:
|
|
rows, checkpoints = _choices(db, values)
|
|
wanted = str(args.get("template") or "")
|
|
chosen = _pick(rows, wanted, context.image_workflow_id, values)
|
|
if chosen is None:
|
|
return ToolOutcome(
|
|
"No image workflow has been set up on this instance.",
|
|
{**event, "status": "error", "error": "No workflow."},
|
|
)
|
|
template = json.loads(json.dumps(chosen.workflow_json or {}))
|
|
template_slug, template_name = chosen.slug, chosen.name
|
|
except ToolOutcome: # pragma: no cover - defensive
|
|
raise
|
|
except Exception as exc: # noqa: BLE001
|
|
log.exception("could not resolve an image workflow")
|
|
return ToolOutcome(
|
|
f"The image workflow could not be read: {exc}",
|
|
{**event, "status": "error", "error": str(exc)},
|
|
)
|
|
|
|
checkpoint = _checkpoint(str(args.get("model") or ""), context.image_checkpoint, checkpoints)
|
|
if checkpoint is None:
|
|
return ToolOutcome(
|
|
"No checkpoint is available. An administrator has to list them on the "
|
|
"image generation page.",
|
|
{**event, "status": "error", "error": "No checkpoint."},
|
|
)
|
|
|
|
given = {name: args.get(name) for name in workflow.PLACEHOLDERS if name in args}
|
|
given["model"] = checkpoint
|
|
given["prompt"] = prompt
|
|
|
|
reviewer = _reviewer(context)
|
|
tries = int(values.get("max_tries") or 1) if reviewer else 1
|
|
preserve = bool(values.get("preserve_vram"))
|
|
|
|
attempts: list[Attempt] = []
|
|
kept: tuple[bytes, dict[str, Any]] | None = None
|
|
|
|
try:
|
|
for number in range(1, tries + 1):
|
|
if preserve:
|
|
await _unload_llm(context)
|
|
|
|
params = workflow.resolve({**given, "seed": args.get("seed") if number == 1 else None})
|
|
refs = await comfy.await_images(
|
|
config, await comfy.submit(config, workflow.fill(template, params))
|
|
)
|
|
if not refs:
|
|
raise comfy.ComfyError("ComfyUI finished but saved no image.")
|
|
payload = await comfy.fetch_image(config, refs[0])
|
|
|
|
if preserve:
|
|
await comfy.free(config)
|
|
|
|
if reviewer is None:
|
|
attempts.append(Attempt(number, params["seed"], kept=True))
|
|
kept = (payload, params)
|
|
break
|
|
|
|
endpoint, model_id = reviewer
|
|
keep, reason = await _review(context, endpoint, model_id, prompt, payload)
|
|
last = number == tries
|
|
attempts.append(Attempt(number, params["seed"], kept=keep or last, verdict=reason))
|
|
if keep or last:
|
|
kept = (payload, params)
|
|
break
|
|
except comfy.ComfyError as exc:
|
|
if preserve:
|
|
await comfy.free(config)
|
|
return ToolOutcome(
|
|
f"The image could not be generated: {exc.message}",
|
|
{**event, "status": "error", "error": exc.message},
|
|
)
|
|
|
|
if preserve:
|
|
await comfy.free(config)
|
|
if kept is None: # pragma: no cover - the loop always keeps its last attempt
|
|
return ToolOutcome(
|
|
"Nothing was generated.", {**event, "status": "error", "error": "No image."}
|
|
)
|
|
|
|
payload, params = kept
|
|
try:
|
|
with session_scope() as db:
|
|
attachment = files_service.store(
|
|
db,
|
|
user_id=context.owner_id,
|
|
chat_id=context.chat_id,
|
|
payload=payload,
|
|
filename=f"{template_slug}-{params['seed']}.png",
|
|
# What ComfyUI made, at the size it made it. See `_keep_image`.
|
|
keep_original=True,
|
|
source_label="Image generation",
|
|
source_path=f"{checkpoint} · seed {params['seed']}",
|
|
)
|
|
attachment_id = attachment.id
|
|
width, height = attachment.width, attachment.height
|
|
except Exception as exc: # noqa: BLE001
|
|
log.exception("could not store a generated image")
|
|
return ToolOutcome(
|
|
f"The image was generated but could not be saved: {exc}",
|
|
{**event, "status": "error", "error": str(exc)},
|
|
)
|
|
|
|
return ToolOutcome(
|
|
_describe(prompt, template_name, checkpoint, params, attempts),
|
|
{
|
|
**event,
|
|
"status": "ok",
|
|
"detail": f"{template_name} · {checkpoint}",
|
|
"text": _transcript(params, attempts),
|
|
# Bound to the reply by `generation._persist`, the single writer. A
|
|
# runner may create the row; only the loop may say which turn owns
|
|
# it.
|
|
"attachment_id": attachment_id,
|
|
"image": {"id": attachment_id, "width": width, "height": height},
|
|
},
|
|
)
|
|
|
|
|
|
def _pick(rows: list[Any], wanted: str, chat_default: str, values: dict[str, Any]) -> Any:
|
|
"""The workflow to use: asked for, then the chat's, then the instance's."""
|
|
by_slug = {row.slug: row for row in rows}
|
|
if wanted and wanted in by_slug:
|
|
return by_slug[wanted]
|
|
by_id = {row.id: row for row in rows}
|
|
if chat_default and chat_default in by_id:
|
|
return by_id[chat_default]
|
|
fallback = str(values.get("default_workflow_id") or "")
|
|
if fallback and fallback in by_id:
|
|
return by_id[fallback]
|
|
return rows[0] if rows else None
|
|
|
|
|
|
def _checkpoint(wanted: str, chat_default: str, available: list[str]) -> str | None:
|
|
"""The checkpoint to draw with, on the same ladder.
|
|
|
|
A name the instance does not have is ignored rather than passed through: it
|
|
would reach ComfyUI, be refused, and cost a round to discover -- and the
|
|
model was shown the list it may choose from.
|
|
"""
|
|
if wanted and wanted in available:
|
|
return wanted
|
|
if chat_default and chat_default in available:
|
|
return chat_default
|
|
return available[0] if available else None
|
|
|
|
|
|
def _describe(
|
|
prompt: str, template: str, checkpoint: str, params: dict[str, Any], attempts: list[Attempt]
|
|
) -> str:
|
|
"""What the model reads back.
|
|
|
|
It is told the image is already on screen, because otherwise the commonest
|
|
next thing it does is offer to show it -- and there is nothing it could do
|
|
to comply.
|
|
"""
|
|
lines = [
|
|
"The image has been generated and is shown to them. It is not a link and "
|
|
"needs no further action.",
|
|
f"Prompt: {prompt}",
|
|
f"Template {template}, checkpoint {checkpoint}, "
|
|
f"{params['width']}x{params['height']}, seed {params['seed']}, "
|
|
f"{params['steps']} steps, cfg {params['cfg']}.",
|
|
]
|
|
if len(attempts) > 1:
|
|
rejected = [a for a in attempts if not a.kept]
|
|
lines.append(
|
|
f"It took {len(attempts)} attempts; the earlier ones were rejected on review "
|
|
f"({'; '.join(a.verdict for a in rejected if a.verdict) or 'no reason given'})."
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _transcript(params: dict[str, Any], attempts: list[Attempt]) -> str:
|
|
"""What the reader sees when they open the tool block.
|
|
|
|
The rejected attempts are recorded here and their images are not kept. A
|
|
transcript full of pictures somebody's model decided against is noise, and
|
|
the disk they would occupy buys nothing -- what is worth knowing is that it
|
|
took three goes and why the first two did not do.
|
|
"""
|
|
lines = [
|
|
f"seed {params['seed']} · {params['steps']} steps · cfg {params['cfg']} · "
|
|
f"{params['sampler']}/{params['scheduler']} · denoise {params['denoise']}"
|
|
]
|
|
if len(attempts) > 1:
|
|
lines.append("")
|
|
for attempt in attempts:
|
|
state = "kept" if attempt.kept else "rejected"
|
|
reason = f" — {attempt.verdict}" if attempt.verdict else ""
|
|
lines.append(f"Attempt {attempt.number} (seed {attempt.seed}): {state}{reason}")
|
|
return "\n".join(lines)
|