195 lines
8.8 KiB
Python
195 lines
8.8 KiB
Python
import unittest
|
|
from unittest import mock
|
|
|
|
try:
|
|
import torch
|
|
except Exception:
|
|
torch = None
|
|
|
|
try:
|
|
from .test_flux2klein_helpers import (
|
|
BrokenClipStub,
|
|
FakeMask,
|
|
FakeVAE,
|
|
FakeTensor,
|
|
NativeEditModelStub,
|
|
NativeFluxStub,
|
|
StandardClipStub,
|
|
import_plugin_submodule,
|
|
load_plugin_package,
|
|
)
|
|
except ImportError:
|
|
from test_flux2klein_helpers import (
|
|
BrokenClipStub,
|
|
FakeMask,
|
|
FakeVAE,
|
|
FakeTensor,
|
|
NativeEditModelStub,
|
|
NativeFluxStub,
|
|
StandardClipStub,
|
|
import_plugin_submodule,
|
|
load_plugin_package,
|
|
)
|
|
|
|
|
|
class TestFlux2KleinNode(unittest.TestCase):
|
|
def setUp(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSFlux2KleinEditTextEncode"]
|
|
self.node = node_cls()
|
|
self.module = import_plugin_submodule(plugin, "flux2klein.lls_flux2klein_edit_text_encode")
|
|
|
|
def test_fallback_path_builds_conditioning_latent_and_custom_output(self):
|
|
conditioning, latent, custom_output, main_image, output_mask = self.node.encode(
|
|
clip=StandardClipStub(),
|
|
vae=FakeVAE(latent_channels=4, downscale_ratio=8),
|
|
image1=FakeTensor((1, 600, 400, 3), label="main"),
|
|
prompt="make it cinematic",
|
|
ref_longest_edge=1024,
|
|
resize_mode="keep_original",
|
|
mask_mode="use_mask",
|
|
image2=FakeTensor((1, 256, 512, 3), label="ref2"),
|
|
image3=None,
|
|
mask=FakeMask((1, 300, 200), mask_bbox=(10, 10, 40, 40), mask_area_ratio=0.1, label="mask"),
|
|
)
|
|
|
|
self.assertEqual(
|
|
self.module._build_flux2klein_prompt("make it cinematic", 2),
|
|
"image1: <|vision_start|><|image_pad|><|vision_end|> "
|
|
"image2: <|vision_start|><|image_pad|><|vision_end|> make it cinematic",
|
|
)
|
|
self.assertEqual(
|
|
conditioning[0][0],
|
|
"cond::image1: <|vision_start|><|image_pad|><|vision_end|> "
|
|
"image2: <|vision_start|><|image_pad|><|vision_end|> make it cinematic",
|
|
)
|
|
self.assertEqual(len(conditioning[0][1]["reference_latents"]), 2)
|
|
self.assertEqual(conditioning[0][1]["reference_latents"][0].shape, (1, 4, 75, 50))
|
|
self.assertEqual(conditioning[0][1]["reference_latents"][1].shape, (1, 4, 75, 50))
|
|
self.assertEqual(latent["samples"].shape, (1, 4, 75, 50))
|
|
self.assertEqual(latent["noise_mask"].shape, (1, 75, 50))
|
|
self.assertEqual(main_image.shape, (1, 600, 400, 3))
|
|
self.assertEqual(output_mask.shape, (1, 600, 400))
|
|
self.assertEqual(custom_output["node_name"], "LLS Flux2Klein Edit Text Encode")
|
|
self.assertEqual(custom_output["prompt"], "make it cinematic")
|
|
self.assertEqual(custom_output["main_image"], "image1")
|
|
self.assertEqual(custom_output["reference_images"], ["image2"])
|
|
self.assertEqual(custom_output["num_reference_images"], 1)
|
|
self.assertEqual(custom_output["conditioning_backend"], "flux2klein_multivision_clip")
|
|
self.assertTrue(custom_output["has_mask"])
|
|
self.assertEqual(custom_output["latent_mode"], "image1_reference_latent")
|
|
self.assertEqual(custom_output["latent_has_noise_mask"], True)
|
|
self.assertEqual(custom_output["vision_image_count"], 2)
|
|
self.assertEqual(custom_output["reference_latent_count"], 2)
|
|
self.assertEqual(self.module._count_present_reference_images(main_image, None, None), 1)
|
|
|
|
def test_longest_edge_resize_scales_images_and_inverts_mask(self):
|
|
conditioning, latent, custom_output, main_image, output_mask = self.node.encode(
|
|
clip=StandardClipStub(),
|
|
vae=FakeVAE(latent_channels=16, downscale_ratio=16),
|
|
image1=FakeTensor((1, 800, 400, 3), label="main"),
|
|
prompt="change outfit",
|
|
ref_longest_edge=1024,
|
|
resize_mode="longest_edge",
|
|
mask_mode="invert_mask",
|
|
image2=FakeTensor((1, 400, 1600, 3), label="ref2"),
|
|
image3=FakeTensor((1, 600, 300, 3), label="ref3"),
|
|
mask=FakeMask((1, 800, 400), mask_bbox=(20, 20, 180, 380), mask_area_ratio=0.25, label="mask"),
|
|
)
|
|
|
|
self.assertEqual(
|
|
conditioning[0][0],
|
|
"cond::image1: <|vision_start|><|image_pad|><|vision_end|> "
|
|
"image2: <|vision_start|><|image_pad|><|vision_end|> "
|
|
"image3: <|vision_start|><|image_pad|><|vision_end|> change outfit",
|
|
)
|
|
self.assertEqual(len(conditioning[0][1]["reference_latents"]), 3)
|
|
self.assertEqual(main_image.shape, (1, 1024, 512, 3))
|
|
self.assertEqual(latent["samples"].shape, (1, 16, 64, 32))
|
|
self.assertEqual(latent["noise_mask"].shape, (1, 64, 32))
|
|
self.assertEqual(output_mask.shape, (1, 1024, 512))
|
|
self.assertAlmostEqual(output_mask.mask_area_ratio, 0.75)
|
|
self.assertEqual(custom_output["reference_images"], ["image2", "image3"])
|
|
self.assertEqual(custom_output["num_reference_images"], 2)
|
|
self.assertEqual(custom_output["resize_mode"], "longest_edge")
|
|
self.assertEqual(custom_output["mask_mode"], "invert_mask")
|
|
self.assertEqual(custom_output["vision_image_count"], 3)
|
|
self.assertEqual(custom_output["reference_latent_count"], 3)
|
|
self.assertEqual(custom_output["latent_has_noise_mask"], True)
|
|
|
|
def test_mask_mode_none_outputs_full_mask_but_no_noise_mask(self):
|
|
conditioning, latent, custom_output, main_image, output_mask = self.node.encode(
|
|
clip=StandardClipStub(),
|
|
vae=FakeVAE(latent_channels=16, downscale_ratio=16),
|
|
image1=FakeTensor((1, 512, 512, 3), label="main"),
|
|
prompt="replace background",
|
|
ref_longest_edge=1024,
|
|
resize_mode="keep_original",
|
|
mask_mode="none",
|
|
image2=FakeTensor((1, 512, 512, 3), label="ref2"),
|
|
image3=None,
|
|
mask=FakeMask((1, 512, 512), mask_bbox=(0, 0, 256, 256), mask_area_ratio=0.25, label="mask"),
|
|
)
|
|
|
|
self.assertEqual(main_image.shape, (1, 512, 512, 3))
|
|
self.assertEqual(output_mask.shape, (1, 512, 512))
|
|
if hasattr(output_mask, "mask_area_ratio"):
|
|
self.assertEqual(output_mask.mask_area_ratio, 1.0)
|
|
else:
|
|
self.assertTrue(torch is not None)
|
|
self.assertTrue(torch.all(output_mask == 1))
|
|
self.assertNotIn("noise_mask", latent)
|
|
self.assertEqual(custom_output["latent_has_noise_mask"], False)
|
|
|
|
def test_flux2_official_path_uses_empty_latent_and_reference_latent_nodes(self):
|
|
NativeFluxStub.reset()
|
|
NativeEditModelStub.reset()
|
|
|
|
with mock.patch.object(self.module, "native_flux_nodes", NativeFluxStub, create=True), mock.patch.object(
|
|
self.module,
|
|
"native_edit_model_nodes",
|
|
NativeEditModelStub,
|
|
create=True,
|
|
):
|
|
conditioning, latent, custom_output, main_image, output_mask = self.node.encode(
|
|
clip=StandardClipStub(),
|
|
vae=FakeVAE(latent_channels=128, downscale_ratio=16),
|
|
image1=FakeTensor((1, 832, 1248, 3), label="main"),
|
|
prompt="replace background with a clean kitchen scene",
|
|
ref_longest_edge=1248,
|
|
resize_mode="keep_original",
|
|
mask_mode="use_mask",
|
|
image2=FakeTensor((1, 640, 640, 3), label="ref2"),
|
|
image3=None,
|
|
mask=FakeMask((1, 832, 1248), mask_bbox=(50, 40, 500, 700), mask_area_ratio=0.35, label="mask"),
|
|
)
|
|
|
|
self.assertEqual(NativeFluxStub.last_empty_call, {"width": 1248, "height": 832, "batch_size": 1})
|
|
self.assertEqual(len(NativeEditModelStub.last_calls), 2)
|
|
self.assertEqual(latent["samples"].shape, (1, 128, 52, 78))
|
|
self.assertEqual(latent["noise_mask"].shape, (1, 52, 78))
|
|
self.assertEqual(custom_output["latent_mode"], "empty_flux2_latent")
|
|
self.assertEqual(custom_output["conditioning_backend"], "flux2klein_multivision_clip+reference_latent")
|
|
self.assertEqual(custom_output["reference_latent_count"], 2)
|
|
self.assertEqual(main_image.shape, (1, 832, 1248, 3))
|
|
self.assertEqual(output_mask.shape, (1, 832, 1248))
|
|
|
|
def test_missing_clip_backend_raises_clear_error(self):
|
|
with self.assertRaisesRegex(RuntimeError, "vision-aware|text encoding backend"):
|
|
self.node.encode(
|
|
clip=BrokenClipStub(),
|
|
vae=FakeVAE(),
|
|
image1=FakeTensor((1, 512, 512, 3), label="main"),
|
|
prompt="replace background",
|
|
ref_longest_edge=1024,
|
|
resize_mode="keep_original",
|
|
mask_mode="none",
|
|
image2=None,
|
|
image3=None,
|
|
mask=None,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|