Files
DemonGatanjieuandClaude Opus 5.5 46eca1b594 Run a picture's own workflow, keep a task on one page, read a big canvas whole
- run_workflow takes a gallery picture as its workflow (mcp_fill.opening
  says what a name opens, for open_workflow too); a picture's workflow
  opens with an id of its own, and the tab's mark goes into the workflow
  before it loads, so ComfyUI finishing the tab does not drop it.
- With ComfyUI open in two browser tabs a task's actions were split
  between them (one opened the task's tab, the other could not find it):
  the page that did the last action gets the next first now, another
  only when it is gone (mcp_bridge.py, pages keep an id of their own).
- describe_canvas on a big canvas: titles off, then nodes as a line each
  with their settings, so every value stays readable.
- search_models names the models that did not fit; begin_task's tool
  directory goes by name alone when the answer would be too long.
- ask_user refuses a "test" question with made-up options too.
- Exam round recorded: question 2 with nothing filed away, a big canvas
  read right in 3 calls, question 3 in 12 calls.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-11-07 01:16:17 +08:00

123 lines
7.5 KiB
Python

"""What loads a model, for the AI (api/mcp_model_use.py): a model patch filed under ControlNet is
said to belong in model_patches, move_model puts it there, and a copy already there is never
overwritten. Files go to temporary folders."""
import json
import os
import struct
import sys
import tempfile
import unittest
from pathlib import Path
from unittest import mock
PLUGIN = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])]
from api import mcp_canvas_edit, mcp_model_use, mcp_tools, model_placement # noqa: E402
CTX = {"base_url": "http://127.0.0.1:1"}
def safetensors(*names):
raw = json.dumps({name: {"dtype": "F16", "shape": [4, 4], "data_offsets": [0, 0]} for name in names}).encode()
return struct.pack("<Q", len(raw)) + raw
LLLITE = safetensors("lllite_conditioning1.conv1.weight", "lllite_modules.0.down.weight")
CONTROL = safetensors("control_model.input_hint_block.0.weight", "control_model.zero_convs.0.0.weight")
class ModelUseTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
base = os.path.realpath(self.temp.name)
self.roots = {kind: [os.path.join(base, kind)] for kind in ("controlnet", "model_patches")}
for paths in self.roots.values():
os.makedirs(paths[0])
listing = lambda kind: sorted(os.path.relpath(os.path.join(d, f), self.roots[kind][0]) for d, _, names in os.walk(self.roots[kind][0]) for f in names)
info = lambda ctx: {"ControlNetLoader": {"input": {"required": {"control_net_name": [listing("controlnet")]}}},
"ModelPatchLoader": {"input": {"required": {"name": [listing("model_patches")]}}}}
for patch in (
mock.patch.object(model_placement.folder_paths, "folder_names_and_paths", {k: (v, set()) for k, v in self.roots.items()}),
mock.patch.object(model_placement.folder_paths, "get_folder_paths", side_effect=lambda t: self.roots.get(t, [])),
mock.patch("api.model_import.folder_paths.get_folder_paths", side_effect=lambda t: self.roots.get(t, [])),
mock.patch.object(mcp_model_use.folder_paths, "get_filename_list", side_effect=listing),
mock.patch.object(mcp_tools, "resolve_folder_subdir",
side_effect=lambda t, i, s: (self.roots[t][0], os.path.join(self.roots[t][0], s.strip("/")))),
mock.patch.object(mcp_canvas_edit, "object_info", side_effect=info),
mock.patch.object(model_placement, "add_entry"),
mock.patch.object(model_placement, "forget_file_lists"),
mock.patch.object(model_placement, "get_metadata", return_value={}),
):
patch.start()
self.addCleanup(patch.stop)
model_placement._heads.clear()
self.addCleanup(self.temp.cleanup)
def put(self, kind, name, data):
Path(self.roots[kind][0], name).write_bytes(data)
return f"{kind}:0:{name}"
def about(self, model_id):
return mcp_model_use._about(CTX, {"model_id": model_id})
def test_a_loras_author_weight_comes_with_the_search(self):
found = {"total_matches": 2, "models": [{"model_id": "loras:0:a.safetensors", "type": "loras"},
{"model_id": "checkpoints:0:b.safetensors", "type": "checkpoints"}]}
with mock.patch.object(mcp_tools, "search_models", return_value=found), \
mock.patch.object(mcp_model_use, "_about", side_effect=lambda ctx, model: model), \
mock.patch.object(mcp_tools, "_resolve_model", return_value=("loras", 0, "/m", "/m/a.safetensors")), \
mock.patch.object(mcp_tools, "get_metadata", return_value={"hash": "abc"}), \
mock.patch.object(mcp_model_use.model_advice, "author_examples", return_value={"lora_weight": "0.7-0.8"}) as examples:
got = mcp_model_use.search_models(CTX, "a")
self.assertEqual(got["models"][0]["author_weight"], "0.7-0.8")
self.assertNotIn("author_weight", got["models"][1], "only a LoRA has one")
examples.assert_called_once_with("/m/a.safetensors", "loras", "abc")
def test_the_search_answer_stays_readable(self):
lora = lambda at: {"model_id": f"loras:0:style/l{at}.safetensors", "name": f"Style {at}", "file": f"l{at}.safetensors", # noqa: E731
"type": "loras", "folder": "style", "base_model": "Illustrious",
"trigger_words": [f"word{n}" for n in range(9)], "loaded_by": ["LoraLoader"]}
models = [lora(at) for at in range(40)]
models.append({"model_id": "loras:0:wrong.safetensors", "name": "Wrong", "file": "wrong.safetensors", "type": "loras",
"folder": "", "belongs_in": "model_patches", "placement": "moves"})
got = mcp_model_use._compact({"total_matches": 41, "models": models})
self.assertLessEqual(mcp_model_use.mcp_fit.size(got), mcp_model_use.mcp_fit.BUDGET)
self.assertNotIn("loaded_by", got, "a misfiled one shares no loader")
self.assertEqual(got["models"][0]["loaded_by"], ["LoraLoader"])
self.assertIn("named in more_names", got["left_out"])
self.assertEqual(len(got["more_names"]) + len(got["models"]), 41, "every model is there, the rest by name")
self.assertEqual(got["more_names"][-1], "Wrong")
self.assertEqual((got["models"][0]["trigger_words"], got["models"][0]["more_trigger_words"]), (["word0", "word1", "word2", "word3"], 5))
self.assertFalse({"file", "folder", "type"} & set(got["models"][0]), "the model_id says them")
shared = mcp_model_use._compact({"total_matches": 2, "models": [lora(0), lora(1)]})
self.assertEqual((shared["loaded_by"], "loaded_by" in shared["models"][0]), ({"loras": ["LoraLoader"]}, False), "said once")
def test_a_misfiled_patch_is_said_to_belong_elsewhere_and_moves_there(self):
self.put("controlnet", "depth.safetensors", CONTROL)
patch = self.put("controlnet", "anima-lllite-depth.safetensors", LLLITE)
self.assertEqual(self.about("controlnet:0:depth.safetensors")["loaded_by"], ["ControlNetLoader"])
misfiled = self.about(patch)
self.assertEqual(misfiled["belongs_in"], "model_patches")
self.assertNotIn("loaded_by", misfiled) # listed by the ControlNet loader, which cannot read it
self.assertIn("move_model", misfiled["placement"])
moved = mcp_model_use.move_model(CTX, patch, "model_patches")
self.assertEqual(moved["model_id"], "model_patches:0:anima-lllite-depth.safetensors")
self.assertEqual(moved["loaded_by"], ["ModelPatchLoader"])
self.assertTrue(Path(self.roots["model_patches"][0], "anima-lllite-depth.safetensors").is_file())
self.assertFalse(Path(self.roots["controlnet"][0], "anima-lllite-depth.safetensors").exists())
with self.assertRaises(mcp_tools.ToolError): # where it belongs now: nothing to move
mcp_model_use.move_model(CTX, moved["model_id"], "model_patches")
def test_a_copy_already_in_place_is_named_and_never_overwritten(self):
patch = self.put("controlnet", "pose.safetensors", LLLITE)
self.put("model_patches", "pose.safetensors", LLLITE)
self.assertIn("model_patches:0:pose.safetensors", self.about(patch)["placement"])
with self.assertRaises(mcp_tools.ToolError):
mcp_model_use.move_model(CTX, patch, "model_patches")
self.assertTrue(Path(self.roots["controlnet"][0], "pose.safetensors").is_file())
if __name__ == "__main__":
unittest.main()