702 lines
24 KiB
Python
702 lines
24 KiB
Python
import importlib.util
|
|
import pathlib
|
|
import sys
|
|
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parents[1]
|
|
COMFY_ROOT = ROOT.parents[2]
|
|
|
|
|
|
def load_plugin():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"ComfyUI_Krea2_VisualRef",
|
|
ROOT / "__init__.py",
|
|
submodule_search_locations=[str(ROOT)],
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FakeTensor:
|
|
def __init__(self, shape):
|
|
self.shape = shape
|
|
|
|
|
|
class FakeImage:
|
|
def __init__(self, shape):
|
|
self.shape = shape
|
|
|
|
|
|
class FakeKrea2Tokenizer:
|
|
pass
|
|
|
|
|
|
class FakeKrea2TEModel:
|
|
pass
|
|
|
|
|
|
class FakeKrea2Clip:
|
|
def __init__(self):
|
|
self.tokenizer = FakeKrea2Tokenizer()
|
|
self.cond_stage_model = FakeKrea2TEModel()
|
|
self.tokenize_calls = []
|
|
self.encode_calls = []
|
|
|
|
def tokenize(self, text, **kwargs):
|
|
self.tokenize_calls.append((text, kwargs))
|
|
if "images" not in kwargs:
|
|
return {"qwen3vl_4b": [[(151644, 1.0), (198, 1.0)]]}
|
|
row = [(151644, 1.0)]
|
|
for image in kwargs["images"]:
|
|
row.append(({"type": "image", "data": image, "original_type": "image"}, 1.0))
|
|
return {"qwen3vl_4b": [row]}
|
|
|
|
def encode_from_tokens_scheduled(self, tokens):
|
|
self.encode_calls.append(tokens)
|
|
return [[FakeTensor((1, 42, 30720)), {"attention_mask": FakeTensor((1, 42))}]]
|
|
|
|
def generate(self, tokens, **kwargs):
|
|
self.generate_call = (tokens, kwargs)
|
|
return [101, 102, 103]
|
|
|
|
def decode(self, token_ids, skip_special_tokens=True):
|
|
self.decode_call = (token_ids, skip_special_tokens)
|
|
return "soft watercolor storybook illustration, warm pastel palette, gentle brush texture, diffused light"
|
|
|
|
class FakeKrea2MultiStyleClip(FakeKrea2Clip):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.generate_calls = []
|
|
|
|
def generate(self, tokens, **kwargs):
|
|
index = len(self.generate_calls)
|
|
self.generate_calls.append((tokens, kwargs))
|
|
return [index]
|
|
|
|
def decode(self, token_ids, skip_special_tokens=True):
|
|
self.decode_call = (token_ids, skip_special_tokens)
|
|
styles = [
|
|
"fine monochrome ink linework",
|
|
"soft pastel watercolor wash",
|
|
"flat graphic poster shapes",
|
|
]
|
|
index = int(token_ids[0])
|
|
if index >= len(styles):
|
|
return "Style lock: mixed-media illustration; style recipe: 45% fine ink linework, 35% soft pastel watercolor washes, 20% flat poster geometry; hierarchy: ink edges with watercolor surfaces and poster accents; palette: muted warm-cool contrast; line_shape: clean minimal contours; lighting_texture: soft wash grain; mood: calm crafted; avoid: copied subjects, text, logos"
|
|
return styles[index]
|
|
|
|
|
|
class FakeNativeClip:
|
|
def __init__(self):
|
|
self.tokenize_calls = []
|
|
self.encode_calls = []
|
|
|
|
def tokenize(self, text, **kwargs):
|
|
self.tokenize_calls.append((text, kwargs))
|
|
return {"native": [[(text, 1.0)]]}
|
|
|
|
def encode_from_tokens_scheduled(self, tokens):
|
|
self.encode_calls.append(tokens)
|
|
return [[FakeTensor((1, 7, 4096)), {"native": True}]]
|
|
|
|
|
|
class FakeKrea2TensorClip(FakeKrea2Clip):
|
|
def tokenize(self, text, **kwargs):
|
|
self.tokenize_calls.append((text, kwargs))
|
|
if "images" not in kwargs:
|
|
return {"qwen3vl_4b": [[(151644, 1.0), (872, 1.0), (198, 1.0)]]}
|
|
image = kwargs["images"][0]
|
|
return {
|
|
"qwen3vl_4b": [
|
|
[
|
|
(151644, 1.0),
|
|
(8948, 1.0),
|
|
(198, 1.0),
|
|
(151645, 1.0),
|
|
(198, 1.0),
|
|
(151644, 1.0),
|
|
(872, 1.0),
|
|
(198, 1.0),
|
|
(151652, 1.0),
|
|
({"type": "image", "data": image, "original_type": "image"}, 1.0),
|
|
(151653, 1.0),
|
|
(198, 1.0),
|
|
]
|
|
]
|
|
}
|
|
|
|
def encode_from_tokens_scheduled(self, tokens):
|
|
import torch
|
|
|
|
self.encode_calls.append(tokens)
|
|
return [[torch.ones((1, 8, 30720)), {}]]
|
|
|
|
|
|
def test_node_mappings_load():
|
|
plugin = load_plugin()
|
|
assert "Krea2VisualReferenceEncode" in plugin.NODE_CLASS_MAPPINGS
|
|
assert "Krea2VisualReferenceDebug" in plugin.NODE_CLASS_MAPPINGS
|
|
assert "Krea2StylePromptEncode" in plugin.NODE_CLASS_MAPPINGS
|
|
|
|
|
|
def test_node_interfaces():
|
|
plugin = load_plugin()
|
|
encode = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]
|
|
debug = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceDebug"]
|
|
|
|
assert encode.RETURN_TYPES == ("CONDITIONING",)
|
|
assert encode.FUNCTION == "encode"
|
|
input_types = encode.INPUT_TYPES()
|
|
assert "style_image" not in input_types["required"]
|
|
assert "style_image" in input_types["optional"]
|
|
assert "strength" in input_types["required"]
|
|
assert input_types["required"]["strength"][1]["min"] == 0.0
|
|
assert input_types["required"]["strength"][1]["max"] == 1.0
|
|
assert input_types["required"]["strength"][1]["step"] == 0.01
|
|
assert "style_text_mode" in input_types["required"]
|
|
assert input_types["required"]["style_text_mode"][0] == ["auto_style_text_only", "auto_style_text_and_visual", "visual_only"]
|
|
assert "reference_intent" not in input_types["required"]
|
|
assert "max_style_tokens" not in input_types["required"]
|
|
assert "style_seed" not in input_types["required"]
|
|
|
|
assert debug.RETURN_TYPES == ("STRING",)
|
|
assert debug.FUNCTION == "debug"
|
|
|
|
|
|
def test_visual_template_uses_fixed_style_reference_guidance():
|
|
plugin = load_plugin()
|
|
visual_ref = importlib.import_module(f"{plugin.__name__}.nodes.visual_ref")
|
|
|
|
template = visual_ref._build_krea2_visual_template(
|
|
"a ceramic teapot on a wooden table",
|
|
"before_prompt",
|
|
)
|
|
|
|
assert "original-medium reference" in template
|
|
assert "photographic lighting" in template
|
|
assert "real-camera photographic realism" in template
|
|
assert "natural optics" in template
|
|
assert "do not include people" in template
|
|
assert "real camera photographic appearance" in template
|
|
assert "actual medium" in template
|
|
assert "Do not add unrelated stylization" in template
|
|
assert "Do not copy the reference subject" in template
|
|
assert "maximum visual reference strength" in template
|
|
assert "spatial relationships" not in template
|
|
assert "objects and background" not in template
|
|
assert "environment treatment" not in template
|
|
for forbidden in ("cartoon", "anime", "comic", "manga", "line art", "cel shading"):
|
|
assert forbidden not in template.lower()
|
|
|
|
|
|
def test_style_extraction_prompt_uses_structured_style_fields():
|
|
plugin = load_plugin()
|
|
visual_ref = importlib.import_module(f"{plugin.__name__}.nodes.visual_ref")
|
|
|
|
prompt = visual_ref._style_extraction_prompt()
|
|
|
|
assert "transferable visual style" in prompt
|
|
assert "medium:" in prompt
|
|
assert "signature_cues:" in prompt
|
|
assert "line_shape:" in prompt
|
|
assert "texture_surface:" in prompt
|
|
assert "negative_subjects:" in prompt
|
|
assert "Do not include bullet points" in prompt
|
|
|
|
def test_palette_hint_extracts_basic_color_characteristics():
|
|
import torch
|
|
|
|
plugin = load_plugin()
|
|
visual_ref = importlib.import_module(f"{plugin.__name__}.nodes.visual_ref")
|
|
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
|
image[..., 0] = 0.9
|
|
image[..., 1] = 0.65
|
|
image[..., 2] = 0.35
|
|
|
|
hint = visual_ref._extract_palette_hint(image)
|
|
|
|
assert "algorithmic palette hint:" in hint
|
|
assert "warm" in hint
|
|
assert "high-saturation" in hint
|
|
assert "#e5a659" in hint
|
|
|
|
|
|
def test_image_batch_items_splits_tensor_batches():
|
|
import torch
|
|
|
|
plugin = load_plugin()
|
|
visual_ref = importlib.import_module(f"{plugin.__name__}.nodes.visual_ref")
|
|
image = torch.zeros((3, 4, 4, 3), dtype=torch.float32)
|
|
|
|
items = visual_ref._image_batch_items(image)
|
|
|
|
assert len(items) == 3
|
|
assert all(item.shape == (1, 4, 4, 3) for item in items)
|
|
|
|
|
|
def test_encode_passes_image_to_tokenizer_and_tags_conditioning():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=1.0,
|
|
style_image=image,
|
|
placement="after_prompt",
|
|
style_text_mode="visual_only",
|
|
)
|
|
|
|
assert len(clip.tokenize_calls) == 1
|
|
text, kwargs = clip.tokenize_calls[0]
|
|
assert kwargs["images"] == [image]
|
|
assert "<|image_pad|>" in text
|
|
assert text.index("Generation prompt:") < text.index("<|vision_start|>")
|
|
|
|
metadata = conditioning[0][1]["krea2_visual_ref"]
|
|
assert metadata["source"] == "visual_token_mode"
|
|
assert metadata["placement"] == "after_prompt"
|
|
assert metadata["style_image_shape"] == [1, 64, 64, 3]
|
|
assert metadata["expected_feature_width"] == 30720
|
|
|
|
|
|
def test_visual_strength_uses_original_image_and_metadata_strength():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=0.75,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="visual_only",
|
|
)
|
|
|
|
assert len(clip.tokenize_calls) == 1
|
|
text, kwargs = clip.tokenize_calls[0]
|
|
assert kwargs["images"] == [image]
|
|
assert "Reference style guidance:" in text
|
|
|
|
assert len(conditioning) == 1
|
|
metadata = conditioning[0][1]["krea2_visual_ref"]
|
|
assert conditioning[0][1]["strength"] == 0.75 ** 2.5
|
|
assert metadata["strength"] == 0.75
|
|
assert metadata["effective_strength"] == 0.75 ** 2.5
|
|
assert metadata["strength_mode"] == "conditioning_strength_power_2_5"
|
|
assert metadata["style_image_shape"] == [1, 64, 64, 3]
|
|
|
|
|
|
|
|
def test_visual_reference_auto_style_text_only_skips_final_visual_tokens():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=0.75,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="auto_style_text_only",
|
|
)
|
|
|
|
assert len(clip.tokenize_calls) == 2
|
|
analysis_text, analysis_kwargs = clip.tokenize_calls[0]
|
|
native_text, native_kwargs = clip.tokenize_calls[1]
|
|
assert analysis_kwargs["images"] == [image]
|
|
assert "transferable visual style" in analysis_text
|
|
assert native_kwargs == {}
|
|
assert "TRANSFERABLE STYLE LOCK" in native_text
|
|
assert "Mandatory style-transfer instruction" in native_text
|
|
assert "soft watercolor storybook illustration" in native_text
|
|
assert len(conditioning) == 1
|
|
metadata = conditioning[0][1]["krea2_visual_ref"]
|
|
assert conditioning[0][1]["strength"] == 1.0
|
|
assert metadata["source"] == "style_text_mode"
|
|
assert metadata["style_text_mode"] == "auto_style_text_only"
|
|
assert "soft watercolor storybook illustration" in metadata["style_text"]
|
|
assert metadata["style_text_raw"] == "soft watercolor storybook illustration, warm pastel palette, gentle brush texture, diffused light"
|
|
assert metadata["palette_hint"] == ""
|
|
assert metadata["visual_placeholders"] == 0
|
|
assert metadata["visual_span"] is None
|
|
|
|
|
|
def test_visual_reference_auto_style_text_only_adds_palette_hint_for_tensor_image():
|
|
import torch
|
|
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
|
image[..., 0] = 0.9
|
|
image[..., 1] = 0.65
|
|
image[..., 2] = 0.35
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=0.5,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="auto_style_text_only",
|
|
)
|
|
|
|
native_text, native_kwargs = clip.tokenize_calls[1]
|
|
metadata = conditioning[0][1]["krea2_visual_ref"]
|
|
assert native_kwargs == {}
|
|
assert "algorithmic palette hint:" in native_text
|
|
assert "algorithmic palette hint:" in metadata["style_text"]
|
|
assert "#e5a659" in metadata["palette_hint"]
|
|
assert metadata["style_text_raw"] == "soft watercolor storybook illustration, warm pastel palette, gentle brush texture, diffused light"
|
|
|
|
|
|
def test_visual_reference_auto_style_text_only_fuses_batched_reference_styles():
|
|
import torch
|
|
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2MultiStyleClip()
|
|
image = torch.zeros((3, 4, 4, 3), dtype=torch.float32)
|
|
image[0, ..., 0] = 0.9
|
|
image[0, ..., 1] = 0.65
|
|
image[0, ..., 2] = 0.35
|
|
image[1, ..., 0] = 0.35
|
|
image[1, ..., 1] = 0.55
|
|
image[1, ..., 2] = 0.95
|
|
image[2, ..., :] = 0.25
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=0.75,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="auto_style_text_only",
|
|
)
|
|
|
|
assert len(clip.generate_calls) == 4
|
|
assert len(clip.tokenize_calls) == 5
|
|
for index in range(3):
|
|
analysis_kwargs = clip.tokenize_calls[index][1]
|
|
assert analysis_kwargs["images"][0].shape == (1, 4, 4, 3)
|
|
|
|
fusion_text, fusion_kwargs = clip.tokenize_calls[3]
|
|
assert fusion_kwargs == {}
|
|
assert "Synthesize one strong style lock" in fusion_text
|
|
assert "style hierarchy" in fusion_text
|
|
assert "style recipe:" in fusion_text
|
|
assert "Style lock:" in fusion_text
|
|
assert "reference 1: fine monochrome ink linework" in fusion_text
|
|
native_text, native_kwargs = clip.tokenize_calls[4]
|
|
metadata = conditioning[0][1]["krea2_visual_ref"]
|
|
assert native_kwargs == {}
|
|
assert metadata["style_image_count"] == 3
|
|
assert metadata["style_text_raw_items"] == [
|
|
"fine monochrome ink linework",
|
|
"soft pastel watercolor wash",
|
|
"flat graphic poster shapes",
|
|
]
|
|
assert len(metadata["palette_hint_items"]) == 3
|
|
assert "reference 1: fine monochrome ink linework" in metadata["style_text_raw"]
|
|
assert "reference 2: soft pastel watercolor wash" in metadata["style_text_raw"]
|
|
assert "reference 3: flat graphic poster shapes" in metadata["style_text_raw"]
|
|
assert "reference 1:" in metadata["palette_hint"]
|
|
assert "reference 2:" in metadata["palette_hint"]
|
|
assert "reference 3:" in metadata["palette_hint"]
|
|
assert metadata["style_text_fusion"] == "model_synthesized"
|
|
assert metadata["style_text"].startswith("Style lock: mixed-media illustration")
|
|
assert "TRANSFERABLE STYLE LOCK" in native_text
|
|
assert "mixed-media illustration" in native_text
|
|
assert "reference 1:" not in native_text
|
|
|
|
|
|
def test_visual_reference_auto_style_text_combines_text_and_visual_conditioning():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=0.75,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="auto_style_text_and_visual",
|
|
)
|
|
|
|
assert len(clip.tokenize_calls) == 3
|
|
analysis_text, analysis_kwargs = clip.tokenize_calls[0]
|
|
native_text, native_kwargs = clip.tokenize_calls[1]
|
|
visual_text, visual_kwargs = clip.tokenize_calls[2]
|
|
assert analysis_kwargs["images"] == [image]
|
|
assert "transferable visual style" in analysis_text
|
|
assert native_kwargs == {}
|
|
assert "TRANSFERABLE STYLE LOCK" in native_text
|
|
assert "Mandatory style-transfer instruction" in native_text
|
|
assert visual_kwargs["images"] == [image]
|
|
assert "TRANSFERABLE STYLE LOCK" in visual_text
|
|
assert "low-weight style calibrator" in visual_text
|
|
assert "Suppress depicted people" in visual_text
|
|
assert len(conditioning) == 2
|
|
metadata = conditioning[1][1]["krea2_visual_ref"]
|
|
assert metadata["style_text_mode"] == "auto_style_text_and_visual"
|
|
assert "soft watercolor storybook illustration" in metadata["style_text"]
|
|
assert metadata["conditioning_prompt_length"] > metadata["prompt_length"]
|
|
assert conditioning[1][1]["strength"] == 0.75 * 0.54
|
|
assert metadata["strength_mode"] == "conditioning_strength_text_visual_balanced"
|
|
assert metadata["effective_strength"] == 0.75 * 0.54
|
|
|
|
|
|
def test_visual_reference_auto_style_text_and_visual_uses_all_batched_images():
|
|
import torch
|
|
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2MultiStyleClip()
|
|
image = torch.zeros((3, 4, 4, 3), dtype=torch.float32)
|
|
image[0, ..., 0] = 0.9
|
|
image[1, ..., 1] = 0.9
|
|
image[2, ..., 2] = 0.9
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=0.75,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="auto_style_text_and_visual",
|
|
)
|
|
|
|
assert len(clip.generate_calls) == 4
|
|
visual_text, visual_kwargs = clip.tokenize_calls[-1]
|
|
assert len(visual_kwargs["images"]) == 3
|
|
assert visual_text.count("<|vision_start|>") == 3
|
|
assert "low-weight style calibrator" in visual_text
|
|
assert "Suppress depicted people" in visual_text
|
|
assert len(conditioning) == 2
|
|
metadata = conditioning[1][1]["krea2_visual_ref"]
|
|
assert metadata["style_image_count"] == 3
|
|
assert metadata["visual_placeholders"] == 3
|
|
assert metadata["style_text_fusion"] == "model_synthesized"
|
|
assert metadata["strength_mode"] == "conditioning_strength_text_visual_balanced"
|
|
assert conditioning[1][1]["strength"] == 0.75 * 0.54
|
|
|
|
def test_style_prompt_encode_extracts_style_then_encodes_text_only_prompt():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2StylePromptEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
conditioning, style_prompt = node.encode(
|
|
clip=clip,
|
|
style_image=image,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
style_strength=0.65,
|
|
)
|
|
|
|
assert len(clip.tokenize_calls) == 2
|
|
analysis_text, analysis_kwargs = clip.tokenize_calls[0]
|
|
final_text, final_kwargs = clip.tokenize_calls[1]
|
|
assert analysis_kwargs["images"] == [image]
|
|
assert "transferable visual style" in analysis_text
|
|
assert final_kwargs == {}
|
|
assert "a ceramic teapot on a wooden table" in final_text
|
|
assert "soft watercolor storybook illustration" in final_text
|
|
assert "TRANSFERABLE STYLE LOCK" in final_text
|
|
assert "Mandatory style-transfer instruction" in final_text
|
|
assert "soft watercolor storybook illustration" in style_prompt
|
|
assert conditioning[0][0].shape == (1, 42, 30720)
|
|
|
|
|
|
def test_style_prompt_strength_zero_uses_plain_text_prompt():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2StylePromptEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
conditioning, style_prompt = node.encode(
|
|
clip=clip,
|
|
style_image=image,
|
|
prompt="plain text prompt",
|
|
style_strength=0.0,
|
|
)
|
|
|
|
assert clip.tokenize_calls == [("plain text prompt", {})]
|
|
assert style_prompt == "plain text prompt"
|
|
assert conditioning[0][0].shape == (1, 42, 30720)
|
|
|
|
|
|
def test_encode_without_image_matches_native_clip_path():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeNativeClip()
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="plain text prompt",
|
|
strength=1.0,
|
|
placement="before_prompt",
|
|
)
|
|
|
|
assert clip.tokenize_calls == [("plain text prompt", {})]
|
|
assert len(clip.encode_calls) == 1
|
|
assert conditioning[0][0].shape == (1, 7, 4096)
|
|
assert conditioning[0][1] == {"native": True}
|
|
assert "krea2_visual_ref" not in conditioning[0][1]
|
|
|
|
|
|
def test_visual_strength_does_not_scale_embeddings():
|
|
import torch
|
|
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2TensorClip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="a ceramic teapot on a wooden table",
|
|
strength=0.5,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="visual_only",
|
|
)
|
|
|
|
visual_cond, metadata = conditioning[0]
|
|
assert torch.all(visual_cond == 1.0)
|
|
assert metadata["strength"] == 0.5 ** 2.5
|
|
assert metadata["krea2_visual_ref"]["strength"] == 0.5
|
|
assert metadata["krea2_visual_ref"]["effective_strength"] == 0.5 ** 2.5
|
|
assert metadata["krea2_visual_ref"]["strength_mode"] == "conditioning_strength_power_2_5"
|
|
assert metadata["krea2_visual_ref"]["visual_span"] == [1, 4]
|
|
|
|
|
|
|
|
def test_visual_reference_rejects_empty_prompt_with_image():
|
|
import pytest
|
|
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeKrea2Clip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
with pytest.raises(RuntimeError, match="non-empty generation prompt"):
|
|
node.encode(
|
|
clip=clip,
|
|
prompt=" ",
|
|
strength=0.75,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
style_text_mode="visual_only",
|
|
)
|
|
|
|
def test_visual_strength_zero_uses_native_text_path_even_with_image():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceEncode"]()
|
|
clip = FakeNativeClip()
|
|
image = FakeImage((1, 64, 64, 3))
|
|
|
|
(conditioning,) = node.encode(
|
|
clip=clip,
|
|
prompt="plain text prompt",
|
|
strength=0.0,
|
|
style_image=image,
|
|
placement="before_prompt",
|
|
)
|
|
|
|
assert clip.tokenize_calls == [("plain text prompt", {})]
|
|
assert conditioning[0][1] == {"native": True}
|
|
|
|
|
|
def test_debug_reports_valid_and_invalid_widths():
|
|
plugin = load_plugin()
|
|
node = plugin.NODE_CLASS_MAPPINGS["Krea2VisualReferenceDebug"]()
|
|
|
|
(valid_info,) = node.debug([[FakeTensor((1, 12, 30720)), {"krea2_visual_ref": {"source": "test"}}]])
|
|
assert "krea2_width_valid: True" in valid_info
|
|
assert "sequence_length: 12" in valid_info
|
|
|
|
(invalid_info,) = node.debug([[FakeTensor((1, 12, 4096)), {}]])
|
|
assert "krea2_width_valid: False" in invalid_info
|
|
assert "cannot be reliably inferred" in invalid_info
|
|
|
|
def test_real_krea2_tokenizer_keeps_image_token_after_prefix_strip():
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
|
|
import torch
|
|
from comfy.text_encoders.krea2 import Krea2Tokenizer
|
|
|
|
plugin = load_plugin()
|
|
visual_ref = importlib.import_module(f"{plugin.__name__}.nodes.visual_ref")
|
|
|
|
image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
template = visual_ref._build_krea2_visual_template(
|
|
"a ceramic teapot on a wooden table",
|
|
"before_prompt",
|
|
)
|
|
|
|
tokens = Krea2Tokenizer().tokenize_with_weights(template, images=[image])
|
|
row = tokens["qwen3vl_4b"][0]
|
|
|
|
image_indices = [
|
|
index
|
|
for index, item in enumerate(row)
|
|
if isinstance(item[0], dict) and item[0].get("type") == "image"
|
|
]
|
|
assert len(image_indices) == 1
|
|
|
|
count_im_start = 0
|
|
template_end = -1
|
|
for index, item in enumerate(row):
|
|
token = item[0]
|
|
if isinstance(token, int) and token == 151644 and count_im_start < 2:
|
|
template_end = index
|
|
count_im_start += 1
|
|
assert count_im_start == 2
|
|
if row[template_end + 1][0] == 872 and row[template_end + 2][0] == 198:
|
|
template_end += 3
|
|
|
|
assert image_indices[0] >= template_end
|
|
|
|
def test_real_krea2_tokenizer_native_prompt_matches_official_text_template():
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
|
|
from comfy.text_encoders.krea2 import KREA2_TEMPLATE, Krea2Tokenizer
|
|
|
|
prompt = "a ceramic teapot on a wooden table"
|
|
tokenizer = Krea2Tokenizer()
|
|
native = tokenizer.tokenize_with_weights(prompt)["qwen3vl_4b"][0]
|
|
explicit = tokenizer.tokenize_with_weights(KREA2_TEMPLATE.format(prompt))["qwen3vl_4b"][0]
|
|
|
|
assert [item[0] for item in native] == [item[0] for item in explicit]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|