86 lines
3.5 KiB
Python
86 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from flux2_klein_conditioning.nodes import (
|
|
Flux2KleinPromptReferenceBalance,
|
|
Flux2KleinReferenceAppearanceBalancer,
|
|
Flux2KleinReferenceLatentMixer,
|
|
)
|
|
from flux2_klein_conditioning.reference import gaussian_blur_per_channel
|
|
|
|
|
|
def test_gaussian_blur_preserves_constant_tensor():
|
|
x = torch.ones(1, 4, 8, 8, dtype=torch.float32)
|
|
y = gaussian_blur_per_channel(x, 2)
|
|
assert torch.allclose(x, y, atol=1e-6, rtol=1e-6)
|
|
|
|
|
|
def test_reference_mixer_updates_all_refs_by_default():
|
|
node = Flux2KleinReferenceLatentMixer()
|
|
ref_a = torch.ones(1, 4, 2, 2)
|
|
ref_b = torch.ones(1, 4, 2, 2) * 2
|
|
conditioning = [(torch.zeros(1, 4, 4), {"reference_latents": [ref_a, ref_b]})]
|
|
out, = node.control(conditioning, reference_keep=0.0, replace_mode="zeros", channel_mask_end=4)
|
|
refs = out[0][1]["reference_latents"]
|
|
assert torch.count_nonzero(refs[0]) == 0
|
|
assert torch.count_nonzero(refs[1]) == 0
|
|
|
|
|
|
def test_reference_mixer_can_target_single_ref():
|
|
node = Flux2KleinReferenceLatentMixer()
|
|
ref_a = torch.ones(1, 4, 2, 2)
|
|
ref_b = torch.ones(1, 4, 2, 2) * 2
|
|
conditioning = [(torch.zeros(1, 4, 4), {"reference_latents": [ref_a, ref_b]})]
|
|
out, = node.control(
|
|
conditioning,
|
|
reference_keep=0.0,
|
|
replace_mode="zeros",
|
|
channel_mask_end=4,
|
|
target_reference_index=1,
|
|
)
|
|
refs = out[0][1]["reference_latents"]
|
|
assert torch.count_nonzero(refs[0]) == ref_a.numel()
|
|
assert torch.count_nonzero(refs[1]) == 0
|
|
|
|
|
|
def test_prompt_reference_balance_reduces_reference_when_balance_is_text_only():
|
|
node = Flux2KleinPromptReferenceBalance()
|
|
cond = torch.ones(1, 8, 4)
|
|
ref = torch.ones(1, 4, 2, 2)
|
|
meta = {"attention_mask": torch.tensor([[1, 1, 1, 1, 1, 0, 0, 0]]), "reference_latents": [ref]}
|
|
out, = node.balance([(cond, meta)], balance=1.0, replace_mode="zeros", skip_bos=False)
|
|
assert torch.count_nonzero(out[0][1]["reference_latents"][0]) == 0
|
|
|
|
|
|
def test_reference_mixer_lowpass_mode_preserves_nonzero_coarse_signal():
|
|
node = Flux2KleinReferenceLatentMixer()
|
|
ref = torch.arange(1, 1 + 1 * 4 * 6 * 6, dtype=torch.float32).reshape(1, 4, 6, 6)
|
|
conditioning = [(torch.zeros(1, 4, 4), {"reference_latents": [ref]})]
|
|
out, = node.control(conditioning, reference_keep=0.0, replace_mode="lowpass_reference", channel_mask_end=4)
|
|
mixed = out[0][1]["reference_latents"][0]
|
|
assert torch.count_nonzero(mixed) > 0
|
|
assert not torch.allclose(mixed, ref)
|
|
|
|
|
|
def test_reference_appearance_balancer_boosts_coarse_component_and_reduces_detail():
|
|
node = Flux2KleinReferenceAppearanceBalancer()
|
|
base_x = torch.linspace(0.0, 1.0, 8).view(1, 1, 1, 8).expand(1, 4, 8, 8)
|
|
base_y = torch.linspace(0.0, 1.0, 8).view(1, 1, 8, 1).expand(1, 4, 8, 8)
|
|
coarse = base_x + base_y
|
|
checker = ((torch.arange(8).view(1, 1, 8, 1) + torch.arange(8).view(1, 1, 1, 8)) % 2).float() * 2.0 - 1.0
|
|
detail = checker.expand_as(coarse) * 0.25
|
|
ref = coarse + detail
|
|
conditioning = [(torch.zeros(1, 4, 4), {"reference_latents": [ref]})]
|
|
|
|
out, = node.balance(conditioning, appearance_scale=1.5, detail_scale=0.25, blur_radius=2, channel_mask_end=4)
|
|
modified = out[0][1]["reference_latents"][0]
|
|
|
|
input_low = gaussian_blur_per_channel(ref, 2)
|
|
output_low = gaussian_blur_per_channel(modified, 2)
|
|
input_high = ref - input_low
|
|
output_high = modified - output_low
|
|
|
|
assert output_low.abs().mean() > input_low.abs().mean()
|
|
assert output_high.abs().mean() < input_high.abs().mean()
|