DemonGatanjieuandClaude Opus 5.5
dd9c3ee2f5
- tags_from_picture: WD14 tagger (the source's model and thresholds) on a picture the user handed; run_status now carries what nodes showed as words (texts: a tagger's tags), with a summary and next step for a run that made words but no pictures. - Optional nodes: a block node marked optional goes when its pack is missing, and a link "match|decode" takes the first one there. The four upscales end in KJNodes' ColorMatch (mkl, the source's) when it is installed; the panel says "better with" it otherwise. - LanPaint's next step keeps the prompt about the whole picture; Follow a pose says OpenPose misreads anime (use control_depth), as the source's author did. - The directory has a size cap: past it, block names alone, so collected ways never make begin_task long; shorter headings and few-word lines. Live-tested on the isolated 8199 ComfyUI: LanPaint (area changed 14.6, outside 0.0, 82 s), tiled 2x (167 s), depth ControlNet (follows the layout), OpenPose (misreads an anime picture), background removal (8 s), tags (166 s with the model download), upscale with ColorMatch (73 s). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
249 lines
15 KiB
Python
249 lines
15 KiB
Python
"""A run an AI started (api/ai_task_runs.py): where it stands, read from a stand-in ComfyUI queue
|
|
and progress registry, on its step and for the AI's run_status. Temporary folders only."""
|
|
import sys
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1]), str(PLUGIN / "tests")]
|
|
|
|
from api import ai_gates, ai_task_runs, ai_tasks # noqa: E402
|
|
from test_ai_tasks import isolate, ok # noqa: E402
|
|
|
|
|
|
class State(Enum):
|
|
Pending = "pending"
|
|
Running = "running"
|
|
Finished = "finished"
|
|
|
|
|
|
class Queue:
|
|
def __init__(self, running=(), pending=()):
|
|
self.running, self.pending = list(running), list(pending)
|
|
|
|
def get_current_queue(self):
|
|
return self.running, self.pending
|
|
|
|
|
|
def registry(prompt_id, nodes):
|
|
prompt = {"3": {"class_type": "KSampler", "_meta": {"title": "采样器"}}, "4": {"class_type": "CheckpointLoaderSimple"},
|
|
"8": {"class_type": "VAEDecode"}, "9": {"class_type": "SaveImage"}}
|
|
return SimpleNamespace(prompt_id=prompt_id, nodes=nodes,
|
|
dynprompt=SimpleNamespace(original_prompt=prompt, get_node=lambda node_id: prompt[node_id]))
|
|
|
|
|
|
class ProgressTests(unittest.TestCase):
|
|
def setUp(self):
|
|
isolate()
|
|
|
|
def test_queued_then_running_with_the_node_and_its_steps(self):
|
|
queue = Queue(running=[(5, "other")], pending=[(7, "mine"), (6, "first")])
|
|
with mock.patch.object(ai_task_runs, "_queue", return_value=queue):
|
|
self.assertEqual(ai_task_runs.progress_of(["mine"]), {"state": "queued", "ahead": 2})
|
|
queue.running, queue.pending = [(7, "mine")], []
|
|
nodes = {"4": {"state": State.Finished, "value": 1, "max": 1}, "3": {"state": State.Running, "value": 14, "max": 28}}
|
|
with mock.patch.object(ai_task_runs, "_registry", return_value=registry("mine", nodes)):
|
|
got = ai_task_runs.progress_of(["mine"])
|
|
self.assertEqual(got, {"state": "running", "prompt": 1, "prompts": 1, "nodes_done": 1, "nodes": 4,
|
|
"node": "采样器", "value": 14, "max": 28, "percent": 50})
|
|
with mock.patch.object(ai_task_runs, "_queue", side_effect=RuntimeError("no server")):
|
|
self.assertIsNone(ai_task_runs.progress_of(["mine"]))
|
|
|
|
def test_progress_reaches_the_step_and_run_status_waits_for_the_end(self):
|
|
ai_tasks.record("Claude", "run_workflow", {"batch": 1}, ok({"queued": 1}), time.time())
|
|
task_id = ai_tasks.list_tasks()[0]["id"]
|
|
step_id = ai_tasks.get_task(task_id)["steps"][0]["id"]
|
|
with mock.patch.object(ai_tasks, "_notify") as told:
|
|
ai_task_runs._show_progress(task_id, step_id, {"state": "running", "value": 3, "max": 20})
|
|
ai_task_runs._show_progress(task_id, step_id, {"state": "running", "value": 3, "max": 20})
|
|
self.assertEqual(told.call_count, 1, "the page hears only of changes")
|
|
self.assertEqual(told.call_args[0][1]["progress"]["value"], 3)
|
|
status = ai_task_runs.status("Claude")
|
|
self.assertEqual((status["state"], status["progress"]["max"]), ("running", 20))
|
|
self.assertIn("end your turn", status["next"], "by default the user tells the AI when it is done")
|
|
with mock.patch.object(ai_gates, "run_wait", return_value="wait"):
|
|
self.assertIn("run_status(wait_seconds=55)", ai_task_runs.status("Claude")["next"])
|
|
history = {"p1": {"status": {"completed": True}, "outputs": {
|
|
"9": {"images": [{"filename": "a.png", "subfolder": "x", "type": "output"}]},
|
|
"12": {"images": [{"filename": "pose_00001_.png", "type": "temp"}]}, # a PreviewImage: the AI's to look at
|
|
"14": {"tags": ["1girl, red hair, smile"]}}}} # a tagger's words
|
|
threading.Timer(0.3, ai_task_runs._finish_run, (task_id, step_id, ["p1"], history)).start()
|
|
status = ai_task_runs.status("Claude", wait_seconds=5)
|
|
self.assertEqual((status["state"], status["pictures"], status["previews"]), ("done", ["x/a.png"], ["pose_00001_.png"]))
|
|
self.assertIn("folder='temp'", status["previews_note"])
|
|
self.assertEqual(status["texts"], [{"node": "14", "text": "1girl, red hair, smile"}], "what nodes showed as words")
|
|
self.assertEqual([item["filename"] for item in ai_tasks.get_task(task_id)["steps"][0]["results"]], ["a.png"], "previews are no results")
|
|
self.assertNotIn("progress", ai_tasks.get_task(task_id)["steps"][0], "a finished run has no progress left")
|
|
|
|
def test_a_run_saves_only_when_something_changes(self):
|
|
ai_tasks.record("A", "run_workflow", {}, ok({}), time.time())
|
|
task_id = ai_tasks.list_tasks()[0]["id"]
|
|
step_id = ai_tasks.get_task(task_id)["steps"][0]["id"]
|
|
with mock.patch.object(ai_tasks, "_save") as saved:
|
|
ai_task_runs._finish_run(task_id, step_id, ["p1"], {})
|
|
ai_task_runs._finish_run(task_id, step_id, ["p1"], {})
|
|
self.assertEqual(saved.call_count, 1)
|
|
|
|
def test_no_run_yet(self):
|
|
ai_tasks.record("A", "describe_canvas", {}, ok({}), time.time(), True)
|
|
self.assertEqual(ai_task_runs.status("A")["state"], "none") # an answer: the task shows no failed step
|
|
|
|
def run_step(self):
|
|
step = ai_tasks.record("A", "run_workflow", {}, ok({"queued": 1}), time.time())
|
|
return ai_tasks.active_task_id("A"), step["id"]
|
|
|
|
def test_no_queued_prompt_is_not_success(self):
|
|
task, step = self.run_step()
|
|
ai_task_runs._finish_run(task, step, [], {})
|
|
self.assertEqual(ai_task_runs.status("A")["state"], "stopped")
|
|
self.assertFalse(ai_tasks.get_task(task)["steps"][-1]["ok"])
|
|
|
|
def test_failure_is_recorded_even_when_other_prompts_are_pending(self):
|
|
task, step = self.run_step()
|
|
ai_task_runs._finish_run(task, step, ["one", "two"], {})
|
|
ai_task_runs._finish_run(task, step, ["one", "two"], {"one": {"status": {"status_str": "error"}}})
|
|
self.assertFalse(ai_tasks.get_task(task)["steps"][-1]["ok"])
|
|
self.assertIn("error", ai_task_runs.status("A"))
|
|
|
|
def test_watcher_timeout_and_history_failure_stop_the_receipt(self):
|
|
# Run the actual watcher synchronously; replace only time and host queue boundaries.
|
|
def inline_thread(target, **kwargs):
|
|
return SimpleNamespace(start=target)
|
|
for fail_history in (False, True):
|
|
isolate()
|
|
task, step = self.run_step()
|
|
queue = mock.Mock()
|
|
queue.get_history.side_effect = OSError("unavailable")
|
|
with mock.patch.object(ai_task_runs.threading, "Thread", side_effect=inline_thread), \
|
|
mock.patch.object(ai_task_runs, "_queue", return_value=queue), \
|
|
mock.patch.object(ai_tasks, "RUN_WATCH_SECONDS", 60 if fail_history else 0):
|
|
ai_task_runs.watch(step, {"prompt_ids": ["mine"]})
|
|
status = ai_task_runs.status("A")
|
|
self.assertEqual(status["state"], "stopped")
|
|
self.assertIn("history" if fail_history else "timed out", status["error"])
|
|
|
|
def test_batch_progress_keeps_total_and_current_index(self):
|
|
task, step = self.run_step()
|
|
queue = Queue(running=[(2, "two")])
|
|
with mock.patch.object(ai_task_runs, "_queue", return_value=queue), mock.patch.object(ai_task_runs, "_registry", return_value=None):
|
|
self.assertEqual(ai_task_runs.progress_of(["one", "two"]), {"state": "running", "prompt": 2, "prompts": 2})
|
|
|
|
def test_two_tasks_collect_only_their_receipts_and_keep_batch_order(self):
|
|
def record(filename):
|
|
return {"status": {"completed": True}, "outputs": {"1": {"images": [{"filename": filename, "type": "output"}]}}}
|
|
history = {pid: record(pid + ".png") for pid in ("z-first", "a-second", "other-ai", "manual")}
|
|
queue = mock.Mock()
|
|
queue.get_history.side_effect = lambda prompt_id: {prompt_id: history[prompt_id]}
|
|
with mock.patch.object(ai_task_runs.threading, "Thread", side_effect=lambda target, **kw: SimpleNamespace(start=target)), \
|
|
mock.patch.object(ai_task_runs, "_queue", return_value=queue):
|
|
task, step = self.run_step()
|
|
ai_task_runs.watch(step, {"prompt_ids": ["z-first", "a-second"]})
|
|
second = ai_tasks.begin_task("B", "other")["id"]
|
|
other = ai_tasks.record("B", "run_workflow", {}, ok({}), time.time())
|
|
ai_task_runs.watch(other["id"], {"prompt_ids": ["other-ai"], "submission_error": "Only 1 of 2 prompts accepted"})
|
|
self.assertEqual(ai_task_runs.status("A")["pictures"], ["z-first.png", "a-second.png"])
|
|
self.assertEqual(ai_task_runs.status("B")["pictures"], ["other-ai.png"])
|
|
self.assertEqual(ai_task_runs.status("B")["state"], "error")
|
|
self.assertNotIn(mock.call(prompt_id="manual"), queue.get_history.call_args_list)
|
|
|
|
|
|
class NewsTests(unittest.TestCase):
|
|
"""The AI hears where its run stands in its other answers, and once that it ended."""
|
|
|
|
def setUp(self):
|
|
isolate()
|
|
self.step = ai_tasks.record("A", "run_workflow", {}, ok({"queued": 1}), time.time())
|
|
self.task = ai_tasks.active_task_id("A")
|
|
with ai_tasks._lock:
|
|
task = ai_tasks._load()[self.task]
|
|
task["steps"][-1]["prompt_ids"] = ["p1"]
|
|
ai_tasks._save(task)
|
|
|
|
def told(self):
|
|
return ai_task_runs.tell({"content": [{"type": "text", "text": "{}"}], "structuredContent": {"n": 1}}, self.task)
|
|
|
|
def test_while_it_goes_every_answer_says_where_it_stands(self):
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._show_progress(self.task, self.step["id"], {"state": "running", "prompt": 1, "prompts": 1, "node": "KSampler",
|
|
"value": 14, "max": 28, "percent": 50, "nodes_done": 3, "nodes": 9})
|
|
for _ in range(2):
|
|
result = self.told()
|
|
run = result["structuredContent"]["your_run"]
|
|
self.assertEqual((run["state"], run["progress"]["percent"]), ("running", 50))
|
|
self.assertIn("KSampler step 14/28, 50%, 3 of 9 nodes done", run["summary"])
|
|
self.assertIn("your_run", result["content"][0]["text"], "the text says it too")
|
|
self.assertEqual(result["structuredContent"]["n"], 1, "the tool's own answer stays")
|
|
|
|
def test_the_end_is_told_once(self):
|
|
history = {"p1": {"status": {"completed": True}, "outputs": {"9": {"images": [{"filename": "a.png", "type": "output"}]}}}}
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], history)
|
|
self.assertIn("finished", self.told()["structuredContent"]["your_run"]["summary"])
|
|
self.assertNotIn("your_run", self.told()["structuredContent"])
|
|
|
|
def test_run_status_telling_the_end_is_enough(self):
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}}})
|
|
self.assertIn("finished", ai_task_runs.status("A")["summary"])
|
|
self.assertNotIn("your_run", self.told()["structuredContent"])
|
|
|
|
def test_the_pictures_come_once_with_the_answer_that_tells_the_end(self):
|
|
history = {"p1": {"status": {"completed": True}, "outputs": {"9": {"images": [
|
|
{"filename": f"{n}.png", "subfolder": "x", "type": "output"} for n in range(6)]}}}}
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], history)
|
|
picture = {"type": "image", "data": "AAAA", "mimeType": "image/webp"}
|
|
with mock.patch.object(ai_task_runs, "_resolve_image", side_effect=lambda rel: (rel, "/out/" + rel)), mock.patch.object(ai_task_runs, "_image_content", return_value=picture):
|
|
before = ai_task_runs.show_pictures({"content": [{"type": "text", "text": "{}"}], "structuredContent": {}}, self.task)
|
|
self.assertEqual(len(before["content"]), 1, "nothing before the AI is told it ended")
|
|
result = ai_task_runs.show_pictures(self.told(), self.task)
|
|
self.assertEqual(result["structuredContent"]["pictures_attached"], ["x/0.png", "x/1.png", "x/2.png", "x/3.png"])
|
|
self.assertEqual([item["type"] for item in result["content"]], ["text"] + ["image"] * ai_task_runs.MAX_SHOWN)
|
|
self.assertIn("pictures_attached", result["content"][0]["text"])
|
|
again = ai_task_runs.show_pictures({"content": [{"type": "text", "text": "{}"}], "structuredContent": {}}, self.task)
|
|
self.assertEqual(len(again["content"]), 1, "once per run")
|
|
|
|
def test_a_line_to_paste_until_the_ai_knows(self):
|
|
with mock.patch.object(ai_tasks, "_notify"), mock.patch.object(ai_gates, "run_wait", return_value="call"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}}})
|
|
self.assertTrue(ai_tasks.get_task(self.task)["steps"][-1]["call_ai"], "the workbench offers the line")
|
|
self.told()
|
|
self.assertNotIn("call_ai", ai_tasks.get_task(self.task)["steps"][-1], "gone once the AI is told")
|
|
|
|
def test_no_line_when_the_ai_waits_or_a_new_run_starts(self):
|
|
with mock.patch.object(ai_tasks, "_notify"), mock.patch.object(ai_gates, "run_wait", return_value="wait"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}}})
|
|
self.assertNotIn("call_ai", ai_tasks.get_task(self.task)["steps"][-1])
|
|
with ai_tasks._lock:
|
|
task = ai_tasks._load()[self.task]
|
|
task["steps"][-1]["call_ai"] = True
|
|
ai_tasks._save(task)
|
|
newer = ai_tasks.record("A", "run_workflow", {}, ok({"queued": 1}), time.time())
|
|
with mock.patch("threading.Thread"):
|
|
ai_task_runs.watch(newer["id"], {"prompt_ids": ["p2"]})
|
|
self.assertNotIn("call_ai", ai_tasks.get_task(self.task)["steps"][0], "an older run offers no line once another starts")
|
|
|
|
def test_an_old_run_told_before_brings_no_pictures(self):
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}, "outputs": {
|
|
"9": {"images": [{"filename": "a.png", "type": "output"}]}}}})
|
|
with ai_tasks._lock:
|
|
task = ai_tasks._load()[self.task]
|
|
task["steps"][-1]["told"] = True # told before pictures came with it: no "shown" mark
|
|
ai_tasks._save(task)
|
|
result = ai_task_runs.show_pictures({"content": [{"type": "text", "text": "{}"}], "structuredContent": {}}, self.task)
|
|
self.assertEqual(len(result["content"]), 1)
|
|
|
|
def test_no_run_no_news(self):
|
|
ai_tasks.begin_task("A", "other", new_task=True)
|
|
self.assertIsNone(ai_task_runs.news(ai_tasks.active_task_id("A")))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|