- 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>
91 lines
4.5 KiB
Python
91 lines
4.5 KiB
Python
"""What using a model takes besides its file (api/model_guide.py): the family's files a drawing
|
|
model loaded in parts needs, found by name in this ComfyUI's folders, and the model details' request,
|
|
which names a model by type, folder index, subfolder and file name only."""
|
|
import asyncio
|
|
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 aiohttp.test_utils import make_mocked_request # noqa: E402
|
|
|
|
from api import metadata, model_guide # noqa: E402
|
|
|
|
|
|
def listing(files):
|
|
return mock.patch.object(model_guide.folder_paths, "get_filename_list", side_effect=lambda folder: files.get(folder, []))
|
|
|
|
|
|
class FamilyTests(unittest.TestCase):
|
|
def test_the_first_family_whose_mark_starts_the_base_model(self):
|
|
self.assertEqual(model_guide.profile("Flux.1 S")["family"], "Flux.1 Schnell")
|
|
self.assertEqual(model_guide.profile("Flux.1 Krea")["family"], "Flux.1")
|
|
self.assertEqual(model_guide.profile("NoobAI")["family"], "Illustrious / NoobAI")
|
|
self.assertEqual(model_guide.profile("ZImageTurbo")["family"], "Z-Image")
|
|
self.assertIsNone(model_guide.profile(""))
|
|
self.assertIsNone(model_guide.profile("Other"))
|
|
|
|
def test_parts_are_found_by_name_and_only_for_a_model_loaded_in_parts(self):
|
|
files = {"text_encoders": ["clip_l.safetensors", "t5/t5-v1_1-xxl-encoder-Q8_0.gguf"], "vae": ["aesthetic_vae.safetensors", "flux/ae.safetensors"]}
|
|
with listing(files):
|
|
self.assertEqual([part["found"] for part in model_guide.parts("unet", "Flux.1 D")],
|
|
["clip_l.safetensors", "t5/t5-v1_1-xxl-encoder-Q8_0.gguf", "flux/ae.safetensors"])
|
|
self.assertEqual(model_guide.parts("checkpoints", "Flux.1 D"), [], "a checkpoint carries its own")
|
|
self.assertEqual(model_guide.parts("diffusion_models", "SDXL"), [])
|
|
with listing({"vae": ["aesthetic_vae.safetensors"]}):
|
|
self.assertEqual(model_guide.parts("unet", "Chroma")[1]["found"], "", "a short name must be the whole name")
|
|
|
|
def test_usage_for_the_ai(self):
|
|
with listing({}):
|
|
self.assertEqual(model_guide.usage("diffusion_models", "Chroma")["needs"],
|
|
["text_encoders: t5xxl (none here by that name)", "vae: ae (Flux VAE) (none here by that name)"])
|
|
self.assertEqual(model_guide.usage("loras", "Pony")["family"], "Pony")
|
|
self.assertEqual(model_guide.usage("loras", "Hunyuan Video"), {})
|
|
|
|
|
|
class RequestTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.folder = os.path.realpath(self.temp.name)
|
|
self.addCleanup(self.temp.cleanup)
|
|
def resolve(kind, at, sub):
|
|
if kind != "loras" or at != 0:
|
|
raise KeyError(kind)
|
|
return self.folder, self.folder
|
|
patch = mock.patch.object(model_guide, "resolve_folder_subdir", side_effect=resolve)
|
|
patch.start()
|
|
self.addCleanup(patch.stop)
|
|
metadata.clear_metadata_cache()
|
|
|
|
def ask(self, filename, kind="loras", at="0"):
|
|
request = make_mocked_request("GET", f"/anomalous/model_guide?type={kind}&path_idx={at}&subfolder=&filename={filename}")
|
|
response = asyncio.run(model_guide.api_model_guide(request))
|
|
return response.status, json.loads(response.body) if response.status == 200 else None
|
|
|
|
def test_a_lora_scanned_before_gets_its_words_read_now(self):
|
|
header = {"lora_unet_input_blocks_4_1_proj_in.lora_down.weight": {"dtype": "F16", "shape": [4, 4], "data_offsets": [0, 0]},
|
|
"__metadata__": {"ss_tag_frequency": json.dumps({"d": {"glowstyle": 4}})}}
|
|
raw = json.dumps(header).encode("utf-8")
|
|
Path(self.folder, "glow.safetensors").write_bytes(struct.pack("<Q", len(raw)) + raw)
|
|
status, guide = self.ask("glow.safetensors")
|
|
self.assertEqual((status, guide), (200, {"examples": {}, "words": ["glowstyle"], "parts": []}))
|
|
|
|
def test_only_a_model_file_inside_its_folder(self):
|
|
Path(self.folder, "note.png").write_bytes(b"x")
|
|
self.assertEqual(self.ask("note.png")[0], 404)
|
|
self.assertEqual(self.ask("missing.safetensors")[0], 404)
|
|
self.assertEqual(self.ask("..%2Fother.safetensors")[0], 400)
|
|
self.assertEqual(self.ask("x.safetensors", at="bad")[0], 400)
|
|
self.assertEqual(self.ask("x.safetensors", kind="nope")[0], 400)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|