613 lines
26 KiB
Python
613 lines
26 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import importlib.util
|
|
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())
|
|
|
|
|
|
def _picture(width=1200, height=800, fmt="PNG", mode="RGB") -> bytes:
|
|
from PIL import Image
|
|
|
|
out = io.BytesIO()
|
|
Image.new(mode, (width, height)).save(out, fmt)
|
|
return out.getvalue()
|
|
|
|
|
|
@unittest.skipIf(importlib.util.find_spec("PIL") is None, "needs Pillow (run with ComfyUI's Python)")
|
|
class LoraCivitaiPictureTests(unittest.TestCase):
|
|
name = "Krea 2/candid.safetensors"
|
|
|
|
def setUp(self):
|
|
LoraCivitaiSidecarTests.setUp(self)
|
|
self.target = self.lora.parent / "candid.preview.jpg"
|
|
|
|
def tearDown(self):
|
|
LoraCivitaiSidecarTests.tearDown(self)
|
|
|
|
def test_saves_a_small_jpeg_the_hover_preview_finds(self):
|
|
from PIL import Image
|
|
|
|
saved = _lora_helpers.save_civitai_picture(self.name, _picture(fmt="WEBP", mode="RGBA"))
|
|
self.assertEqual(saved, self.target)
|
|
with Image.open(saved) as picture:
|
|
self.assertEqual(picture.format, "JPEG")
|
|
self.assertEqual(max(picture.size), _lora_helpers.PICTURE_EDGE)
|
|
self.assertFalse(picture.info.get("exif"))
|
|
self.assertEqual(_lora_helpers.find_thumbnail(self.name), self.target)
|
|
self.assertTrue(_lora_helpers.lora_info(self.name)["has_preview"])
|
|
|
|
def test_never_replaces_a_picture_that_is_there(self):
|
|
own = self.lora.parent / "candid.png"
|
|
own.write_bytes(_picture(10, 10))
|
|
with self.assertRaises(ValueError):
|
|
_lora_helpers.save_civitai_picture(self.name, _picture())
|
|
self.assertFalse(self.target.exists())
|
|
self.assertEqual(own.read_bytes(), _picture(10, 10))
|
|
|
|
def test_refuses_what_is_not_a_picture_or_too_large(self):
|
|
for data in (b"", b"not a picture", b"\x89PNG\r\n\x1a\n broken", None):
|
|
with self.assertRaises(ValueError):
|
|
_lora_helpers.save_civitai_picture(self.name, data)
|
|
with self.assertRaises(ValueError):
|
|
_lora_helpers.save_civitai_picture(self.name, b"x" * (_lora_helpers.MAX_PICTURE_BYTES + 1))
|
|
with self.assertRaises(ValueError):
|
|
_lora_helpers.save_civitai_picture(self.name, _picture(8000, 6000, mode="L"))
|
|
self.assertFalse(self.target.exists())
|
|
|
|
def test_only_next_to_a_listed_lora(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_picture(name, _picture())
|
|
self.assertFalse((Path(self._tmp.name) / "secret.preview.jpg").exists())
|
|
|
|
|
|
class ParseStackTests(unittest.TestCase):
|
|
def test_rows_in_a_switched_off_group_do_not_run(self):
|
|
rows = _lora_helpers.parse_lora_stack(json.dumps([
|
|
{"name": "a.safetensors", "enabled": True},
|
|
{"name": "b.safetensors", "enabled": True, "group": "Look", "group_off": True},
|
|
{"name": "c.safetensors", "enabled": True, "group": "Look", "group_off": "yes"},
|
|
]))
|
|
self.assertEqual([row["enabled"] for row in rows], [True, False, True])
|
|
self.assertNotIn("group_off", rows[1])
|
|
|
|
|
|
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 UnreadableLoraFileTests(unittest.TestCase):
|
|
"""A file in a row that is not a LoRA: named, explained, and caught before a model loads."""
|
|
|
|
ROW = {"name": "import_lora_f16.ckpt", "strength": 1.0, "strength_clip": 1.0,
|
|
"enabled": True, "triggers": ""}
|
|
SQLITE = b"SQLite format 3\x00" + bytes(84)
|
|
|
|
def setUp(self):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(tmp.cleanup)
|
|
self.root = Path(tmp.name)
|
|
self.resolved = patch.object(
|
|
_lora_helpers, "resolve_lora_path", side_effect=lambda name: self.root / name
|
|
)
|
|
self.resolved.start()
|
|
self.addCleanup(self.resolved.stop)
|
|
|
|
def write(self, name: str, data: bytes) -> Path:
|
|
path = self.root / name
|
|
path.write_bytes(data)
|
|
return path
|
|
|
|
def test_each_wrong_file_gets_its_own_plain_reason(self):
|
|
cases = {
|
|
"Draw Things": self.SQLITE,
|
|
"web page": b"\n<!DOCTYPE html><html><head><title>Sign in</title>",
|
|
"text pointer": b"version https://git-lfs.github.com/spec/v1\noid sha256:0\n",
|
|
"is empty": b"",
|
|
}
|
|
for expected, data in cases.items():
|
|
reason = _lora_helpers.unreadable_lora_reason(self.write("case.ckpt", data))
|
|
self.assertIn(expected, reason)
|
|
reason.encode("ascii") # Windows cp1252 consoles
|
|
|
|
def test_a_real_safetensors_start_is_left_alone(self):
|
|
# The file opens with its header length, so its first byte can be any
|
|
# byte at all, '<' and 'S' included.
|
|
for first in (0x20, 0x3C, 0x53, 0x7B):
|
|
head = bytes([first, 0x22]) + bytes(6) + b'{"__metadata__":{"format":"pt"}}'
|
|
path = self.write("real.safetensors", head)
|
|
self.assertEqual(_lora_helpers.unreadable_lora_reason(path), "")
|
|
|
|
def test_only_rows_that_would_load_are_checked(self):
|
|
self.write("import_lora_f16.ckpt", self.SQLITE)
|
|
self.write("off.ckpt", self.SQLITE)
|
|
self.write("parked.ckpt", self.SQLITE)
|
|
rows = [
|
|
dict(self.ROW),
|
|
{**self.ROW, "name": "off.ckpt", "enabled": False},
|
|
{**self.ROW, "name": "parked.ckpt", "strength": 0, "strength_clip": 0},
|
|
{**self.ROW, "name": "not_on_disk.safetensors"},
|
|
]
|
|
self.assertEqual(
|
|
[name for name, _ in _lora_helpers.unreadable_lora_rows(rows)],
|
|
["import_lora_f16.ckpt"],
|
|
)
|
|
|
|
def test_validation_names_the_file_before_anything_loads(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}")
|
|
self.write("import_lora_f16.ckpt", self.SQLITE)
|
|
stack = json.dumps([dict(self.ROW)])
|
|
for mode in ("error", "skip"):
|
|
verdict = AusBossLoraLoader.VALIDATE_INPUTS(stack, on_missing=mode)
|
|
self.assertIsInstance(verdict, str)
|
|
self.assertIn("'import_lora_f16.ckpt'", verdict)
|
|
self.assertIn("Draw Things", verdict)
|
|
self.assertIn(".safetensors", verdict)
|
|
verdict.encode("ascii")
|
|
|
|
def test_a_file_that_fails_to_load_is_named_with_the_loader_error_last(self):
|
|
import types
|
|
|
|
fake_comfy = types.ModuleType("comfy")
|
|
fake_sd = types.ModuleType("comfy.sd")
|
|
fake_comfy.sd = fake_sd
|
|
modules = patch.dict(sys.modules, {"comfy": fake_comfy, "comfy.sd": fake_sd})
|
|
modules.start()
|
|
self.addCleanup(modules.stop)
|
|
self.write("broken.safetensors", bytes(64))
|
|
row = {**self.ROW, "name": "broken.safetensors"}
|
|
failure = RuntimeError("Weights only load failed.\n\nUnsupported operand 83")
|
|
with patch.object(_lora_helpers, "_load_lora_file", side_effect=failure):
|
|
with self.assertRaises(ValueError) as caught:
|
|
_lora_helpers.apply_lora_stack("model", None, [row])
|
|
message = str(caught.exception)
|
|
self.assertTrue(message.startswith("LoRA Loader: 'broken.safetensors' is not a LoRA file"))
|
|
self.assertIn("Download the file again", message)
|
|
self.assertTrue(message.endswith("RuntimeError: Weights only load failed. Unsupported operand 83"))
|
|
self.assertIs(caught.exception.__cause__, failure)
|
|
|
|
|
|
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()
|