296 lines
16 KiB
Python
296 lines
16 KiB
Python
import importlib.util
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
import torch
|
|
import comfy.sampler_helpers
|
|
|
|
|
|
spec = importlib.util.spec_from_file_location("fl_krea_test", Path(__file__).parents[1] / "nodes/conditioning/FL_KreaReference.py")
|
|
m = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(m)
|
|
|
|
|
|
def model():
|
|
result = SimpleNamespace(model=Mock(spec=m.comfy.model_base.Krea2), model_options={}, is_dynamic=lambda: False,
|
|
get_model_object=lambda name: SimpleNamespace(percent_to_sigma=lambda p: 1 - p))
|
|
result.clone = model
|
|
result.set_model_sampler_calc_cond_batch_function = lambda fn: result.model_options.update(sampler_calc_cond_batch_function=fn)
|
|
return result
|
|
|
|
|
|
def conditioning(value):
|
|
return [[torch.full((1, 2, 30720), float(value)), {"value": value}]]
|
|
|
|
|
|
def reference(**kwargs):
|
|
kwargs.setdefault("reference_mode", "full")
|
|
return m.FL_KreaReference.execute(torch.zeros(1, 48, 96, 3), **kwargs).result[0]
|
|
|
|
|
|
class KreaReferenceTests(unittest.TestCase):
|
|
def predict(self, branches, influence=.7, sigma=.5):
|
|
guider = m.KreaReferenceGuider(model(), conditioning(2), branches, influence)
|
|
guider.conds = guider.original_conds
|
|
guider.inner_model = object()
|
|
calls = []
|
|
|
|
def sample(inner, x, timestep, negative, positive, cfg, **kwargs):
|
|
calls.append(positive[0]["value"])
|
|
return torch.full_like(x, positive[0]["value"])
|
|
|
|
x = torch.zeros(1, 16, 2, 2)
|
|
with patch.object(m.comfy.samplers, "sampling_function", side_effect=sample):
|
|
result = guider.predict_noise(x, torch.tensor([sigma]))
|
|
return result, calls
|
|
|
|
def test_zero_influence_and_no_references_use_baseline(self):
|
|
for branches, strength in (([], .7), ([(reference(), conditioning(8))], 0)):
|
|
output, calls = self.predict(branches, strength)
|
|
self.assertEqual(calls, [2])
|
|
torch.testing.assert_close(output, torch.full_like(output, 2))
|
|
|
|
def test_single_reference_endpoint_skips_baseline(self):
|
|
output, calls = self.predict([(reference(), conditioning(8))], 1)
|
|
self.assertEqual(calls, [8])
|
|
torch.testing.assert_close(output, torch.full_like(output, 8))
|
|
|
|
def test_blend_formula_order_and_weight_scaling(self):
|
|
branches = [(reference(weight=.25), conditioning(4)), (reference(weight=.75), conditioning(8))]
|
|
a, calls = self.predict(branches, .5)
|
|
b, _ = self.predict(list(reversed(branches)), .5)
|
|
c, _ = self.predict([(dict(r, weight=r["weight"] * .5), cond) for r, cond in branches], .5)
|
|
self.assertEqual(calls, [2, 4, 8])
|
|
torch.testing.assert_close(a, torch.full_like(a, 4.5))
|
|
torch.testing.assert_close(a, b)
|
|
torch.testing.assert_close(c, torch.full_like(c, 3.25))
|
|
|
|
def test_single_reference_slider_scales_contribution(self):
|
|
for weight in (0, .05, .25, .5, 1):
|
|
with self.subTest(weight=weight):
|
|
output, _ = self.predict([(reference(weight=weight), conditioning(8))], .7)
|
|
torch.testing.assert_close(output, torch.full_like(output, 2 + .7 * weight * (8 - 2)))
|
|
|
|
def test_removing_reference_does_not_boost_remaining_reference(self):
|
|
a = (reference(weight=.25), conditioning(4))
|
|
b = (reference(weight=.5), conditioning(8))
|
|
both, _ = self.predict([a, b], 1)
|
|
only_b, _ = self.predict([b], 1)
|
|
zero_a, _ = self.predict([(reference(weight=0), conditioning(4)), b], 1)
|
|
torch.testing.assert_close(both - only_b, torch.full_like(both, .25 * (4 - 2)))
|
|
torch.testing.assert_close(only_b, zero_a)
|
|
|
|
def test_combined_weights_above_one_keep_unit_prediction_gain(self):
|
|
output, _ = self.predict([(reference(weight=1), conditioning(4)), (reference(weight=1), conditioning(4))], 1)
|
|
torch.testing.assert_close(output, torch.full_like(output, 6))
|
|
same, _ = self.predict([(reference(weight=1), conditioning(2)), (reference(weight=1), conditioning(2))], 1)
|
|
torch.testing.assert_close(same, torch.full_like(same, 2))
|
|
|
|
def test_schedule_does_not_renormalize_other_reference(self):
|
|
branches = [(reference(weight=.5, start=.6), conditioning(4)), (reference(weight=.5), conditioning(8))]
|
|
a, calls = self.predict(branches, 1, sigma=.7)
|
|
self.assertEqual(calls, [2, 8])
|
|
torch.testing.assert_close(a, torch.full_like(a, 5))
|
|
|
|
def test_envelope_boundaries_and_fades(self):
|
|
bounds = (.8, .6, .4, .2)
|
|
for sigma, expected in ((1, 0), (.8, 0), (.7, .5), (.6, 1), (.5, 1), (.3, .5), (.2, 0), (.1, 0)):
|
|
self.assertAlmostEqual(m.sigma_envelope(sigma, bounds), expected)
|
|
self.assertEqual(m.sigma_envelope(1, (1, 1, 0, 0)), 1)
|
|
self.assertEqual(m.sigma_envelope(0, (1, 1, 0, 0)), 1)
|
|
|
|
def test_repeat_evaluations_follow_sigma_not_call_count(self):
|
|
branches = [(reference(start=.25, end=.75), conditioning(8))]
|
|
for sigma in (.9, .5, .5, .1, .9):
|
|
output, _ = self.predict(branches, 1, sigma)
|
|
torch.testing.assert_close(output, torch.full_like(output, 8 if sigma == .5 else 2))
|
|
|
|
def test_reference_validation(self):
|
|
for kwargs in (dict(start=.5, end=.5), dict(start=-.1), dict(end=1.1), dict(fade=.6),
|
|
dict(weight=-1), dict(weight=1.01), dict(weight=10), dict(weight=float("nan")),
|
|
dict(weight=float("inf")), dict(role="unknown"), dict(resolution=999)):
|
|
with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
|
|
reference(**kwargs)
|
|
with self.assertRaisesRegex(ValueError, "one RGB image"):
|
|
m.FL_KreaReference.execute(torch.zeros(2, 32, 32, 3))
|
|
|
|
def clip(self):
|
|
clip = Mock()
|
|
clip.tokenizer = Mock(spec=m.Krea2Tokenizer)
|
|
clip.tokenize.side_effect = lambda text, **kwargs: (text, kwargs)
|
|
clip.encode_from_tokens_scheduled.side_effect = lambda tokens: conditioning(len(tokens[0]))
|
|
return clip
|
|
|
|
def test_disabled_and_zero_weight_references_are_not_encoded(self):
|
|
clip = self.clip()
|
|
with patch.object(m, "encode_reference", return_value=conditioning(8)) as encode:
|
|
guider = m.FL_KreaReferenceGuider.execute(model(), clip, "dog", references={
|
|
"reference_0": reference(enabled=False), "reference_1": reference(weight=0),
|
|
"reference_2": reference(),
|
|
}).result[0]
|
|
self.assertEqual(encode.call_count, 1)
|
|
self.assertEqual(len(guider.references), 1)
|
|
self.assertEqual(guider.references[0][1], 1)
|
|
|
|
def test_zero_strength_skips_reference_encoding(self):
|
|
clip = self.clip()
|
|
with patch.object(m, "encode_reference") as encode:
|
|
m.FL_KreaReferenceGuider.execute(model(), clip, "dog", influence=0, references={"reference_0": reference()})
|
|
encode.assert_not_called()
|
|
clip.tokenize.assert_called_once_with("dog")
|
|
|
|
def test_same_shape_images_and_clip_changes_are_reencoded(self):
|
|
clip = self.clip()
|
|
ref = reference()
|
|
encoded_images = []
|
|
|
|
def tokenize(text, **kwargs):
|
|
if "images" in kwargs:
|
|
encoded_images.append(kwargs["images"][0].clone())
|
|
return text, kwargs
|
|
|
|
clip.tokenize.side_effect = tokenize
|
|
for image in (torch.zeros_like(ref["image"]), torch.ones_like(ref["image"])):
|
|
m.FL_KreaReferenceGuider.execute(model(), clip, "dog", references={"reference_0": dict(ref, image=image)})
|
|
other_clip = self.clip()
|
|
m.FL_KreaReferenceGuider.execute(model(), other_clip, "dog", references={"reference_0": ref})
|
|
self.assertEqual(len(encoded_images), 2)
|
|
self.assertFalse(torch.equal(*encoded_images))
|
|
self.assertEqual(other_clip.encode_from_tokens_scheduled.call_count, 2)
|
|
|
|
def test_each_branch_sees_one_image_and_preserves_taps(self):
|
|
clip = self.clip()
|
|
refs = {"reference_0": reference(), "reference_1": reference(role="palette")}
|
|
guider = m.FL_KreaReferenceGuider.execute(model(), clip, "dog", references=refs).result[0]
|
|
calls = clip.tokenize.call_args_list
|
|
self.assertEqual(len(calls), 3)
|
|
for call in calls[1:]:
|
|
self.assertEqual(len(call.kwargs["images"]), 1)
|
|
self.assertEqual(call.args[0].count("<|image_pad|>"), 1)
|
|
self.assertIn("dog", call.args[0])
|
|
self.assertIn(m.ROLES["palette"], calls[2].args[0])
|
|
for cond in guider.original_conds.values():
|
|
self.assertEqual(cond[0]["cross_attn"].shape[-1], 30720)
|
|
self.assertTrue(torch.all(cond[0]["cross_attn"] != 0))
|
|
|
|
def test_model_and_guidance_mismatch_rejected_before_encoding(self):
|
|
clip = self.clip()
|
|
bad = model()
|
|
bad.model = object()
|
|
with self.assertRaisesRegex(ValueError, "Krea 2"):
|
|
m.FL_KreaReferenceGuider.execute(bad, clip, "dog")
|
|
bad = model()
|
|
bad.model_options["sampler_cfg_function"] = lambda args: None
|
|
with self.assertRaisesRegex(ValueError, "guidance patches"):
|
|
m.FL_KreaReferenceGuider.execute(bad, clip, "dog")
|
|
clip.encode_from_tokens_scheduled.assert_not_called()
|
|
|
|
def test_schema_has_optional_growing_references(self):
|
|
schema = m.FL_KreaReferenceGuider.INPUT_TYPES()
|
|
self.assertEqual(schema["optional"]["references"][1]["template"]["min"], 0)
|
|
|
|
def test_ksampler_outputs_preserve_guider_and_original_model(self):
|
|
original = model()
|
|
guider, patched, cond = m.FL_KreaReferenceGuider.execute(original, self.clip(), "dog", references={
|
|
"reference_0": reference(weight=.25), "reference_1": reference(weight=.5),
|
|
}).result
|
|
self.assertIsInstance(guider, m.KreaReferenceGuider)
|
|
self.assertEqual(original.model_options, {})
|
|
self.assertIs(patched.model_options["sampler_calc_cond_batch_function"], m.reference_cond_batch)
|
|
self.assertEqual(len(cond), 3)
|
|
self.assertNotIn("fl_krea_reference", cond[0][1])
|
|
self.assertEqual(cond[1][1]["fl_krea_reference"][1], .7 * .25)
|
|
|
|
def test_ksampler_matches_guider_weights_schedules_and_cfg(self):
|
|
x = torch.zeros(1, 16, 2, 2)
|
|
|
|
def batch(inner, conds, x, sigma, options):
|
|
return [torch.full_like(x, cond[0]["value"] if cond else 0) for cond in conds]
|
|
|
|
for weight in (0, .05, .5, 1):
|
|
for sigma in (.9, .5, .1):
|
|
refs = [reference(weight=weight, start=.2, end=.8, fade=.25), reference(weight=.3)]
|
|
branches = list(zip(refs, (conditioning(8), conditioning(4))))
|
|
expected, _ = self.predict(branches, .7, sigma)
|
|
clip = self.clip()
|
|
clip.encode_from_tokens_scheduled.side_effect = lambda tokens: conditioning(2)
|
|
with patch.object(m, "encode_reference", side_effect=[cond for ref, cond in branches if ref["weight"] > 0]):
|
|
_, patched, cond = m.FL_KreaReferenceGuider.execute(model(), clip, "dog", references=dict(enumerate(refs))).result
|
|
prepared = comfy.sampler_helpers.convert_cond(cond)
|
|
negative = comfy.sampler_helpers.convert_cond(conditioning(-2))
|
|
with patch.object(m.comfy.samplers, "calc_cond_batch", side_effect=batch):
|
|
for cfg in (1., 2.):
|
|
actual = m.comfy.samplers.sampling_function(object(), x, torch.tensor([sigma]), negative, prepared, cfg, model_options=patched.model_options)
|
|
torch.testing.assert_close(actual, -2 + (expected + 2) * cfg)
|
|
|
|
def test_batch_hook_leaves_plain_and_missing_conditioning_unchanged(self):
|
|
cond = [{"value": 2}]
|
|
with patch.object(m.comfy.samplers, "calc_cond_batch", side_effect=lambda model, conds, *args: conds):
|
|
out = m.reference_cond_batch(dict(model=object(), conds=[cond, None], input=None, sigma=torch.tensor([.5]), model_options={}))
|
|
self.assertEqual(out, [cond, None])
|
|
|
|
def test_average_blends_predictions_identically_for_both_sampler_outputs(self):
|
|
cases = [
|
|
([], [], 2, 2),
|
|
([reference()], [8], 8, 8),
|
|
([reference() for _ in range(4)], [4, 6, 8, 10], 22, 7),
|
|
([reference(), reference(weight=.5)], [4, 8], 7, 4.5),
|
|
([reference(weight=0), reference()], [4, 8], 8, 5),
|
|
([reference(enabled=False), reference()], [4, 8], 8, 8),
|
|
([reference(start=.6), reference()], [4, 8], 8, 5),
|
|
([reference(start=.25, end=.75, fade=.5), reference()], [4, 8], 10, 6),
|
|
]
|
|
x = torch.zeros(1, 16, 2, 2)
|
|
for refs, values, additive, average in cases:
|
|
for mode in ("add", "average"):
|
|
for amount in (0, .5, 1):
|
|
for influence in (0, .7, 1):
|
|
with self.subTest(values=values, mode=mode, amount=amount, influence=influence):
|
|
clip = self.clip()
|
|
clip.encode_from_tokens_scheduled.side_effect = lambda tokens: conditioning(2)
|
|
encoded = [conditioning(value) for ref, value in zip(refs, values) if ref["enabled"] and ref["weight"] > 0]
|
|
with patch.object(m, "encode_reference", side_effect=encoded):
|
|
guider, patched, cond = m.FL_KreaReferenceGuider.execute(
|
|
model(), clip, "dog", influence, dict(enumerate(refs)), mode, amount).result
|
|
guider.conds = guider.original_conds
|
|
guider.inner_model = object()
|
|
with patch.object(m.comfy.samplers, "sampling_function", side_effect=lambda inner, x, t, neg, pos, cfg, **kwargs: torch.full_like(x, pos[0]["value"])):
|
|
guided = guider.predict_noise(x, torch.tensor([.5]))
|
|
with patch.object(m.comfy.samplers, "calc_cond_batch", side_effect=lambda inner, conds, x, sigma, options: [torch.full_like(x, cond[0]["value"] if cond else 0) for cond in conds]):
|
|
sampled = m.comfy.samplers.sampling_function(
|
|
object(), x, torch.tensor([.5]), None, comfy.sampler_helpers.convert_cond(cond), 1,
|
|
model_options=patched.model_options)
|
|
blended = additive + (average - additive) * amount if mode == "average" else additive
|
|
expected = torch.full_like(x, 2 + influence * (blended - 2))
|
|
torch.testing.assert_close(guided, expected)
|
|
torch.testing.assert_close(sampled, expected)
|
|
|
|
def test_invalid_blend_controls_fail_before_encoding(self):
|
|
clip = self.clip()
|
|
for kwargs in (dict(blend_mode="unknown"), dict(average_amount=-.1), dict(average_amount=1.1), dict(average_amount=float("nan"))):
|
|
with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
|
|
m.FL_KreaReferenceGuider.execute(model(), clip, "dog", **kwargs)
|
|
clip.encode_from_tokens_scheduled.assert_not_called()
|
|
|
|
def test_context_mode_removes_visual_prefix_and_copies_suffix(self):
|
|
clip = self.clip()
|
|
clip.tokenize.side_effect = None
|
|
clip.tokenize.return_value = {"qwen3vl_4b": [[(151652, 1), ({"type": "image"}, 1), (151653, 1), (12, 1), (13, 1)]]}
|
|
tensor = torch.randn(1, 12, 30720)
|
|
mask = torch.ones(1, 12)
|
|
clip.encode_from_tokens_scheduled.side_effect = None
|
|
clip.encode_from_tokens_scheduled.return_value = [[tensor, {"attention_mask": mask, "other": 7}]]
|
|
result = m.encode_reference(clip, "dog", reference(reference_mode="context"))
|
|
torch.testing.assert_close(result[0][0], tensor[:, -2:])
|
|
self.assertEqual(result[0][0].shape[-1], 30720)
|
|
self.assertNotEqual(result[0][0].untyped_storage().data_ptr(), tensor.untyped_storage().data_ptr())
|
|
self.assertEqual(result[0][1]["attention_mask"].shape, (1, 2))
|
|
self.assertEqual(result[0][1]["other"], 7)
|
|
self.assertEqual(mask.shape, (1, 12))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|