Files
BlackVortexAI-bv_nodepack/tests/test_krea_preparation_cache.py
BlackVortexAIandClaude Fable 5.1 95ace90041 Add BV Text writing assistance, preparation caches and the 1.4.4 hardening
Features: BV Text node with management-gated writing assistance (translate,
improve, undo, prompt profiles); reuse of prepared MODEL/CLIP state for Krea 2
and Anima LLLite consumers; Regional refactoring.

Security boundaries: one wall-clock budget per Remote LLM request enforced by
a socket-closing watchdog across connect, TLS handshake, headers, chunk framing
and body; writing assistance uses the profile's configured timeout; a shared
bounded JSON reader with route-specific limits on every JSON route; endpoint
binding negative tests for the assist path.

Release and packaging: SECURITY.md; bundled-library notices and a readable
bundle header; Registry scan replica as a hard test with a pinned count; a
packaged-file-set test; publish only from a validated push to main with
provenance checks, pinned comfy-cli, an input manifest and a byte-exact
comparison of the stored archive; npm audit and a whole js/ rebuild gate in
Validate; design references, brand sources and a type declaration excluded
from the package. README and changelog for 1.4.4.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-13 16:51:38 +02:00

221 lines
10 KiB
Python

"""Preparation cache regression: replay changed downstream work with stable inputs."""
import copy
import ast
import os
import sys
import tempfile
import unittest
import uuid
from pathlib import Path
from unittest.mock import Mock
from test_basis_model_patcher import AutomaticLoraServiceTests, Model, Clip, subject
from util.lora_preparation_cache import LoraPreparationCache
class PreparationCacheTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.path = Path(self.temp.name) / "weights.safetensors"
self.path.write_bytes(b"fixture")
self.entries = [(str(self.path), 1., .5)]
self.model, self.clip = Model(), Clip()
self.model.model_options = {"transformer_options": {}}
self.model.patches_uuid = uuid.uuid4()
self.cache = LoraPreparationCache()
def prepare(model, clip, entries):
model, clip = model.clone(), clip.clone()
model.patches_uuid = uuid.uuid4()
return model, clip
self.prepare = Mock(side_effect=prepare)
def apply(self):
return self.cache.apply(self.model, self.clip, self.entries, self.prepare)
def test_repeated_preparation_preserves_patch_uuid_and_isolates_outputs(self):
first, _ = self.apply()
identifier = first.patches_uuid
first.model_options["transformer_options"]["attention"] = "downstream"
first.set_attachments("downstream", {})
second, _ = self.apply()
self.assertEqual(self.prepare.call_count, 1)
self.assertEqual(second.patches_uuid, identifier)
self.assertEqual(second.model_options, self.model.model_options)
self.assertEqual(second.attachments, {})
self.assertIsNot(first, second)
def test_invalidation_and_single_entry(self):
self.apply()
mutations = [
lambda: setattr(self.model, "patches_uuid", uuid.uuid4()),
lambda: self.model.model_options.update(extra=True),
lambda: setattr(self.clip, "layer_idx", -2),
lambda: self.path.write_bytes(b"new weights"),
lambda: self.entries.append((str(self.path), .25, -.5)),
lambda: self.entries.reverse(),
lambda: self.entries.pop(),
lambda: setattr(self, "model", self.model.clone()),
lambda: setattr(self, "clip", self.clip.clone()),
]
for count, mutate in enumerate(mutations, 2):
mutate()
self.apply()
self.assertEqual(self.prepare.call_count, count)
other = LoraPreparationCache()
other.apply(self.model, self.clip, self.entries, self.prepare)
self.assertEqual(self.prepare.call_count, len(mutations) + 2)
def test_opaque_extension_bypasses_and_failure_is_not_cached(self):
self.model.attachments["foreign"] = object()
self.apply(); self.apply()
self.assertEqual(self.prepare.call_count, 2)
self.model.attachments.clear()
self.prepare.side_effect = RuntimeError("failed")
with self.assertRaises(RuntimeError): self.apply()
self.assertIsNone(self.cache._entry)
def test_missing_file_does_not_hit_old_preparation(self):
self.apply()
self.path.unlink()
self.apply()
self.assertEqual(self.prepare.call_count, 2)
self.assertIsNone(self.cache._entry)
def test_strength_change_and_return_to_previous_entry(self):
original = list(self.entries)
self.apply()
self.entries = [(str(self.path), .5, .5)]
self.apply()
self.entries = original
self.apply()
self.assertEqual(self.prepare.call_count, 3)
def test_cyclic_options_bypass(self):
self.model.model_options["cycle"] = self.model.model_options
self.apply()
self.apply()
self.assertEqual(self.prepare.call_count, 2)
def test_clip_immutable_object_patches_hit_and_invalidate(self):
import torch
self.clip.patcher.object_patches = {"manual_cast_dtype": torch.float32}
self.apply(); self.apply()
self.assertEqual(self.prepare.call_count, 1)
self.clip.patcher.object_patches["manual_cast_dtype"] = torch.bfloat16
self.apply()
self.assertEqual(self.prepare.call_count, 2)
self.clip.patcher.object_patches = {"flag": True}
self.apply()
self.clip.patcher.object_patches = {"flag": 1}
self.apply()
self.assertEqual(self.prepare.call_count, 4)
def test_mutable_object_patches_still_bypass(self):
import torch
for value in ({}, [], object(), torch.zeros(1), lambda: None):
self.clip.patcher.object_patches = {"extension": value}
self.apply()
self.assertIsNone(self.cache._entry)
class CachedAutomaticTests(AutomaticLoraServiceTests):
def test_occurrences_validated_on_hits_and_empty_clears(self):
cache = LoraPreparationCache()
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "lora"
path.write_bytes(b"fixture")
resources = self.resources()
resources[0]["stack"] = [(str(path), -.5, .25), (str(path), 1., .25)]
def run(model=None, data=None):
return subject.apply_automatic_resources(model or self.model, self.clip,
resources if data is None else data, preparation_cache=cache)
first, _ = run()
second, _ = run()
self.assertEqual(self.apply.call_count, 1)
self.assertEqual(len(self.apply.call_args.args[2]), 2)
self.assertEqual(first.attachments, second.attachments)
changed = copy.deepcopy(resources)
changed[0]["stack"][0] = (str(path), .7, .25)
with self.assertRaises(ValueError): run(first, changed)
run(data=[])
self.assertIsNone(cache._entry)
class ConsumerPreparationTests(unittest.TestCase):
def test_all_consumers_share_helper_but_not_cache_instances(self):
source = Path(__file__).parents[1] / "py/nodes/bv_regional.py"
tree = ast.parse(source.read_text(encoding="utf-8"))
names = {"BVRegionalNativeConditioningNode", "BVRegionalSDXLAttentionNode",
"BVRegionalZImageAttentionNode", "BVRegionalFlux2KleinAttentionNode",
"BVRegionalKrea2AttentionNode", "BVRegionalAnimaConditioningNode"}
base = next(node for node in tree.body if isinstance(node, ast.ClassDef)
and node.name == "_GlobalLoraPreparation")
# Execute the real shared helper without importing GPU backend modules.
base = copy.deepcopy(base)
constructor = next(node for node in base.body if getattr(node, "name", None) == "__init__")
constructor.body = [node for node in constructor.body if not isinstance(node, ast.ImportFrom)]
apply = Mock(return_value=("model", "clip"))
namespace = {"LoraPreparationCache": LoraPreparationCache, "apply_global_patches": apply}
exec(compile(ast.fix_missing_locations(ast.Module(body=[base], type_ignores=[])), str(source), "exec"), namespace)
caches = []
for node in tree.body:
if not isinstance(node, ast.ClassDef) or node.name not in names:
continue
self.assertEqual([item.id for item in node.bases], ["_GlobalLoraPreparation"])
self.assertFalse(any(getattr(item, "name", None) == "__init__" for item in node.body))
calls = [item for item in ast.walk(node) if isinstance(item, ast.Call)
and isinstance(item.func, ast.Attribute) and item.func.attr == "_prepare_global_loras"]
self.assertEqual(len(calls), 1, node.name)
instance = type(node.name, (namespace["_GlobalLoraPreparation"],), {})()
instance._prepare_global_loras(None, "clip", "regional")
cache = apply.call_args.kwargs["preparation_cache"]
instance._prepare_global_loras(None, "clip", "changed regional")
self.assertIs(apply.call_args.kwargs["preparation_cache"], cache)
caches.append(cache)
self.assertEqual(len({id(cache) for cache in caches}), len(names))
@unittest.skipUnless(os.environ.get("BV_NATIVE_COMFY_ROOT"), "requires native CPU process")
class NativePreparationTests(unittest.TestCase):
def test_native_patch_identity_and_clone_isolation(self):
sys.path.insert(0, os.environ["BV_NATIVE_COMFY_ROOT"])
import comfy.options
comfy.options.enable_args_parsing()
sys.argv = [sys.argv[0], "--cpu"]
import torch
import comfy.model_patcher
base = comfy.model_patcher.ModelPatcher(
torch.nn.Linear(2, 2), torch.device("cpu"), torch.device("cpu"))
base.set_model_compute_dtype(torch.float32)
def prepare(model, clip, entries):
result = model.clone()
result.add_patches({"weight": ("diff", (torch.ones(2, 2),))})
return result, None
apply = Mock(side_effect=prepare)
cache = LoraPreparationCache()
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "fixture"
path.write_bytes(b"fixture")
entries = [(str(path), 1., 0.)]
first, _ = cache.apply(base, None, entries, apply)
identifier = first.patches_uuid
first.model_options["transformer_options"]["changed"] = True
first.set_attachments("downstream", {"mutable": []})
second, _ = cache.apply(base, None, entries, apply)
self.assertEqual(apply.call_count, 1)
self.assertEqual(identifier, second.patches_uuid)
self.assertEqual(second.model_options, base.model_options)
self.assertFalse(second.attachments)
self.assertFalse(base.patches)
self.assertTrue(second.patches)
self.assertEqual(second.object_patches["manual_cast_dtype"], torch.float32)
base.set_model_compute_dtype(torch.bfloat16)
third, _ = cache.apply(base, None, entries, apply)
self.assertEqual(apply.call_count, 2)
self.assertEqual(third.object_patches["manual_cast_dtype"], torch.bfloat16)
if __name__ == "__main__":
unittest.main()