ausboss
da60955801
The pack's server still never contacts a host. The info card's button has the browser ask civitai.com about one file by its SHA256, then the server saves the answer beside the LoRA as <model>.civitai.info after checking it names the same file (size, shape and hash). The old Python HTTP call stays gone. release_preflight now allows a raw fetch() only in js/cache_guard and js/lora_loader. SECURITY.md, AGENTS.md, the README and the LoRA Loader help page describe the change.
451 lines
19 KiB
Python
451 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes import _lora_helpers
|
|
from nodes._lora_helpers import base_model_family, patch_total
|
|
|
|
|
|
class StubPatcher:
|
|
def __init__(self, patches):
|
|
self.patches = patches
|
|
|
|
|
|
class StubClip:
|
|
def __init__(self, patches):
|
|
self.patcher = StubPatcher(patches)
|
|
|
|
|
|
class PatchTotalTests(unittest.TestCase):
|
|
def test_counts_patch_entries_not_keys(self):
|
|
self.assertEqual(patch_total(StubPatcher({"a": [1], "b": [1, 1]}), None), 3)
|
|
|
|
def test_sums_model_and_clip(self):
|
|
self.assertEqual(patch_total(StubPatcher({"a": [1]}), StubClip({"t": [1, 1]})), 3)
|
|
|
|
def test_empty_patchers_are_zero_not_none(self):
|
|
self.assertEqual(patch_total(StubPatcher({}), StubClip({})), 0)
|
|
|
|
def test_missing_clip_is_tolerated(self):
|
|
self.assertEqual(patch_total(StubPatcher({"a": [1]}), None), 1)
|
|
|
|
def test_unreadable_patcher_disables_the_check(self):
|
|
self.assertIsNone(patch_total(object(), None))
|
|
|
|
|
|
class BaseModelFamilyTests(unittest.TestCase):
|
|
def test_reads_declarative_keys(self):
|
|
self.assertEqual(base_model_family({"modelspec.architecture": "flux-1-dev/lora"}), "Flux")
|
|
self.assertEqual(
|
|
base_model_family({"modelspec.architecture": "stable-diffusion-xl-v1-base/lora"}),
|
|
"SDXL",
|
|
)
|
|
self.assertEqual(base_model_family({"ss_base_model_version": "sd_v1"}), "SD1.5")
|
|
|
|
def test_source_filename_never_implies_a_family(self):
|
|
# A filename ending _v1 or containing xl is not an architecture; mining
|
|
# it mislabeled working LoRAs as SD1.5 / SDXL.
|
|
self.assertEqual(base_model_family({"ss_sd_model_name": "krea2_turbo_v1.safetensors"}), "")
|
|
self.assertEqual(base_model_family({"ss_sd_model_name": "mymodel_xl_fix.safetensors"}), "")
|
|
|
|
def test_unknown_family_reports_what_the_file_declares(self):
|
|
self.assertEqual(base_model_family({"modelspec.architecture": "krea2/lora"}), "krea2")
|
|
self.assertEqual(base_model_family({"ss_base_model_version": "krea2"}), "krea2")
|
|
|
|
def test_no_metadata_is_empty(self):
|
|
self.assertEqual(base_model_family({}), "")
|
|
|
|
|
|
class NoEffectWarningTests(unittest.TestCase):
|
|
def setUp(self):
|
|
_lora_helpers._no_effect_warned.clear()
|
|
self.addCleanup(_lora_helpers._no_effect_warned.clear)
|
|
|
|
def _warn(self, name, family=""):
|
|
buffer = io.StringIO()
|
|
with contextlib.redirect_stdout(buffer):
|
|
_lora_helpers._warn_no_effect(name, family)
|
|
return buffer.getvalue()
|
|
|
|
def test_names_the_lora_and_stays_ascii(self):
|
|
message = self._warn("anime_outline.safetensors")
|
|
self.assertIn("anime_outline.safetensors", message)
|
|
self.assertIn("matched nothing", message)
|
|
message.encode("ascii") # Windows cp1252 consoles
|
|
|
|
def test_includes_declared_base_model_when_known(self):
|
|
self.assertIn("'krea2'", self._warn("x.safetensors", "krea2"))
|
|
|
|
def test_warns_once_per_lora(self):
|
|
self.assertNotEqual(self._warn("x.safetensors"), "")
|
|
self.assertEqual(self._warn("x.safetensors"), "")
|
|
|
|
|
|
class _FolderPaths:
|
|
def __init__(self, lora_root: Path, user_root: Path):
|
|
self.lora_root = lora_root
|
|
self.user_root = user_root
|
|
|
|
def get_full_path(self, kind: str, name: str) -> str | None:
|
|
if kind != "loras":
|
|
return None
|
|
path = self.lora_root / name
|
|
return str(path) if path.is_file() else None
|
|
|
|
def get_folder_paths(self, kind: str) -> list[str]:
|
|
return [str(self.lora_root)] if kind == "loras" else []
|
|
|
|
def get_user_directory(self) -> str:
|
|
return str(self.user_root)
|
|
|
|
|
|
class LoraCivitaiSidecarTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self._tmp = tempfile.TemporaryDirectory()
|
|
base = Path(self._tmp.name)
|
|
self.lora_root = base / "models" / "loras"
|
|
self.user_root = base / "user"
|
|
self.lora = self.lora_root / "Krea 2" / "candid.safetensors"
|
|
self.lora.parent.mkdir(parents=True)
|
|
self.lora.write_bytes(b"not a real safetensors file")
|
|
self.folder_paths = _FolderPaths(self.lora_root, self.user_root)
|
|
self.folder_paths_patch = patch.object(_lora_helpers, "folder_paths", self.folder_paths)
|
|
self.folder_paths_patch.start()
|
|
|
|
def tearDown(self):
|
|
self.folder_paths_patch.stop()
|
|
self._tmp.cleanup()
|
|
|
|
def test_reads_a_standard_sidecar_created_by_another_comfyui_tool(self):
|
|
payload = {
|
|
"id": "456",
|
|
"modelId": "123",
|
|
"baseModel": "Krea 2",
|
|
"trainedWords": [" candid style ", "film grain"],
|
|
"model": {"name": "Candid Slider"},
|
|
}
|
|
self.lora.with_suffix(".civitai.info").write_text(
|
|
json.dumps(payload), encoding="utf-8"
|
|
)
|
|
|
|
cached = _lora_helpers.load_civitai_cache("Krea 2/candid.safetensors")
|
|
info = _lora_helpers.lora_info("Krea 2/candid.safetensors")
|
|
|
|
self.assertEqual(cached["title"], "Candid Slider")
|
|
self.assertEqual(cached["trained_words"], ["candid style", "film grain"])
|
|
self.assertEqual(cached["model_id"], 123)
|
|
self.assertEqual(cached["version_id"], 456)
|
|
self.assertTrue(info["has_civitai"])
|
|
self.assertEqual(info["civitai_triggers"], ["candid style", "film grain"])
|
|
self.assertEqual(info["civitai_model_id"], 123)
|
|
self.assertEqual(info["civitai_version_id"], 456)
|
|
|
|
def test_lora_info_reports_file_facts(self):
|
|
info = _lora_helpers.lora_info("Krea 2/candid.safetensors")
|
|
|
|
stat = self.lora.stat()
|
|
self.assertEqual(info["size_bytes"], stat.st_size)
|
|
self.assertEqual(info["mtime"], stat.st_mtime)
|
|
|
|
def test_invalid_sidecar_is_ignored(self):
|
|
self.lora.with_suffix(".civitai.info").write_text("{}", encoding="utf-8")
|
|
|
|
self.assertEqual(
|
|
_lora_helpers.load_civitai_cache("Krea 2/candid.safetensors"), {}
|
|
)
|
|
self.assertFalse(_lora_helpers.lora_info("Krea 2/candid.safetensors")["has_civitai"])
|
|
|
|
def _civitai_answer(self, sha=None, **extra):
|
|
import hashlib
|
|
|
|
sha = sha or hashlib.sha256(self.lora.read_bytes()).hexdigest().upper()
|
|
return {
|
|
"id": 456,
|
|
"modelId": 123,
|
|
"baseModel": "Krea 2",
|
|
"trainedWords": ["candid style"],
|
|
"model": {"name": "Candid Slider"},
|
|
"files": [{"name": "candid.safetensors", "hashes": {"SHA256": sha}}],
|
|
**extra,
|
|
}
|
|
|
|
def test_hash_route_gives_the_files_sha256(self):
|
|
import hashlib
|
|
|
|
expected = hashlib.sha256(self.lora.read_bytes()).hexdigest()
|
|
self.assertEqual(_lora_helpers.lora_sha256("Krea 2/candid.safetensors"), expected)
|
|
# Second call is served from the (mtime, size) cache in the user folder.
|
|
self.assertEqual(_lora_helpers.lora_sha256("Krea 2/candid.safetensors"), expected)
|
|
self.assertTrue((self.user_root / "ausboss" / "lora_hashes.json").is_file())
|
|
|
|
def test_saves_an_answer_that_describes_this_file(self):
|
|
info = _lora_helpers.save_civitai_sidecar("Krea 2/candid.safetensors", self._civitai_answer())
|
|
|
|
self.assertEqual(info["title"], "Candid Slider")
|
|
sidecar = self.lora.with_suffix(".civitai.info")
|
|
self.assertEqual(json.loads(sidecar.read_text(encoding="utf-8"))["modelId"], 123)
|
|
self.assertTrue(_lora_helpers.lora_info("Krea 2/candid.safetensors")["has_civitai"])
|
|
self.assertEqual(list(self.lora.parent.glob("*.tmp*")), [])
|
|
|
|
def test_refuses_an_answer_for_a_different_file(self):
|
|
with self.assertRaises(ValueError):
|
|
_lora_helpers.save_civitai_sidecar(
|
|
"Krea 2/candid.safetensors", self._civitai_answer(sha="0" * 64)
|
|
)
|
|
self.assertFalse(self.lora.with_suffix(".civitai.info").exists())
|
|
|
|
def test_refuses_answers_in_the_wrong_shape_or_size(self):
|
|
for bad in ([1], "text", None, {}, {"files": "x"}):
|
|
with self.assertRaises(ValueError):
|
|
_lora_helpers.save_civitai_sidecar("Krea 2/candid.safetensors", bad)
|
|
huge = self._civitai_answer(padding="x" * (_lora_helpers.MAX_CIVITAI_BYTES + 1))
|
|
with self.assertRaises(ValueError):
|
|
_lora_helpers.save_civitai_sidecar("Krea 2/candid.safetensors", huge)
|
|
self.assertFalse(self.lora.with_suffix(".civitai.info").exists())
|
|
|
|
def test_cannot_write_outside_the_lora_folders(self):
|
|
outside = Path(self._tmp.name) / "secret.safetensors"
|
|
outside.write_bytes(b"x")
|
|
for name in ("../../secret.safetensors", str(outside), "Krea 2/missing.safetensors"):
|
|
with self.assertRaises(Exception):
|
|
_lora_helpers.save_civitai_sidecar(name, self._civitai_answer())
|
|
self.assertFalse(outside.with_suffix(".civitai.info").exists())
|
|
|
|
|
|
|
|
|
|
class TestMatchLoraName(unittest.TestCase):
|
|
AVAILABLE = [
|
|
"krea2/Beauty.safetensors",
|
|
"styles/Beauty.safetensors",
|
|
"klein/outpaint_v1.safetensors",
|
|
]
|
|
|
|
def test_exact_wins(self):
|
|
self.assertEqual(
|
|
_lora_helpers.match_lora_name("klein/outpaint_v1.safetensors", self.AVAILABLE),
|
|
("klein/outpaint_v1.safetensors", "exact"),
|
|
)
|
|
|
|
def test_case_insensitive_full_path(self):
|
|
self.assertEqual(
|
|
_lora_helpers.match_lora_name("KLEIN/Outpaint_V1.safetensors", self.AVAILABLE),
|
|
("klein/outpaint_v1.safetensors", "remapped"),
|
|
)
|
|
|
|
def test_moved_file_resolves_by_unique_basename(self):
|
|
# The whole point: a file reorganized into another folder (or a
|
|
# workflow from another machine's layout) lands on the same file.
|
|
self.assertEqual(
|
|
_lora_helpers.match_lora_name("old\\path\\outpaint_v1.safetensors", self.AVAILABLE),
|
|
("klein/outpaint_v1.safetensors", "remapped"),
|
|
)
|
|
|
|
def test_extension_ignored_in_basename_match(self):
|
|
self.assertEqual(
|
|
_lora_helpers.match_lora_name("outpaint_v1.ckpt", self.AVAILABLE),
|
|
("klein/outpaint_v1.safetensors", "remapped"),
|
|
)
|
|
|
|
def test_ambiguous_basename_stays_put(self):
|
|
name, status = _lora_helpers.match_lora_name(
|
|
"elsewhere/Beauty.safetensors", self.AVAILABLE
|
|
)
|
|
self.assertEqual((name, status), ("elsewhere/Beauty.safetensors", "ambiguous"))
|
|
|
|
def test_missing_is_reported(self):
|
|
self.assertEqual(
|
|
_lora_helpers.match_lora_name("gone.safetensors", self.AVAILABLE)[1],
|
|
"missing",
|
|
)
|
|
|
|
|
|
class MissingRowPolicyTests(unittest.TestCase):
|
|
"""on_missing: error (the default) stops before sampling, skip warns and keeps going."""
|
|
|
|
ROW = {"name": "gone.safetensors", "strength": 1.0, "strength_clip": 1.0,
|
|
"enabled": True, "triggers": ""}
|
|
|
|
def setUp(self):
|
|
_lora_helpers._missing_warned.clear()
|
|
self.addCleanup(_lora_helpers._missing_warned.clear)
|
|
# apply_lora_stack imports comfy.sd lazily; the missing-file branch runs
|
|
# before any of it is used, so an empty stand-in module is enough.
|
|
import types
|
|
|
|
fake_comfy = types.ModuleType("comfy")
|
|
fake_sd = types.ModuleType("comfy.sd")
|
|
fake_comfy.sd = fake_sd
|
|
self.modules = patch.dict(sys.modules, {"comfy": fake_comfy, "comfy.sd": fake_sd})
|
|
self.modules.start()
|
|
self.addCleanup(self.modules.stop)
|
|
self.gone = patch.object(
|
|
_lora_helpers, "resolve_lora_path",
|
|
side_effect=ValueError("LoRA file not found in models/loras: gone.safetensors"),
|
|
)
|
|
self.gone.start()
|
|
self.addCleanup(self.gone.stop)
|
|
|
|
def test_missing_rows_ignores_disabled_and_zero_rows(self):
|
|
rows = [
|
|
dict(self.ROW),
|
|
{**self.ROW, "name": "off.safetensors", "enabled": False},
|
|
{**self.ROW, "name": "parked.safetensors", "strength": 0, "strength_clip": 0},
|
|
]
|
|
self.assertEqual(
|
|
[name for name, _ in _lora_helpers.missing_lora_rows(rows)],
|
|
["gone.safetensors"],
|
|
)
|
|
|
|
def test_skip_warns_once_and_keeps_going(self):
|
|
buffer = io.StringIO()
|
|
with contextlib.redirect_stdout(buffer):
|
|
model, clip = _lora_helpers.apply_lora_stack(
|
|
"model", None, [dict(self.ROW)], on_missing="skip"
|
|
)
|
|
_lora_helpers.apply_lora_stack("model", None, [dict(self.ROW)], on_missing="skip")
|
|
self.assertEqual((model, clip), ("model", None))
|
|
self.assertEqual(buffer.getvalue().count("skipping 'gone.safetensors'"), 1)
|
|
|
|
def test_skip_mode_reports_the_rows_it_left_out(self):
|
|
missing = []
|
|
rows = [
|
|
dict(self.ROW),
|
|
{**self.ROW, "name": "parked.safetensors", "strength": 0, "strength_clip": 0},
|
|
]
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
_lora_helpers.apply_lora_stack("model", None, rows, on_missing="skip", missing=missing)
|
|
self.assertEqual([row["name"] for row in missing], ["gone.safetensors"])
|
|
|
|
def test_a_row_skipped_as_missing_adds_no_trigger_words(self):
|
|
# A LoRA that never loaded must not put its words in the prompt. A
|
|
# row parked at strength 0 keeps its words, as the node docs promise,
|
|
# so comparing with and without it changes only the weights.
|
|
try:
|
|
from nodes.node_lora_loader import AusBossLoraLoader
|
|
except Exception as exc: # pragma: no cover - needs the package importable
|
|
self.skipTest(f"node module not importable offline: {exc}")
|
|
|
|
def resolve(name):
|
|
if name == "gone.safetensors":
|
|
raise ValueError("LoRA file not found in models/loras: gone.safetensors")
|
|
return Path(name)
|
|
|
|
rows = [
|
|
{**self.ROW, "name": "kept.safetensors", "triggers": "kept word"},
|
|
{**self.ROW, "triggers": "ghost word"},
|
|
{**self.ROW, "name": "parked.safetensors", "strength": 0, "strength_clip": 0,
|
|
"triggers": "parked word"},
|
|
{**self.ROW, "name": "off.safetensors", "enabled": False, "triggers": "off word"},
|
|
]
|
|
sys.modules["comfy.sd"].load_lora_for_models = lambda model, clip, *_args: (model, clip)
|
|
with (
|
|
patch.object(_lora_helpers, "resolve_lora_path", side_effect=resolve),
|
|
patch.object(_lora_helpers, "_load_lora_file", return_value={}),
|
|
contextlib.redirect_stdout(io.StringIO()),
|
|
):
|
|
_model, _clip, triggers = AusBossLoraLoader().load_loras(
|
|
"model", json.dumps(rows), on_missing="skip"
|
|
)
|
|
self.assertEqual(triggers, "kept word, parked word")
|
|
|
|
def test_the_on_missing_choices_come_from_the_helper(self):
|
|
try:
|
|
from nodes.node_lora_loader import AusBossLoraLoader
|
|
except Exception as exc: # pragma: no cover - needs the package importable
|
|
self.skipTest(f"node module not importable offline: {exc}")
|
|
choices = AusBossLoraLoader.INPUT_TYPES()["optional"]["on_missing"][0]
|
|
self.assertEqual(choices, list(_lora_helpers.MISSING_MODES))
|
|
self.assertEqual(choices, ["skip", "error"]) # saved graphs store these
|
|
|
|
def test_error_is_the_default_and_names_the_file_and_the_switch(self):
|
|
for kwargs in ({}, {"on_missing": "error"}):
|
|
with self.assertRaises(ValueError) as caught:
|
|
_lora_helpers.apply_lora_stack("model", None, [dict(self.ROW)], **kwargs)
|
|
message = str(caught.exception)
|
|
self.assertIn("gone.safetensors", message)
|
|
self.assertIn("Stop on missing LoRA", message)
|
|
self.assertIn("settings menu", message)
|
|
self.assertIn("on_missing: skip", message)
|
|
message.encode("ascii") # Windows cp1252 consoles
|
|
|
|
def test_node_validation_blocks_by_default_and_passes_in_skip_mode(self):
|
|
try:
|
|
from nodes.node_lora_loader import AusBossLoraLoader
|
|
except Exception as exc: # pragma: no cover - needs the package importable
|
|
self.skipTest(f"node module not importable offline: {exc}")
|
|
stack = json.dumps([dict(self.ROW)])
|
|
self.assertIs(AusBossLoraLoader.VALIDATE_INPUTS(stack, on_missing="skip"), True)
|
|
for kwargs in ({}, {"on_missing": "error"}):
|
|
verdict = AusBossLoraLoader.VALIDATE_INPUTS(stack, **kwargs)
|
|
self.assertIsInstance(verdict, str)
|
|
self.assertIn("gone.safetensors", verdict)
|
|
self.assertIn("settings menu", verdict)
|
|
verdict.encode("ascii")
|
|
# The input's declared default matches the code paths above.
|
|
on_missing = AusBossLoraLoader.INPUT_TYPES()["optional"]["on_missing"]
|
|
self.assertEqual(on_missing[1]["default"], "error")
|
|
|
|
|
|
class ResolveLoraPathTests(unittest.TestCase):
|
|
"""The file a name resolves to, and the check that keeps it in models/loras."""
|
|
|
|
def setUp(self):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(tmp.cleanup)
|
|
base = Path(tmp.name)
|
|
self.root = base / "loras"
|
|
self.root.mkdir()
|
|
self.outside = base / "other_drive"
|
|
self.outside.mkdir()
|
|
(self.root / "plain.safetensors").write_bytes(b"x")
|
|
(self.outside / "real.safetensors").write_bytes(b"x")
|
|
# A LoRA kept on another drive and linked into models/loras, which
|
|
# ComfyUI's own loaders accept.
|
|
(self.root / "linked.safetensors").symlink_to(self.outside / "real.safetensors")
|
|
self.full_path = lambda kind, name: (
|
|
str(self.root / name) if (self.root / name).is_file() else None
|
|
)
|
|
fake = type("FakeFolderPaths", (), {})()
|
|
fake.get_full_path = lambda kind, name: self.full_path(kind, name)
|
|
fake.get_folder_paths = lambda kind: [str(self.root)]
|
|
fake.get_filename_list = lambda kind: ["linked.safetensors", "plain.safetensors"]
|
|
patcher = patch.object(_lora_helpers, "folder_paths", fake)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def test_plain_file_resolves(self):
|
|
path = _lora_helpers.resolve_lora_path("plain.safetensors")
|
|
self.assertEqual(path.name, "plain.safetensors")
|
|
self.assertTrue(path.is_file())
|
|
|
|
def test_symlinked_file_resolves_inside_the_folder(self):
|
|
path = _lora_helpers.resolve_lora_path("linked.safetensors")
|
|
self.assertEqual(path.parent, self.root)
|
|
self.assertTrue(path.is_file())
|
|
|
|
def test_symlinked_file_is_not_reported_missing(self):
|
|
rows = [{"name": "linked.safetensors", "strength": 1.0, "strength_clip": 1.0, "enabled": True}]
|
|
self.assertEqual(_lora_helpers.missing_lora_rows(rows), [])
|
|
|
|
def test_path_outside_the_folders_is_refused(self):
|
|
self.full_path = lambda kind, name: str(self.root / ".." / "other_drive" / "real.safetensors")
|
|
with self.assertRaisesRegex(ValueError, "escapes the loras folders"):
|
|
_lora_helpers.resolve_lora_path("../other_drive/real.safetensors")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|