Files
DemonGatanjieu-Anomalous_Mo…/tests/test_mcp_model_use.py
T
DemonGatanjieuandClaude Opus 5.5 414e76c935 Read a LoRA's guessed trigger words from its header, and tell the AI what a model needs
- header_words.py: modelspec.trigger_phrase, else the tags most training pictures share
  (ss_tag_frequency, generic ones left out); the scanner keeps them as
  anomalous_header_words when Civitai gives no trigger words, metadata gives headerWords,
  model search finds them.
- model_kind: Flux Schnell told from Dev (no guidance_in), Z-Image from Lumina 2 (caption
  embedder width, a GGUF by parameter count); RULES 3 so kept header answers are read again.
- api/model_guide.py + model_profiles.json: GET /anomalous/model_guide for the model details
  (example settings, header words, the text encoder and VAE a diffusion models file needs,
  found by name), and get_model's header_words and usage for the AI.
- model_advice.example_settings gives the author's example settings as values;
  author_examples says them in words as before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-11-08 22:43:51 +08:00

145 lines
9.4 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_guide, 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", "headerWords": ["glowstyle"]}), \
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.assertEqual(got["models"][0]["header_words"], ["glowstyle"], "no trigger words from Civitai: what its header suggests")
self.assertNotIn("author_weight", got["models"][1], "only a LoRA has one")
examples.assert_called_once_with("/m/a.safetensors", "loras", "abc")
def test_get_model_tells_the_header_words_and_the_familys_usage(self):
def details(model_id, kind, base, meta):
detail = {"model_id": model_id, "type": kind, "base_model": base, "trigger_words": [], "sha256": ""}
with mock.patch.object(mcp_tools, "get_model", return_value=(detail, None)), \
mock.patch.object(mcp_tools, "_resolve_model", return_value=(kind, 0, "/m", "/m/a.safetensors")), \
mock.patch.object(mcp_tools, "get_metadata", return_value=meta), \
mock.patch.object(mcp_model_use, "_about", side_effect=lambda ctx, model: model), \
mock.patch.object(mcp_model_use.model_advice, "author_examples", return_value={}), \
mock.patch.object(model_guide.folder_paths, "get_filename_list",
side_effect=lambda folder: {"text_encoders": ["Anima/qwen_3_06b_base.safetensors"], "vae": []}[folder]):
return mcp_model_use.get_model(CTX, model_id)[0]
anima = details("diffusion_models:0:a.safetensors", "diffusion_models", "", {"base_guess": "Anima"})
self.assertEqual(anima["usage"]["needs"], ["text_encoders: qwen_3_06b (here: Anima/qwen_3_06b_base.safetensors)",
"vae: qwen_image_vae (none here by that name)"])
self.assertIn("batch", anima["usage"]["notes"][0])
self.assertNotIn("header_words", anima, "words are a LoRA's")
lora = details("loras:0:a.safetensors", "loras", "Pony", {"headerWords": ["glowstyle"]})
self.assertEqual(lora["header_words"], ["glowstyle"])
self.assertEqual(lora["usage"]["family"], "Pony")
self.assertNotIn("needs", lora["usage"], "a LoRA needs no text encoder of its own")
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()