Privacy mode left the AI blind twice over: the whole prompt was a token, so it could not help with light, lens or style, and model and LoRA file names still went out as written. - The prompt box has a first field, the gist: the picture in words the user writes for the AI to read. seal_answer hands it on as written beside the tokens of the private prompts; either alone can be handed in. - ai_sealed_models.py: while the Gates window's "Model names in a private task" is on (the default), each model of the library is one token ([M1], sealed as its model_id) however an answer writes it: model_id, loader name, file name in a text, card name. A model's own entry drops name, file, folder, links and hash and keeps type, base model, the author's settings and the user's note. expand writes a model token back as the model_id, as the loader's name among a node's values, as the card's name inside the AI's words. The task's menu lists what each token stands for. - The model editor has "Note for the AI" (ai_note, note_for_ai in the tools, searched by search_models): what a model is for, in words the AI may read. - A combo's name is sealed too: unless renamed it is made of its model's name and the start of its prompt. - The privacy skill says how to write around a gist and how to pick and pass on model tokens. Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
185 lines
10 KiB
Python
185 lines
10 KiB
Python
"""The user's layer comes first: editor edits, covers and what a scan may touch."""
|
|
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])]
|
|
|
|
import isolated_layout # noqa: F401 - side files beside the model, whatever this computer chose
|
|
import scraper
|
|
from aiohttp.test_utils import make_mocked_request
|
|
|
|
from api import metadata as metadata_module, model_media, model_metadata, scan_summary
|
|
from api.metadata import get_metadata
|
|
|
|
CIVITAI = {"id": 7, "modelId": 3, "name": "v2", "model": {"name": "Found"},
|
|
"images": [{"url": "https://example.invalid/cover.png"}]}
|
|
|
|
|
|
def write(path, data):
|
|
Path(path).write_text(json.dumps(data), encoding="utf-8")
|
|
|
|
|
|
class Request(dict): # aiohttp requests also hold values for the middleware
|
|
def __init__(self, body):
|
|
super().__init__()
|
|
self.body = body
|
|
|
|
async def json(self):
|
|
return self.body
|
|
|
|
|
|
class LayerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.dir = Path(self.temp.name)
|
|
self.model = self.dir / "m.safetensors"
|
|
self.model.write_bytes(b"model")
|
|
self.base = str(self.dir / "m")
|
|
metadata_module.clear_metadata_cache()
|
|
|
|
def tearDown(self):
|
|
self.temp.cleanup()
|
|
|
|
def scan(self, *flags, cover=b"new civitai image"):
|
|
"""Runs the scraper with Civitai and downloads faked; returns what went to the network."""
|
|
calls = {"fetch": 0, "download": 0}
|
|
|
|
def fetch(_hash):
|
|
calls["fetch"] += 1
|
|
return json.loads(json.dumps(CIVITAI))
|
|
|
|
def download(_url, base_path):
|
|
calls["download"] += 1
|
|
Path(base_path + ".png").write_bytes(cover)
|
|
return base_path + ".png"
|
|
|
|
with mock.patch.object(sys, "argv", ["scraper.py", str(self.dir), *flags]), \
|
|
mock.patch.object(scraper, "fetch_civitai_info", side_effect=fetch), \
|
|
mock.patch.object(scraper, "download_media", side_effect=download), \
|
|
mock.patch.object(scraper, "calculate_sha256", side_effect=lambda p: "b" * 64 if p == str(self.model) else open(p, "rb").read().hex()), \
|
|
mock.patch.object(scraper.urllib.request, "urlopen", side_effect=OSError("no network in tests")):
|
|
scraper.main()
|
|
metadata_module.clear_metadata_cache()
|
|
return calls
|
|
|
|
def test_an_edited_but_unscanned_model_still_counts_as_new_and_gets_scanned(self):
|
|
write(self.base + ".civitai.info", {"anomalous_custom_name": "Mine"}) # how older versions saved edits
|
|
with mock.patch.object(scan_summary, "get_active_scan_paths", return_value=[str(self.dir)]):
|
|
self.assertEqual(scan_summary._summarize()["new"], 1)
|
|
self.assertEqual(self.scan("--virtual-rename")["fetch"], 1)
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual(meta["info_source"], "civitai")
|
|
self.assertEqual(meta["custom_name"], "Mine") # the old edit still wins over Civitai's name
|
|
|
|
def test_user_layer_wins_and_an_emptied_field_stays_empty(self):
|
|
write(self.base + ".info", {"id": 7, "modelId": 3, "name": "v2", "anomalous_custom_name": "Found_v2"})
|
|
write(self.base + ".anomalous.json", {"format": 1, "custom_name": "Mine", "custom_notes": ""})
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual(meta["custom_name"], "Mine")
|
|
self.assertEqual(meta["custom_notes"], "")
|
|
self.assertEqual(sorted(meta["user_fields"]), ["custom_name", "custom_notes"])
|
|
self.assertEqual(meta["info_source"], "civitai")
|
|
|
|
def test_info_source_tells_local_inference_and_unscanned_apart(self):
|
|
self.assertEqual(get_metadata(str(self.model))["info_source"], "")
|
|
write(self.base + ".info", {"id": -1, "modelId": -1, "name": "m"})
|
|
metadata_module.clear_metadata_cache()
|
|
self.assertEqual(get_metadata(str(self.model))["info_source"], "local")
|
|
|
|
def test_editor_writes_only_the_user_layer(self):
|
|
write(self.base + ".civitai.info", {"id": 7, "modelId": 3, "name": "from another tool"})
|
|
before = Path(self.base + ".civitai.info").read_text(encoding="utf-8")
|
|
with mock.patch.object(model_metadata, "resolve_folder_subdir", return_value=(str(self.dir), str(self.dir))):
|
|
asyncio.run(model_metadata.api_update_metadata(Request({
|
|
"type": "loras", "filename": "m.safetensors", "custom_name": "Mine", "custom_notes": "note", "ai_note": " two-person poses "})))
|
|
self.assertEqual(Path(self.base + ".civitai.info").read_text(encoding="utf-8"), before)
|
|
user = json.loads(Path(self.base + ".anomalous.json").read_text(encoding="utf-8"))
|
|
self.assertEqual((user["custom_name"], user["custom_notes"], user["ai_note"]), ("Mine", "note", "two-person poses"))
|
|
metadata_module.clear_metadata_cache()
|
|
self.assertEqual(get_metadata(str(self.model))["ai_note"], "two-person poses", "the note for an AI is the user's too")
|
|
|
|
def edit(self, **fields):
|
|
with mock.patch.object(model_metadata, "resolve_folder_subdir", return_value=(str(self.dir), str(self.dir))):
|
|
return asyncio.run(model_metadata.api_update_metadata(Request({"type": "loras", "filename": "m.safetensors", **fields})))
|
|
|
|
def test_the_editor_keeps_a_source_link_and_refuses_what_is_no_web_address(self):
|
|
self.assertEqual(self.edit(custom_source_url="https://www.liblib.art/modelinfo/abc").status, 200)
|
|
user = json.loads(Path(self.base + ".anomalous.json").read_text(encoding="utf-8"))
|
|
self.assertEqual(user["source_url"], "https://www.liblib.art/modelinfo/abc")
|
|
for bad in ("javascript:alert(1)", "file:///C:/x.png", "https://"):
|
|
self.assertEqual(self.edit(custom_source_url=bad).status, 400)
|
|
self.assertEqual(json.loads(Path(self.base + ".anomalous.json").read_text(encoding="utf-8")), user)
|
|
self.edit(custom_source_url="") # emptied on purpose
|
|
metadata_module.clear_metadata_cache()
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual((meta["source_url"], "source_url" in meta["user_fields"]), ("", True))
|
|
|
|
def test_the_editor_learns_whether_the_cover_shown_is_civitais(self):
|
|
def covers():
|
|
request = make_mocked_request("GET", "/anomalous/model_covers?type=loras&path_idx=0&subfolder=/&filename=m.safetensors")
|
|
with mock.patch.object(model_media, "resolve_folder_subdir", return_value=(str(self.dir), str(self.dir))):
|
|
return json.loads(asyncio.run(model_media.api_model_covers(request)).text)
|
|
|
|
self.assertEqual(covers(), {"civitai": "", "custom": False}) # no cover at all
|
|
Path(self.base + ".preview.png").write_bytes(b"my own picture")
|
|
self.assertEqual(covers(), {"civitai": "", "custom": True})
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"civitai image")
|
|
found = covers()
|
|
self.assertTrue(found["custom"])
|
|
self.assertIn("filename=m.civitai_bak.png", found["civitai"])
|
|
Path(self.base + ".preview.png").write_bytes(b"civitai image") # a copy of Civitai's
|
|
self.assertFalse(covers()["custom"])
|
|
|
|
def test_a_scan_keeps_what_a_loras_header_suggests_when_civitai_gives_no_words(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({"10_glow": {"glowstyle": 12, "1girl": 12, "hat": 2}})}}
|
|
raw = json.dumps(header).encode("utf-8")
|
|
self.model.write_bytes(struct.pack("<Q", len(raw)) + raw)
|
|
self.scan()
|
|
self.assertEqual(json.loads(Path(self.base + ".info").read_text(encoding="utf-8"))["anomalous_header_words"], ["glowstyle"])
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual((meta["trainedWords"], meta["headerWords"]), ([], ["glowstyle"]))
|
|
with mock.patch.dict(CIVITAI, trainedWords=["glow"]): # Civitai has words now: the guess goes
|
|
self.scan("--force-overwrite")
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual((meta["trainedWords"], meta["headerWords"]), (["glow"], []))
|
|
self.assertNotIn("anomalous_header_words", json.loads(Path(self.base + ".info").read_text(encoding="utf-8")))
|
|
|
|
def test_a_scan_does_not_refetch_what_is_on_disk(self):
|
|
write(self.base + ".info", {**CIVITAI, "anomalous_custom_name": "Found_v2"})
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"old civitai image")
|
|
Path(self.base + ".preview.png").write_bytes(b"old civitai image")
|
|
self.assertEqual(self.scan("--virtual-rename"), {"fetch": 0, "download": 0})
|
|
|
|
def test_fetching_again_keeps_the_users_cover(self):
|
|
write(self.base + ".info", CIVITAI)
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"old civitai image")
|
|
Path(self.base + ".preview.png").write_bytes(b"my own picture")
|
|
self.scan("--force-overwrite", "--virtual-rename")
|
|
self.assertEqual(Path(self.base + ".preview.png").read_bytes(), b"my own picture")
|
|
self.assertEqual(Path(self.base + ".civitai_bak.png").read_bytes(), b"new civitai image")
|
|
|
|
def test_fetching_again_replaces_an_unchanged_civitai_cover(self):
|
|
write(self.base + ".info", CIVITAI)
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"old civitai image")
|
|
Path(self.base + ".preview.png").write_bytes(b"old civitai image")
|
|
with mock.patch.object(scraper, "move_to_trash") as trash:
|
|
self.scan("--force-overwrite")
|
|
# The cover is Civitai's: the copy kept aside goes, the new picture replaces it in place.
|
|
trash.assert_called_once_with(self.base + ".civitai_bak.png")
|
|
self.assertEqual(Path(self.base + ".preview.png").read_bytes(), b"new civitai image")
|
|
info = json.loads(Path(self.base + ".info").read_text(encoding="utf-8"))
|
|
self.assertEqual(len(info["anomalous_cover_sha256"]), 64) # known as Civitai's from now on
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|