From 0930d07d4fea7cfa29ade4a1d3394ecda18a9616 Mon Sep 17 00:00:00 2001 From: DemonGatanjieu <207458273+DemonGatanjieu@users.noreply.github.com> Date: Sat, 24 Oct 2026 01:21:27 +0800 Subject: [PATCH] Give the AI a closer look and a run's previews, so it need not script pictures In one hour Gemini ran 63 shell commands next to 23 AMB calls: checking files, cropping and comparing parts of pictures, running OpenPose itself to see a skeleton, and painting pixels with OpenCV (it overwrote an output picture six times). - api/mcp_look.py zoom_image: part of a picture in the output, input or temp folder shown larger, the same part of a second picture beside it; read only. - get_image_info reads the temp folder too; run_status lists a run's PreviewImage pictures (ai_task_runs.outputs_of(record, "temp"), kept as the step's previews, not its results). - begin_task and the repaint note: AMB's tools cover looking; change pictures only by running workflows, never by editing picture files with code of one's own. Co-Authored-By: Claude Opus 5.5 --- ARCHITECTURE.md | 3 ++ CHANGELOG.md | 2 + api/action_catalog.py | 1 + api/ai_panels.py | 3 +- api/ai_task_runs.py | 24 ++++++++--- api/mcp_look.py | 88 ++++++++++++++++++++++++++++++++++++++ api/mcp_server.py | 3 +- api/mcp_task_tools.py | 7 ++- api/mcp_tools.py | 30 +++++++------ docs/guides/mcp.md | 1 + tests/test_ai_task_runs.py | 8 +++- tests/test_mcp_look.py | 62 +++++++++++++++++++++++++++ 12 files changed, 208 insertions(+), 24 deletions(-) create mode 100644 api/mcp_look.py create mode 100644 tests/test_mcp_look.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 64d629d..49ba978 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -161,6 +161,9 @@ covered by the `styles.css` manifest. ComfyUI's websocket and waits for its answer; `web/modules/mcp_bridge.js` is the page side, doing it with Anomalous's canvas code (one Ctrl+Z step each, logged as the AI's; the outcome carries the log entry's id as `_activity`, for the AI workbench's undo). + `api/mcp_look.py` is zoom_image, the AI's closer look: part of a picture in the output, input + or temp folder shown larger, the same part of a second picture beside it, read only (so an AI + need not crop or compare pictures with code of its own). `api/mcp_model_use.py` offers search_models and get_model (mcp_tools.py finds the files) with what loads each model: the node types whose drop-down lists it (from `/object_info`), or, for a file whose header says it belongs in another folder (model_placement.py), where, diff --git a/CHANGELOG.md b/CHANGELOG.md index b8152b0..002c940 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,6 +34,8 @@ - A model copied to another drive is checked at its start, middle and end (a megabyte each) besides its size before the original goes; no fingerprint needed, no felt cost. - 挪动、导入模型时，就算目标位置在最后一刻冒出一个同名文件，也绝不会被盖掉：这次挪动会失败，两边的文件都原样保留（之前在 Windows 上理论上有极小的机会覆盖）。 - Moving or importing a model never replaces a file that took the name at the last moment: the move fails and both files stay (on Windows a rename could replace it before). +- **AI 能放大看图、看预览，不用再自己写代码**：新工具 `zoom_image` 把一张图的某一块放大给 AI 看，还能把两张图的同一块并排比较（改之前和改之后）；输出、输入、临时预览的图都能看，只看不改。出图结果里会列出这次运行的预览图（比如 OpenPose 骨骼、遮罩的预览），AI 用它们就能检查中间结果，不必自己去跑预处理器。AI 开任务、收到遮罩时都会被明确告知：要改图就跑 ComfyUI 的工作流，绝不用自己的代码去画、去改、去覆盖图片文件。 +- zoom_image shows part of a picture larger (output, input or temp folder), with the same part of a second picture to compare, reading only; run_status lists a run's PreviewImage pictures; begin_task and the repaint note tell the AI to change pictures only by running workflows, never by editing picture files with its own code. - **AI 能直接看输入文件夹里的图**：`get_image_info` 加了 `folder=input`，能看 ComfyUI 输入文件夹里的图（你标注时 AMB 存的重绘遮罩图就在那里）：尺寸、要重画的区域在哪、占多大，也能附上图本身。交给 AI 遮罩图时也明说“它已经在输入文件夹里了，不用跑命令检查”，Gemini 不会再为这个弹窗问你。 - get_image_info with folder=input reads a picture in ComfyUI's input folder (the inpaint picture Anomalous writes for a repaint): its size and the area LoadImage gives as MASK, optionally attached; the marks message says the inpaint picture is in place, so the AI need not check it with a command. - **AI 少去翻文件**：AMB 工具返回给 AI 的内容去掉了多余的空格、画布里 ComfyUI 自己的内部字段不再列出、操作记录默认只给最近 10 条，这样 Antigravity 不必把结果另存成文件再让 Gemini 去读。AI 开始任务时也会被告知：模型、图片、节点、画布和运行都有 AMB 的工具，不要自己去读 ComfyUI 的文件夹、源码或接口（那样每一步都要你点允许）。 diff --git a/api/action_catalog.py b/api/action_catalog.py index bd0b4d5..e048365 100644 --- a/api/action_catalog.py +++ b/api/action_catalog.py @@ -26,6 +26,7 @@ TOOLS = { "scan_report": Entry("look", "library", None), "search_images": Entry("look", "picture", None), "get_image_info": Entry("look", "picture", None), + "zoom_image": Entry("look", "picture", None), "list_combos": Entry("look", "combo", None), "list_workflows": Entry("look", "canvas", None), "list_saved_prompts": Entry("look", "prompt", None), diff --git a/api/ai_panels.py b/api/ai_panels.py index 6c91e96..3b2ac8f 100644 --- a/api/ai_panels.py +++ b/api/ai_panels.py @@ -40,7 +40,8 @@ INTERACTIVE = ("choices", "input", "picks") INPAINT_NOTE = ( "The user painted where to repaint. LoadImage(image=inpaint_image) gives the picture and, as MASK, that area. " "inpaint_image is already in ComfyUI's input folder, ready: do not check it with a command; " - "get_image_info(image=inpaint_image, folder='input', include_image=true) shows it and the area. " + "get_image_info(image=inpaint_image, folder='input', include_image=true) shows it and the area, zoom_image a part " + "of it larger. Repaint with a workflow; do not paint the picture with code of your own. " "Use the picture's own model and prompt (get_image_info on image), e.g. new_workflow: CheckpointLoaderSimple, " "LoadImage, CLIPTextEncode x2, VAEEncode(pixels, vae) -> SetLatentNoiseMask(samples, mask) -> KSampler " "(denoise 0.5-0.8 keeps the look; 1.0 with VAEEncodeForInpaint(grow_mask_by 8) replaces it) -> VAEDecode -> " diff --git a/api/ai_task_runs.py b/api/ai_task_runs.py index 4e34875..a872f9b 100644 --- a/api/ai_task_runs.py +++ b/api/ai_task_runs.py @@ -20,6 +20,7 @@ from .mcp_tools import tool_text RUN_POLL_SECONDS = 2.0 MAX_STATUS_WAIT = 55 +MAX_PREVIEWS = 8 # the last PreviewImage pictures of a run the AI is told of def _queue(): @@ -27,15 +28,16 @@ def _queue(): return PromptServer.instance.prompt_queue -def outputs_of(record_): - """Saved pictures and videos in one history record.""" +def outputs_of(record_, folder="output"): + """Saved pictures and videos in one history record; with `folder` "temp", what PreviewImage + nodes showed (a pose skeleton, a mask): the AI's to look at, not the user's results.""" out = [] for output in (record_.get("outputs") or {}).values(): for key in ("images", "gifs", "videos"): for item in output.get(key) or []: - if isinstance(item, dict) and item.get("type") == "output" and item.get("filename"): + if isinstance(item, dict) and item.get("type") == folder and item.get("filename"): out.append({"type": "image", "filename": item["filename"], "subfolder": item.get("subfolder") or "", - "folder": "output"}) + "folder": folder}) return out @@ -103,7 +105,7 @@ def _finish_run(task_id, step_id, prompts, history, queued=None): if not prompts: _stop_run(task_id, step_id, "没有找到已排队的运行 / No queued run was found; completion cannot be confirmed") return True - images, failed, done = [], False, 0 + images, previews, failed, done = [], [], False, 0 with ai_tasks._lock: task = ai_tasks._load().get(task_id) stopped = next((item.get("stop") for item in (task or {}).get("steps", []) if item["id"] == step_id), None) @@ -125,6 +127,7 @@ def _finish_run(task_id, step_id, prompts, history, queued=None): for image in outputs_of(record_): if image not in images: # a prompt ComfyUI answers from its cache repeats the same file images.append(image) + previews += [image for image in outputs_of(record_, "temp") if image not in previews] with ai_tasks._lock: task = ai_tasks._load().get(task_id) step = next((item for item in (task or {}).get("steps", []) if item["id"] == step_id), None) @@ -133,10 +136,13 @@ def _finish_run(task_id, step_id, prompts, history, queued=None): failed = failed or bool(step.get("submission_error")) stopped = step.get("stop") results = [item for item in step["results"] if item["type"] != "image"] + images[:ai_tasks.MAX_RESULTS] + previews = previews[-MAX_PREVIEWS:] if (done < len(prompts) and results == step["results"] and step.get("prompts") == len(prompts) - and (not failed or step.get("ok") is False)): + and previews == step.get("previews", []) and (not failed or step.get("ok") is False)): return False # nothing new to keep: no write step["results"] = results + if previews: + step["previews"] = previews step["state"] = "running" if done < len(prompts) else ("cancelled" if stopped else "error" if failed else "done") step["prompts"] = len(prompts) if failed: @@ -369,6 +375,12 @@ def status(client, session="", task_id=None, wait_seconds=0, step_id=None): for item in run.get("results") or [] if item.get("type") == "image"] if pictures: answer["pictures"] = pictures + previews = ["/".join(part for part in (item.get("subfolder"), item["filename"]) if part) for item in run.get("previews") or []] + if previews: + answer["previews"] = previews + answer["previews_note"] = ("What the run's PreviewImage nodes showed (pose skeletons, masks…), in ComfyUI's temp " + "folder: zoom_image(image, folder='temp') or get_image_info(image, folder='temp', " + "include_image=true) shows them; there is no need to run a preprocessor yourself.") if run.get("error"): answer["error"] = run["error"] answer["note"] = { diff --git a/api/mcp_look.py b/api/mcp_look.py new file mode 100644 index 0000000..9902fbb --- /dev/null +++ b/api/mcp_look.py @@ -0,0 +1,88 @@ +"""zoom_image, the AI's closer look: part of a picture (in ComfyUI's output, input or temp folder) +cut out and shown larger, and the same part of a second picture beside it to compare a before +and an after, attached to the answer. It reads only: the pictures stay as they are. With it an AI +need not crop, enlarge or compare pictures with code of its own (each such command asks the user +to allow it). The region is in fractions of the picture, as the user's marks give theirs.""" + +import base64 +import io + +from PIL import Image + +from .mcp_tools import ToolError, _resolve_folder_image, _tool + +MIN_SIZE, MAX_SIZE, DEFAULT_SIZE = 256, 1024, 768 +MAX_ENLARGE = 4 # a tiny part is shown at most this many times its size +MIN_PIXELS = 8 + + +def _box(region, width, height): + """Region {x, y, width, height} (fractions 0-1, default the whole picture) as a pixel box.""" + region = region if isinstance(region, dict) else {} + try: + x, y = (min(1.0, max(0.0, float(region.get(key, 0)))) for key in ("x", "y")) + w, h = (min(1.0, max(0.0, float(region.get(key, 1)))) for key in ("width", "height")) + except (TypeError, ValueError): + raise ToolError("region is {x, y, width, height}, each a fraction 0-1 of the picture.") + left, top = round(x * width), round(y * height) + right, bottom = min(width, round((x + w) * width)), min(height, round((y + h) * height)) + if right - left < MIN_PIXELS or bottom - top < MIN_PIXELS: + raise ToolError(f"The region is under {MIN_PIXELS} pixels across; give a larger one.") + return left, top, right, bottom + + +def _part(path, region, size): + """(facts, MCP image item) of `region` of the picture at `path`, its longer side `size` at most.""" + try: + with Image.open(path) as picture: + box = _box(region, picture.width, picture.height) + part = picture.crop(box).convert("RGB") + whole = f"{picture.width}x{picture.height}" + except (OSError, ValueError, Image.DecompressionBombError): + raise ToolError("The picture cannot be read.") + scale = min(size / max(part.size), MAX_ENLARGE) + shown = part.resize((max(1, round(part.width * scale)), max(1, round(part.height * scale))), Image.LANCZOS) + data = io.BytesIO() + shown.save(data, "JPEG", quality=88) + facts = {"picture": whole, "region": f"x {box[0]}-{box[2]}, y {box[1]}-{box[3]} px", "shown": f"{shown.width}x{shown.height}"} + return facts, {"type": "image", "data": base64.b64encode(data.getvalue()).decode("ascii"), "mimeType": "image/jpeg"} + + +def zoom_image(image, region=None, folder="output", compare_with="", compare_folder="", size=DEFAULT_SIZE): + try: + size = min(MAX_SIZE, max(MIN_SIZE, int(size))) + except (TypeError, ValueError): + raise ToolError(f"size is {MIN_SIZE}-{MAX_SIZE} pixels.") + rel, path = _resolve_folder_image(image, folder) + facts, item = _part(path, region, size) + answer = {"image": rel, "folder": folder, **facts} + extra = [item] + if compare_with: + other_folder = compare_folder or folder + other_rel, other_path = _resolve_folder_image(compare_with, other_folder) + other_facts, other_item = _part(other_path, region, size) + answer["compare_with"] = {"image": other_rel, "folder": other_folder, **other_facts} + answer["note"] = "Two pictures are attached: first image, then compare_with, the same region of each." + extra.append(other_item) + return answer, extra + + +_REGION = {"type": "object", "description": "Fractions 0-1 of the picture, as the user's marks give theirs; default the whole picture.", + "properties": {name: {"type": "number", "minimum": 0, "maximum": 1} for name in ("x", "y", "width", "height")}} +_FOLDER = {"type": "string", "enum": ["output", "input", "temp"], "default": "output"} + +LOOK_TOOLS = dict([ + _tool("zoom_image", "Look closer at a picture", + "Part of a picture shown larger, attached: to check a hand, a face, a seam or a mark's area at full detail. " + "folder: output (results), input (what LoadImage reads, e.g. inpaint_image), temp (a run's previews). " + "compare_with: a second picture, the same region attached beside it (a before and an after). Reads only. " + "Use it instead of cropping or comparing pictures with code of your own; change a picture only by running " + "a ComfyUI workflow (inpaint and the like), never by editing its file.", + zoom_image, {"image": {"type": "string", "description": "As search_images, get_image_info or run_status name it."}, + "region": _REGION, "folder": _FOLDER, + "compare_with": {"type": "string", "description": "A second picture to show the same region of."}, + "compare_folder": {"type": "string", "enum": ["output", "input", "temp"], + "description": "compare_with's folder; default the same as folder."}, + "size": {"type": "integer", "minimum": MIN_SIZE, "maximum": MAX_SIZE, "default": DEFAULT_SIZE}}, + ["image"]), +]) diff --git a/api/mcp_server.py b/api/mcp_server.py index d5f5e05..d23464f 100644 --- a/api/mcp_server.py +++ b/api/mcp_server.py @@ -28,10 +28,11 @@ from .local_guard import is_local as _is_local, loopback_host as _loopback_host from .mcp_actions import ACTIONS from .mcp_canvas_edit import CANVAS_TOOLS from .mcp_model_use import MODEL_TOOLS +from .mcp_look import LOOK_TOOLS from .mcp_task_tools import TASK_TOOLS from .mcp_tools import TOOLS as READ_TOOLS, call_tool -TOOLS = {**READ_TOOLS, **MODEL_TOOLS, **ACTIONS, **CANVAS_TOOLS, **TASK_TOOLS, **mcp_apps.APP_TOOLS} +TOOLS = {**READ_TOOLS, **MODEL_TOOLS, **LOOK_TOOLS, **ACTIONS, **CANVAS_TOOLS, **TASK_TOOLS, **mcp_apps.APP_TOOLS} CAPABILITIES = { "tools": {"listChanged": False}, "prompts": {"listChanged": False}, "resources": {"listChanged": False}, "extensions": {mcp_apps.EXTENSION: {"mimeTypes": [mcp_apps.MIME]}}, # the card in the conversation (mcp_apps.py) diff --git a/api/mcp_task_tools.py b/api/mcp_task_tools.py index 6673c02..a792d39 100644 --- a/api/mcp_task_tools.py +++ b/api/mcp_task_tools.py @@ -29,8 +29,11 @@ def begin_task(ctx, title, new_task=False): "or to hand in pictures you made another way; report with done=true at the end. " "The user's models, pictures, node types, canvas and runs are all in AMB's tools (search_models " "says which node loads a model, move_model fixes a misfiled one, find_nodes lists node types, " - "search_images and get_image_info read pictures): do not read ComfyUI's folders or source, " - "or call its routes, yourself; each of those asks the user to allow it in your app."} + "search_images and get_image_info read pictures, zoom_image shows a part larger or " + "compares two, run_status lists a run's previews): do not read ComfyUI's folders or " + "source, run preprocessors or call its routes yourself; each of those asks the user to " + "allow it in your app. Change pictures only by running ComfyUI workflows (inpaint and " + "the like): never edit, paint over or overwrite picture files with code of your own."} box = ai_toolbox.effective(ai_tasks.get_task(task["id"])) if box: answer["toolbox"] = ai_toolbox.words(box) diff --git a/api/mcp_tools.py b/api/mcp_tools.py index c844ad2..f510027 100644 --- a/api/mcp_tools.py +++ b/api/mcp_tools.py @@ -118,20 +118,25 @@ def _resolve_image(image): return rel, path -def _resolve_input_image(image): - """The picture at `image` ('subfolder/file.png') in ComfyUI's input folder, where LoadImage reads.""" +def _resolve_folder_image(image, folder): + """(rel, path) of picture `image` ('subfolder/file.png') in ComfyUI's output, input (where + LoadImage reads) or temp folder (where PreviewImage writes), checked to stay inside it.""" + if folder == "output": + return _resolve_image(image) + if folder not in ("input", "temp"): + raise ToolError("folder is output, input or temp.") rel = str(image or "").replace("\\", "/").strip("/") try: - path = resolve_within(folder_paths.get_input_directory(), rel) + path = resolve_within(getattr(folder_paths, f"get_{folder}_directory")(), rel) except ValueError: - raise ToolError("image is a path under ComfyUI's input folder, as LoadImage lists it.") + raise ToolError(f"image is a path under ComfyUI's {folder} folder.") if os.path.splitext(path)[1].lower() not in IMAGE_TYPES or not os.path.isfile(path): - raise ToolError(f"No picture at {image} in ComfyUI's input folder.") + raise ToolError(f"No picture at {image} in ComfyUI's {folder} folder.") return rel, path def _input_facts(path): - """Size of an input picture and, when it has an alpha, the area LoadImage gives as MASK (1 - alpha).""" + """Size of an input or preview picture and, when it has an alpha, the area LoadImage gives as MASK (1 - alpha).""" try: with Image.open(path) as picture: facts = {"width": picture.width, "height": picture.height} @@ -323,9 +328,9 @@ def search_images(query="", limit=20, page=1, with_pictures=False): def get_image_info(image, include_image=False, folder="output"): - if folder == "input": - rel, path = _resolve_input_image(image) - info = {"image": rel, "folder": "input", "created": _when(os.path.getmtime(path)), **_input_facts(path)} + if folder != "output": + rel, path = _resolve_folder_image(image, folder) + info = {"image": rel, "folder": folder, "created": _when(os.path.getmtime(path)), **_input_facts(path)} extra = [item for item in [_image_content(path)] if include_image and item] if include_image: info["image_attached"] = bool(extra) @@ -499,11 +504,12 @@ TOOLS = dict([ _tool("get_image_info", "Image generation settings", "The nodes and settings that produced an output image (prompts, models, seed, sampler…) and whether the " "user kept it as a recipe, combo or prompt. With folder=input, a picture in ComfyUI's input folder (what " - "LoadImage reads, e.g. the inpaint_image Anomalous put there): its size and the area LoadImage gives as " - "MASK. include_image attaches a 512 px copy. No need to look at these files any other way.", + "LoadImage reads, e.g. the inpaint_image Anomalous put there), with folder=temp a run's preview (run_status " + "lists them): its size and the area LoadImage gives as MASK. include_image attaches a 512 px copy; " + "zoom_image shows a part larger. No need to look at these files any other way.", get_image_info, {"image": {"type": "string", "description": "From search_images, e.g. 'ComfyUI_00012_.png'; with folder=input, the name LoadImage takes."}, "include_image": {"type": "boolean", "default": False}, - "folder": {"type": "string", "enum": ["output", "input"], "default": "output"}}, ["image"]), + "folder": {"type": "string", "enum": ["output", "input", "temp"], "default": "output"}}, ["image"]), _tool("list_combos", "Combos", "The user's combos (搭配): a main model, LoRAs and a prompt kept together, or (kind \"node " "structure\") model nodes and text nodes saved from the canvas with their links, listed by slot.", list_combos, diff --git a/docs/guides/mcp.md b/docs/guides/mcp.md index 5920df6..9ad4c2f 100644 --- a/docs/guides/mcp.md +++ b/docs/guides/mcp.md @@ -86,6 +86,7 @@ claude mcp add --transport http anomalous http://127.0.0.1:8188/anomalous/mcp |---|---| | library_overview | 各类模型的数量、底模分布、扫描情况，存了多少搭配、工作流和提示词 | | search_models / get_model | 按名字、触发词、底模、来源找模型；看某个模型的 C 站信息、描述、你写的备注，可以附上封面。每个模型都写明哪些节点能加载它；放错文件夹的（比如名字叫 ControlNet、其实是模型补丁的 Anima LLLite）会说明该放哪、放在这里加载不了 | +| zoom_image | 把一张图的某一块放大给 AI 看（手、脸、接缝、标注的区域），也可以把两张图的同一块并排比较（改之前和改之后）。输出、输入、预览图都能看，只看不改。AI 不必再自己写代码裁图、比图 | | scan_report | 哪些模型匹配上了 C 站，哪些没有、原因是什么，哪些还没扫，上次扫描做了什么 | | search_images / get_image_info | 按提示词、模型、参数找出图；看某张图的全部生成参数，可以附上缩略图。`folder=input` 时看 ComfyUI 输入文件夹里的图（比如你标注时 AMB 存的重绘遮罩图）：尺寸、要重画的区域在哪、占多大，也能附上图，不用 AI 自己跑命令去看 | | list_combos | 你的搭配（主模型 + LoRA + 提示词） | diff --git a/tests/test_ai_task_runs.py b/tests/test_ai_task_runs.py index 4b7331f..cbca1c0 100644 --- a/tests/test_ai_task_runs.py +++ b/tests/test_ai_task_runs.py @@ -66,10 +66,14 @@ class ProgressTests(unittest.TestCase): status = ai_task_runs.status("Claude") self.assertEqual((status["state"], status["progress"]["max"]), ("running", 20)) self.assertIn("instead of guessing", status["note"]) - history = {"p1": {"status": {"completed": True}, "outputs": {"9": {"images": [{"filename": "a.png", "subfolder": "x", "type": "output"}]}}}} + 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 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"]), ("done", ["x/a.png"])) + self.assertEqual((status["state"], status["pictures"], status["previews"]), ("done", ["x/a.png"], ["pose_00001_.png"])) + self.assertIn("folder='temp'", status["previews_note"]) + 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): diff --git a/tests/test_mcp_look.py b/tests/test_mcp_look.py new file mode 100644 index 0000000..3097845 --- /dev/null +++ b/tests/test_mcp_look.py @@ -0,0 +1,62 @@ +"""zoom_image (api/mcp_look.py): part of a picture shown larger, the same part of a second one to +compare, in the output, input or temp folder, and nothing outside them. Pictures go to temporary folders.""" +import base64 +import io +import os +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from PIL import Image + +PLUGIN = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])] + +from api import mcp_look, mcp_tools # noqa: E402 + + +def shown(item): + return Image.open(io.BytesIO(base64.b64decode(item["data"]))) + + +class ZoomTests(unittest.TestCase): + def setUp(self): + temp = tempfile.TemporaryDirectory() + self.addCleanup(temp.cleanup) + self.dirs = {name: Path(os.path.realpath(temp.name)) / name for name in ("output", "input", "temp")} + for path in self.dirs.values(): + path.mkdir() + for name in self.dirs: + patcher = mock.patch.object(mcp_tools.folder_paths, f"get_{name}_directory", return_value=str(self.dirs[name])) + patcher.start() + self.addCleanup(patcher.stop) + before = Image.new("RGB", (400, 200), "white") + before.paste((255, 0, 0), (100, 50, 200, 150)) # a red square in the region looked at + before.save(self.dirs["output"] / "before.png") + Image.new("RGB", (400, 200), "blue").save(self.dirs["temp"] / "pose_00001_.png") + + def call(self, **arguments): + return mcp_tools.call_tool(mcp_look.LOOK_TOOLS, "zoom_image", arguments) + + def test_a_region_larger_and_compared(self): + result = self.call(image="before.png", region={"x": 0.25, "y": 0.25, "width": 0.25, "height": 0.5}, + compare_with="pose_00001_.png", compare_folder="temp", size=400) + data = result["structuredContent"] + self.assertEqual((data["picture"], data["region"], data["shown"]), ("400x200", "x 100-200, y 50-150 px", "400x400")) + self.assertEqual(data["compare_with"]["folder"], "temp") + first, second = [item for item in result["content"] if item["type"] == "image"] + red, green, _ = shown(first).getpixel((200, 200)) + self.assertTrue(red > 240 and green < 20, (red, green)) # the red square, enlarged (JPEG: near enough) + self.assertGreater(shown(second).getpixel((10, 10))[2], 240) + + def test_whole_picture_and_refusals(self): + self.assertEqual(self.call(image="before.png")["structuredContent"]["region"], "x 0-400, y 0-200 px") + for arguments in ({"image": "../before.png"}, {"image": "missing.png"}, {"image": "before.png", "folder": "input"}, + {"image": "before.png", "region": {"width": 0.001}}, {"image": "before.png", "folder": "models"}): + self.assertTrue(self.call(**arguments)["isError"], arguments) + + +if __name__ == "__main__": + unittest.main()