DemonGatanjieuandClaude Opus 5.5
6d1776d856
Manager reads every .py file of a plugin, tests too, and listed ten names from two test files as AMB's nodes, the only match for Power Lora Loader among them, so a workflow missing it could offer AMB in place of rgthree. The tables are written so its scanner passes them by, and a new test uses the scanner's own patterns to keep any file from looking like one again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
254 lines
15 KiB
Python
254 lines
15 KiB
Python
"""The Tidy tab's "Not used lately" (api/model_usage.py, api/usage_ledger.py): when each model and
|
|
node pack was last used, from the run record, output images and sounds and saved workflows in
|
|
temporary folders."""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])]
|
|
|
|
from api import image_search, model_usage, usage_ledger # noqa: E402
|
|
|
|
DAY = 24 * 3600
|
|
NOW = 1_800_000_000.0
|
|
|
|
|
|
def output(path, when, nodes):
|
|
"""A PNG whose prompt holds `nodes` ({id: (class_type, inputs)}), dated `when`."""
|
|
info = PngInfo()
|
|
info.add_text("prompt", json.dumps({key: {"class_type": kind, "inputs": inputs} for key, (kind, inputs) in nodes.items()}))
|
|
Image.new("RGB", (2, 2)).save(path, pnginfo=info)
|
|
os.utime(path, (when, when))
|
|
|
|
|
|
def node(module, takes=(), makes=()):
|
|
"""A node class of the pack `module`, taking and making those kinds of thing."""
|
|
return type("Node", (), {
|
|
"RELATIVE_PYTHON_MODULE": module, "RETURN_TYPES": tuple(makes),
|
|
"INPUT_TYPES": classmethod(lambda cls: {"required": {f"in{index}": (kind, {}) for index, kind in enumerate(takes)}}),
|
|
})
|
|
|
|
|
|
def stat(born, changed):
|
|
return types.SimpleNamespace(st_birthtime=born, st_ctime=born, st_mtime=changed)
|
|
|
|
|
|
class UsageTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
base = Path(self.temp.name)
|
|
self.out, self.flows, self.recipes = base / "output", base / "workflows", base / "recipes"
|
|
custom = base / "custom_nodes"
|
|
for folder in (self.out / "sub", self.flows, self.recipes):
|
|
folder.mkdir(parents=True)
|
|
packs = ["rgthree-comfy", "ComfyUI-Easy-Use", "forgotten-pack", "ComfyUI-Impact-Pack", "comfyui-impact-subpack",
|
|
"Genie-TTS", "page-only", "ComfyUI-Manager", os.path.basename(model_usage.PLUGIN_DIR)]
|
|
for name in packs:
|
|
(custom / name).mkdir(parents=True)
|
|
(custom / "rgthree-comfy" / "js").mkdir()
|
|
(custom / "page-only" / "web").mkdir()
|
|
self.ledger = {"started": None, "types": {}, "models": {}, "opened": {"types": {}, "models": {}}}
|
|
self.models = [
|
|
{"type": "checkpoints", "root": 0, "rel": "sdxl/fresh.safetensors", "size": 6},
|
|
{"type": "loras", "root": 0, "rel": "old.safetensors", "size": 3},
|
|
{"type": "loras", "root": 0, "rel": "style/kept.safetensors", "size": 2},
|
|
{"type": "loras", "root": 0, "rel": "never.safetensors", "size": 9},
|
|
{"type": "vae", "root": 0, "rel": "hashed.safetensors", "size": 1},
|
|
]
|
|
registered = { # by this name, not as the keyword's value: tests/test_manager_scan.py says why
|
|
"KSampler": node("nodes", ["MODEL"], ["LATENT"]),
|
|
"Power Lora Loader (rgthree)": node("custom_nodes.rgthree-comfy", ["MODEL"], ["MODEL"]),
|
|
"easy loraStack": node("custom_nodes.ComfyUI-Easy-Use"),
|
|
"easy seed": node("custom_nodes.ComfyUI-Easy-Use.py.nodes"),
|
|
"Unused Thing": node("custom_nodes.forgotten-pack"),
|
|
"DetailerForEach": node("custom_nodes.ComfyUI-Impact-Pack", ["IMAGE", "BBOX_DETECTOR"], ["IMAGE"]),
|
|
"UltralyticsDetectorProvider": node("custom_nodes.comfyui-impact-subpack", [], ["BBOX_DETECTOR"]),
|
|
"Genie Speak": node("custom_nodes.Genie-TTS", ["STRING"], ["AUDIO"]),
|
|
"Manager Thing": node("custom_nodes.ComfyUI-Manager"),
|
|
}
|
|
nodes = types.SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS=registered,
|
|
LOADED_MODULE_DIRS={name: str(custom / name) for name in packs},
|
|
EXTENSION_WEB_DIRS={"rgthree-comfy": str(custom / "rgthree-comfy" / "js"), "page-only": str(custom / "page-only" / "web")},
|
|
)
|
|
patches = [
|
|
mock.patch.object(model_usage, "library", return_value=self.models),
|
|
mock.patch.object(model_usage, "_card", side_effect=lambda item: {"name": Path(item["rel"]).stem, "version": "", "preview_url": ""}),
|
|
mock.patch.object(model_usage, "_known_hash", side_effect=lambda item: "AB" * 32 if item["rel"] == "hashed.safetensors" else ""),
|
|
mock.patch.object(model_usage, "part_dirs", return_value={"comfy_workflows": str(self.flows), "recipes": str(self.recipes), "materials": str(base / "nope")}),
|
|
mock.patch.object(model_usage.folder_paths, "get_output_directory", return_value=str(self.out)),
|
|
mock.patch.object(model_usage.folder_paths, "get_folder_paths", side_effect=lambda kind: [str(custom)] if kind == "custom_nodes" else []),
|
|
mock.patch.object(model_usage.usage_ledger, "read", side_effect=lambda: self.ledger),
|
|
mock.patch.object(model_usage, "kept_path", return_value=str(base / "usage_kept.json")),
|
|
mock.patch.dict(sys.modules, {"nodes": nodes}),
|
|
]
|
|
for patch in patches:
|
|
patch.start()
|
|
self.addCleanup(patch.stop)
|
|
image_search._record_cache.clear()
|
|
self.addCleanup(self.temp.cleanup)
|
|
|
|
def test_models_and_packs_by_when_they_were_last_used(self):
|
|
output(self.out / "a.png", NOW - 2 * DAY, {"1": ("CheckpointLoaderSimple", {"ckpt_name": "sdxl\\fresh.safetensors"}),
|
|
"2": ("KSampler", {"seed": 1})})
|
|
output(self.out / "sub" / "b.png", NOW - 100 * DAY, {"1": ("LoraLoader", {"lora_name": "old.safetensors"}),
|
|
"2": ("Power Lora Loader (rgthree)", {"lora_1": {"on": True, "lora": "old.safetensors"}})})
|
|
output(self.out / "c.png", NOW - 40 * DAY, {"1": ("easy seed", {"seed": 2})})
|
|
(self.flows / "portrait.json").write_text(json.dumps({"nodes": [
|
|
{"type": "LoraLoader", "widgets_values": ["style/kept.safetensors", 1, 1]},
|
|
{"type": "easy loraStack", "widgets_values": [True, "simple", 1]}]}), encoding="utf-8")
|
|
(self.recipes / "r.json").write_text(json.dumps({"models": [{"sha256": "ab" * 32}]}), encoding="utf-8")
|
|
|
|
result = model_usage.check(now=NOW)
|
|
models = {Path(item["rel"]).stem: item for item in result["models"]}
|
|
self.assertNotIn("fresh", models) # used two days ago
|
|
self.assertEqual((models["old"]["last_used"], models["old"]["outputs"], models["old"]["saved"]), (NOW - 100 * DAY, 1, 0)) # one image naming it twice
|
|
self.assertEqual((models["kept"]["last_used"], models["kept"]["saved"]), (None, 1)) # only in a saved workflow
|
|
self.assertEqual((models["never"]["last_used"], models["never"]["outputs"], models["never"]["saved"]), (None, 0, 0))
|
|
self.assertEqual(models["hashed"]["saved"], 1) # its SHA-256 is in a recipe
|
|
self.assertEqual(result["seen"], {"outputs": 3, "saved": 2, "capped": False})
|
|
self.assertIsNone(result["ledger_since"])
|
|
|
|
packs = {item["name"]: item for item in result["packs"]}
|
|
self.assertEqual(sorted(packs, key=str.lower), ["ComfyUI-Easy-Use", "ComfyUI-Impact-Pack", "comfyui-impact-subpack", "forgotten-pack",
|
|
"Genie-TTS", "page-only", "rgthree-comfy"]) # ComfyUI's own nodes, the Manager and this plugin are left out
|
|
self.assertEqual((packs["rgthree-comfy"]["last_used"], packs["rgthree-comfy"]["outputs"], packs["rgthree-comfy"]["web"]), (NOW - 100 * DAY, 1, True))
|
|
self.assertEqual((packs["ComfyUI-Easy-Use"]["nodes"], packs["ComfyUI-Easy-Use"]["last_used"], packs["ComfyUI-Easy-Use"]["saved"]),
|
|
(2, NOW - 40 * DAY, 1))
|
|
self.assertEqual((packs["forgotten-pack"]["last_used"], packs["forgotten-pack"]["saved"], packs["forgotten-pack"]["web"]), (None, 0, False))
|
|
self.assertEqual((packs["page-only"]["nodes"], packs["page-only"]["web"]), (0, True)) # works on the page only
|
|
self.assertEqual(packs["comfyui-impact-subpack"]["feeds"], ["ComfyUI-Impact-Pack"]) # makes what only Impact's nodes take
|
|
self.assertEqual(packs["ComfyUI-Impact-Pack"]["feeds"], []) # IMAGE is ComfyUI's own
|
|
self.assertEqual(packs["comfyui-impact-subpack"]["words"], "impact subpack")
|
|
self.assertEqual(packs["ComfyUI-Easy-Use"]["words"], "easy use")
|
|
|
|
def test_the_run_record_and_sounds_count_as_use(self):
|
|
self.ledger = {"started": NOW - 10 * DAY, "types": {"Genie Speak": [NOW - 3 * DAY, 4]},
|
|
"models": {"never.safetensors": [NOW - 31 * DAY, 2]},
|
|
"opened": {"types": {}, "models": {"style/kept.safetensors": [NOW - 45 * DAY, 3]}}}
|
|
sound = self.out / "speech.flac"
|
|
sound.write_bytes(b"")
|
|
os.utime(sound, (NOW - 50 * DAY, NOW - 50 * DAY))
|
|
prompt = {"1": {"class_type": "Unused Thing", "inputs": {"model": "old.safetensors"}}}
|
|
with mock.patch.object(model_usage, "read_audio_tags", return_value={"prompt": json.dumps(prompt)}):
|
|
result = model_usage.check(now=NOW)
|
|
models = {Path(item["rel"]).stem: item for item in result["models"]}
|
|
packs = {item["name"]: item for item in result["packs"]}
|
|
self.assertEqual(result["ledger_since"], NOW - 10 * DAY)
|
|
self.assertEqual((packs["Genie-TTS"]["last_used"], packs["Genie-TTS"]["runs"]), (NOW - 3 * DAY, 4)) # a sound, never an image
|
|
self.assertEqual((models["never"]["last_used"], models["never"]["runs"]), (NOW - 31 * DAY, 2))
|
|
self.assertEqual((models["kept"]["last_used"], models["kept"]["opened"], models["kept"]["runs"]), (NOW - 45 * DAY, 3, 0)) # opened, never run
|
|
self.assertEqual((packs["forgotten-pack"]["last_used"], models["old"]["last_used"]), (NOW - 50 * DAY, NOW - 50 * DAY))
|
|
self.assertEqual(result["seen"]["outputs"], 1)
|
|
|
|
def test_what_came_lately_is_not_judged(self):
|
|
self.models[1]["stat"] = stat(NOW - 5 * DAY, NOW - 400 * DAY) # copied here five days ago: keeps its old date
|
|
self.models[3]["stat"] = stat(NOW - 400 * DAY, NOW - 400 * DAY)
|
|
result = model_usage.check(now=NOW)
|
|
models = {Path(item["rel"]).stem: item for item in result["models"]}
|
|
self.assertNotIn("old", models)
|
|
self.assertEqual(models["never"]["added"], NOW - 400 * DAY)
|
|
|
|
def test_a_pack_can_be_marked_as_still_used(self):
|
|
self.assertEqual(model_usage.keep("ComfyUI-Easy-Use", True), ["ComfyUI-Easy-Use"])
|
|
for name in ("not-installed", "ComfyUI-Manager"): # only a listed pack ComfyUI loaded
|
|
with self.assertRaises(KeyError):
|
|
model_usage.keep(name, True)
|
|
packs = {item["name"]: item for item in model_usage.check(now=NOW)["packs"]}
|
|
self.assertEqual((packs["ComfyUI-Easy-Use"]["kept"], packs["forgotten-pack"]["kept"]), (True, False))
|
|
self.assertEqual(model_usage.keep("ComfyUI-Easy-Use", False), [])
|
|
self.assertEqual(model_usage.keep("gone-since", False), []) # unmarking any name is fine
|
|
|
|
class Request:
|
|
def __init__(self, body):
|
|
self.body = body
|
|
|
|
async def json(self):
|
|
return self.body
|
|
|
|
for body in ({"pack": "ComfyUI-Easy-Use"}, {"pack": "", "kept": True}, ["x"]):
|
|
self.assertEqual(asyncio.run(model_usage.api_keep(Request(body))).status, 400)
|
|
self.assertEqual(asyncio.run(model_usage.api_keep(Request({"pack": "nope", "kept": True}))).status, 404)
|
|
|
|
def test_reading_stops_at_the_limits_and_says_so(self):
|
|
for index in range(3):
|
|
output(self.out / f"{index}.png", NOW - DAY, {"1": ("KSampler", {})})
|
|
with mock.patch.object(model_usage, "MAX_IMAGES", 2):
|
|
result = model_usage.check(now=NOW)
|
|
self.assertEqual((result["seen"]["outputs"], result["seen"]["capped"]), (2, True))
|
|
|
|
def test_a_broken_or_huge_saved_file_is_skipped(self):
|
|
(self.flows / "broken.json").write_text("{not json", encoding="utf-8")
|
|
(self.flows / "ok.json").write_text(json.dumps({"x": "never.safetensors"}), encoding="utf-8")
|
|
with mock.patch.object(model_usage, "MAX_JSON_BYTES", 40):
|
|
(self.flows / "big.json").write_text(json.dumps({"x": "old.safetensors", "pad": "x" * 100}), encoding="utf-8")
|
|
result = model_usage.check(now=NOW)
|
|
models = {Path(item["rel"]).stem: item for item in result["models"]}
|
|
self.assertEqual((models["never"]["saved"], models["old"]["saved"]), (1, 0))
|
|
|
|
|
|
class LedgerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.path = os.path.join(self.temp.name, "usage_ledger.json")
|
|
patch = mock.patch.object(usage_ledger, "ledger_path", return_value=self.path)
|
|
patch.start()
|
|
self.addCleanup(patch.stop)
|
|
usage_ledger._ledger = None
|
|
self.addCleanup(setattr, usage_ledger, "_ledger", None)
|
|
|
|
def test_each_run_is_noted_and_saved(self):
|
|
prompt = {"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "SDXL\\Base.safetensors"}},
|
|
"2": {"class_type": "Power Lora Loader (rgthree)", "inputs": {"lora_1": {"on": True, "lora": "a.safetensors"}}},
|
|
"3": {"class_type": "KSampler", "inputs": {"seed": 5, "model": ["1", 0]}}}
|
|
usage_ledger.note(prompt, when=100.0)
|
|
usage_ledger.note(prompt, when=200.0)
|
|
usage_ledger._timer.cancel()
|
|
usage_ledger._save()
|
|
with open(self.path, encoding="utf-8") as f:
|
|
saved = json.load(f)
|
|
self.assertEqual(saved["started"], 100.0)
|
|
self.assertEqual(saved["types"]["KSampler"], [200.0, 2])
|
|
self.assertEqual(sorted(saved["models"]), ["a.safetensors", "sdxl/base.safetensors"])
|
|
usage_ledger._ledger = None # read back from the file
|
|
self.assertEqual(usage_ledger.read()["models"]["a.safetensors"], [200.0, 2])
|
|
|
|
def test_an_opened_workflow_is_noted_apart_from_runs(self):
|
|
class Request:
|
|
def __init__(self, body):
|
|
self.body = body
|
|
|
|
async def json(self):
|
|
return self.body
|
|
|
|
noted = asyncio.run(usage_ledger.api_opened(Request({"types": ["KSampler", 5, ""], "models": ["Loras\\A.safetensors", "notes.txt"]})))
|
|
usage_ledger._timer.cancel()
|
|
self.assertEqual(json.loads(noted.body)["noted"], 2)
|
|
ledger = usage_ledger.read()
|
|
self.assertEqual((sorted(ledger["opened"]["types"]), sorted(ledger["opened"]["models"])), (["KSampler"], ["loras/a.safetensors"]))
|
|
self.assertEqual((ledger["types"], ledger["models"]), ({}, {})) # not counted as a run
|
|
for body in ({"types": "KSampler", "models": []}, ["KSampler"], {"types": []}):
|
|
self.assertEqual(asyncio.run(usage_ledger.api_opened(Request(body))).status, 400)
|
|
|
|
def test_a_note_never_stops_a_run(self):
|
|
request = {"prompt": "not a prompt", "client_id": "x"}
|
|
self.assertIs(usage_ledger.note_prompt(request), request)
|
|
with mock.patch.object(usage_ledger, "note", side_effect=RuntimeError("boom")):
|
|
self.assertIs(usage_ledger.note_prompt(request), request)
|
|
self.assertIsNone(usage_ledger.read()["started"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|