ausboss
c8383da808
resolve_lora_path followed links before checking that the file sits in models/loras, so a LoRA kept on another drive and linked into the folder was refused as escaping it and then reported missing, which stops the run when the node is set to stop on a missing LoRA. Containment is now checked on the path as written, with '..' folded but links not followed, the way ComfyUI's own loaders treat linked files.
395 lines
16 KiB
Python
395 lines
16 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"])
|
|
|
|
|
|
|
|
|
|
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()
|