diff --git a/__init__.py b/__init__.py index 1ab7708..5bfc0de 100644 --- a/__init__.py +++ b/__init__.py @@ -45,6 +45,7 @@ _SUBPACKAGES: list[str] = [ "conditioning", # 第 2 节：文本编码与条件生成 "sampling", # 第 3 节：图像采样与去噪 "qwen", # 第 4 节：Qwen 一体化节点 + "flux2klein", # 第 4 节（Flux2Klein 子域）：编辑文本编码 "latent", # 第 4 节：Latent 空间操作 "image", # 第 5 节：图像处理与后处理 "repair", # 第 5 节（修复子域）：图像修复 diff --git a/docs/superpowers/plans/2026-05-26-lls-concat-by-target-two-port-ui.md b/docs/superpowers/plans/2026-05-26-lls-concat-by-target-two-port-ui.md new file mode 100644 index 0000000..59bbcfa --- /dev/null +++ b/docs/superpowers/plans/2026-05-26-lls-concat-by-target-two-port-ui.md @@ -0,0 +1,367 @@ +# LLS Concat By Target Two-Port UI Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Change `LLS Concat By Target` so the ComfyUI node shows exactly two inputs labeled `image/mask_A` and `image/mask_B`, while still accepting either `IMAGE` or `MASK` on each side and preserving current concat behavior. + +**Architecture:** Revert the backend schema to two internal wildcard inputs, `a` and `b`, because the `IMAGE,MASK` union type is what causes ComfyUI to render four visible sockets. Add a tiny `web/js` extension that relabels those two internal inputs to the user-facing names `image/mask_A` and `image/mask_B`, while keeping the existing runtime auto-detection and legacy keyword aliases intact. + +**Tech Stack:** Python, unittest, torch, ComfyUI node contracts, lightweight frontend extension JavaScript + +--- + +## File Map + +- `utils/concat_by_target.py` + Restore the schema to two wildcard inputs and keep runtime support for both internal `a` / `b` and legacy `image/mask_A` / `image/mask_B` keyword aliases. +- `web/js/lls_concat_by_target.js` + Register a ComfyUI extension that renames the visible labels for the `a` and `b` inputs. +- `tests/test_concat_by_target_registration.py` + Lock the backend schema contract to two wildcard inputs. +- `tests/test_concat_by_target_node.py` + Keep the runtime contract explicit, including legacy keyword support and the existing image/mask auto-detection behavior. +- `tests/test_concat_by_target_frontend.py` + Lock the new frontend asset and its label-rewrite contract. + +### Task 1: Lock The Two-Port Contract With Failing Tests + +**Files:** +- Modify: `tests/test_concat_by_target_registration.py` +- Create: `tests/test_concat_by_target_frontend.py` + +- [ ] **Step 1: Update the registration test to expect two internal inputs** + +```python +import unittest + +try: + from .test_mask_draw_helpers import load_plugin_package +except ImportError: + from test_mask_draw_helpers import load_plugin_package + + +class TestConcatByTargetRegistration(unittest.TestCase): + def test_plugin_registers_concat_by_target_node(self): + plugin = load_plugin_package() + + self.assertIn("LLSConcatByTarget", plugin.NODE_CLASS_MAPPINGS) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSConcatByTarget"], + "LLS Concat By Target", + ) + + def test_concat_by_target_schema_matches_contract(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSConcatByTarget"] + schema = node_cls.INPUT_TYPES() + required = schema["required"] + optional = schema["optional"] + + self.assertEqual(node_cls.CATEGORY, "LLS/Utils") + self.assertEqual(node_cls.FUNCTION, "concat") + self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "MASK", "INT", "INT")) + self.assertEqual(node_cls.RETURN_NAMES, ("image", "mask", "width", "height")) + self.assertEqual(required["data_type"][0], ["IMAGE", "MASK"]) + self.assertEqual(required["target"][0], ["A", "B"]) + self.assertEqual(required["position"][0], ["top", "bottom", "left", "right"]) + self.assertEqual(required["resize_mode"][0], ["keep_proportion", "stretch", "none"]) + self.assertEqual(required["align"][0], ["start", "center", "end"]) + self.assertEqual(required["match_target_size"][0], "BOOLEAN") + self.assertEqual(required["gap"][0], "INT") + self.assertEqual(required["gap"][1]["default"], 0) + self.assertEqual(required["background_color"][0], "STRING") + self.assertEqual(required["background_color"][1]["default"], "#000000") + self.assertEqual(required["background_value"][0], "FLOAT") + self.assertEqual(required["background_value"][1]["default"], 0.0) + self.assertEqual(required["multiple_of"][0], "INT") + self.assertEqual(required["allow_batch_broadcast"][0], "BOOLEAN") + self.assertEqual(str(optional["a"][0]), "*") + self.assertEqual(str(optional["b"][0]), "*") + self.assertNotIn("image/mask_A", optional) + self.assertNotIn("image/mask_B", optional) + self.assertNotIn("image_a", optional) + self.assertNotIn("image_b", optional) + self.assertNotIn("mask_a", optional) + self.assertNotIn("mask_b", optional) + + +if __name__ == "__main__": + unittest.main() +``` + +- [ ] **Step 2: Add a failing frontend-contract test for the label adapter** + +```python +import unittest + +try: + from .test_mask_draw_helpers import ROOT, load_plugin_package +except ImportError: + from test_mask_draw_helpers import ROOT, load_plugin_package + + +class TestConcatByTargetFrontend(unittest.TestCase): + def test_plugin_exports_web_directory(self): + plugin = load_plugin_package() + self.assertEqual(plugin.WEB_DIRECTORY, "./web") + + def test_frontend_asset_exists(self): + asset = ROOT / "web" / "js" / "lls_concat_by_target.js" + self.assertTrue(asset.exists(), msg=f"Missing frontend asset: {asset}") + + def test_frontend_asset_registers_concat_by_target_extension(self): + asset = (ROOT / "web" / "js" / "lls_concat_by_target.js").read_text(encoding="utf-8") + + self.assertIn("app.registerExtension", asset) + self.assertIn("LLSConcatByTarget", asset) + self.assertIn("LLS Concat By Target", asset) + self.assertIn("image/mask_A", asset) + self.assertIn("image/mask_B", asset) + self.assertIn("beforeRegisterNodeDef", asset) + self.assertIn("onNodeCreated", asset) + self.assertIn("onGraphConfigured", asset) + self.assertIn("onConnectionsChange", asset) + self.assertIn("input.label", asset) + + +if __name__ == "__main__": + unittest.main() +``` + +- [ ] **Step 3: Run the contract tests to verify they fail** + +Run: `python3 -m unittest tests.test_concat_by_target_registration tests.test_concat_by_target_frontend -v` +Expected: FAIL because the schema still exposes `image/mask_A` / `image/mask_B` and `web/js/lls_concat_by_target.js` does not exist yet. + + +### Task 2: Restore The Backend To Two Internal Inputs + +**Files:** +- Modify: `utils/concat_by_target.py` +- Modify: `tests/test_concat_by_target_node.py` +- Test: `tests/test_concat_by_target_registration.py` +- Test: `tests/test_concat_by_target_node.py` + +- [ ] **Step 1: Rename the legacy-keyword runtime test so its purpose stays explicit after the schema change** + +```python +def test_legacy_port_names_are_accepted(self): + image_a = self.make_image(width=2, height=2, value=0.2) + image_b = self.make_image(width=1, height=2, value=0.8) + + image, mask, width, height = self.node.concat( + data_type="IMAGE", + target="A", + position="right", + match_target_size=True, + resize_mode="none", + align="center", + gap=0, + background_color="#000000", + background_value=0.0, + multiple_of=0, + allow_batch_broadcast=True, + **{"image/mask_A": image_a, "image/mask_B": image_b}, + ) + + self.assert_image_output(image, mask, width, height, (1, 2, 3, 3)) + self.assertTrue(torch.allclose(image[:, :, 0:2, :], torch.full((1, 2, 2, 3), 0.2))) + self.assertTrue(torch.allclose(image[:, :, 2:3, :], torch.full((1, 2, 1, 3), 0.8))) +``` + +- [ ] **Step 2: Change the node schema to wildcard `a` / `b` inputs while keeping legacy alias resolution** + +```python +DATA_TYPE_CHOICES = ["IMAGE", "MASK"] +TARGET_CHOICES = ["A", "B"] +POSITION_CHOICES = ["top", "bottom", "left", "right"] +RESIZE_MODE_CHOICES = ["keep_proportion", "stretch", "none"] +ALIGN_CHOICES = ["start", "center", "end"] +PORT_A_NAME = "image/mask_A" +PORT_B_NAME = "image/mask_B" + + +class AnyType(str): + def __ne__(self, other): # pragma: no cover - exercised by ComfyUI type checks + return False + + +ANY_TYPE = AnyType("*") +``` + +```python +@classmethod +def INPUT_TYPES(cls): + return { + "required": { + "data_type": (DATA_TYPE_CHOICES, {"default": "IMAGE"}), + "target": (TARGET_CHOICES, {"default": "A"}), + "position": (POSITION_CHOICES, {"default": "right"}), + "match_target_size": ("BOOLEAN", {"default": True}), + "resize_mode": (RESIZE_MODE_CHOICES, {"default": "keep_proportion"}), + "align": (ALIGN_CHOICES, {"default": "center"}), + "gap": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}), + "background_color": ("STRING", {"default": "#000000", "multiline": False}), + "background_value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "multiple_of": ("INT", {"default": 0, "min": 0, "max": 512, "step": 1}), + "allow_batch_broadcast": ("BOOLEAN", {"default": True}), + }, + "optional": { + "a": (ANY_TYPE,), + "b": (ANY_TYPE,), + }, + } +``` + +```python +def concat( + self, + data_type, + target, + position, + match_target_size, + resize_mode, + align, + gap, + background_color, + background_value, + multiple_of, + allow_batch_broadcast, + a=None, + b=None, + **kwargs, +): + _require_torch() + data_type = str(data_type or "IMAGE") + background_value = max(0.0, min(1.0, float(background_value))) + input_a = kwargs[PORT_A_NAME] if PORT_A_NAME in kwargs else a + input_b = kwargs[PORT_B_NAME] if PORT_B_NAME in kwargs else b + resolved_data_type = resolve_concat_data_type(input_a, input_b, data_type) +``` + +- [ ] **Step 3: Run the backend tests to verify the two-port schema and runtime behavior pass** + +Run: `python3 -m unittest tests.test_concat_by_target_registration tests.test_concat_by_target_node -v` +Expected: PASS, including the legacy keyword test and the existing image/mask behavior tests. + +- [ ] **Step 4: Commit the backend schema change** + +```bash +git add utils/concat_by_target.py tests/test_concat_by_target_registration.py tests/test_concat_by_target_node.py +git commit -m "refactor: restore two-port concat backend schema" +``` + + +### Task 3: Add The Frontend Label Adapter + +**Files:** +- Create: `web/js/lls_concat_by_target.js` +- Test: `tests/test_concat_by_target_frontend.py` +- Test: `tests/test_concat_by_target_registration.py` + +- [ ] **Step 1: Implement the frontend extension that relabels `a` and `b`** + +```javascript +import { app } from "../../scripts/app.js"; + +const EXTENSION_NAME = "LLS.ConcatByTarget"; +const TARGET_NODE_CLASS = "LLSConcatByTarget"; +const TARGET_NODE_DISPLAY_NAME = "LLS Concat By Target"; +const INPUT_LABELS = { + a: "image/mask_A", + b: "image/mask_B", +}; + +function relabelConcatInputs(node) { + let changed = false; + + for (const input of node.inputs ?? []) { + const targetLabel = INPUT_LABELS[input.name]; + if (targetLabel && input.label !== targetLabel) { + input.label = targetLabel; + changed = true; + } + } + + if (changed) { + node.setDirtyCanvas?.(true, true); + } +} + +app.registerExtension({ + name: EXTENSION_NAME, + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== TARGET_NODE_CLASS && nodeData.display_name !== TARGET_NODE_DISPLAY_NAME) { + return; + } + + const previousOnNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function onConcatByTargetNodeCreated() { + const result = previousOnNodeCreated?.apply(this, arguments); + relabelConcatInputs(this); + return result; + }; + + const previousOnGraphConfigured = nodeType.prototype.onGraphConfigured; + nodeType.prototype.onGraphConfigured = function onConcatByTargetGraphConfigured() { + const result = previousOnGraphConfigured?.apply(this, arguments); + relabelConcatInputs(this); + return result; + }; + + const previousOnConnectionsChange = nodeType.prototype.onConnectionsChange; + nodeType.prototype.onConnectionsChange = function onConcatByTargetConnectionsChange() { + const result = previousOnConnectionsChange?.apply(this, arguments); + relabelConcatInputs(this); + return result; + }; + }, +}); +``` + +- [ ] **Step 2: Run the frontend and contract tests to verify the label adapter passes** + +Run: `python3 -m unittest tests.test_concat_by_target_frontend tests.test_concat_by_target_registration -v` +Expected: PASS because the asset exists, registers the extension, and the backend contract now exposes two internal inputs. + +- [ ] **Step 3: Commit the frontend label adapter** + +```bash +git add web/js/lls_concat_by_target.js tests/test_concat_by_target_frontend.py +git commit -m "feat: relabel concat by target ports in frontend" +``` + + +### Task 4: Verify The Final Behavior End-To-End + +**Files:** +- Test: `tests/test_concat_by_target_registration.py` +- Test: `tests/test_concat_by_target_node.py` +- Test: `tests/test_concat_by_target_frontend.py` + +- [ ] **Step 1: Run the full concat test slice** + +Run: `python3 -m unittest tests.test_concat_by_target_registration tests.test_concat_by_target_node tests.test_concat_by_target_frontend -v` +Expected: PASS with the full concat registration, runtime, and frontend suite green. + +- [ ] **Step 2: Perform a manual ComfyUI smoke check** + +Manual: +- Restart or reload ComfyUI so the new `web/js/lls_concat_by_target.js` asset is loaded. +- Create a fresh `LLS Concat By Target` node. +- Confirm the node shows exactly two visible input sockets. +- Confirm their labels are `image/mask_A` and `image/mask_B`. +- Confirm an `IMAGE` can be connected to each socket and the node still executes. +- Confirm a `MASK` can be connected to each socket and the node still executes. + +- [ ] **Step 3: Check old-workflow compatibility explicitly** + +Manual: +- If you have an existing workflow JSON that already contains `LLS Concat By Target`, load it after the backend schema change. +- Confirm whether the old links reconnect automatically. +- If they do not reconnect, record the limitation in the implementation handoff: old graphs need manual reconnection even though direct Python calls still accept `image/mask_A` / `image/mask_B` as legacy keyword aliases. + +- [ ] **Step 4: Leave the worktree clean before merge or review** + +Run: `git status --short` +Expected: no output. diff --git a/flux2klein/__init__.py b/flux2klein/__init__.py new file mode 100644 index 0000000..7c270b5 --- /dev/null +++ b/flux2klein/__init__.py @@ -0,0 +1,8 @@ +""" +node.flux2klein +================ +功能域：Flux2Klein / 图像编辑文本编码 +""" +from .lls_flux2klein_edit_text_encode import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/flux2klein/lls_flux2klein_edit_text_encode.py b/flux2klein/lls_flux2klein_edit_text_encode.py new file mode 100644 index 0000000..5d501a9 --- /dev/null +++ b/flux2klein/lls_flux2klein_edit_text_encode.py @@ -0,0 +1,413 @@ +from __future__ import annotations + +import importlib + +from ..conditioning.nodes import LLSUniversalPromptEncode +from ..repair.repair_utils import ( + clamp_mask_values, + get_image_size, + invert_mask_values, + normalize_mask_values, + resize_image_to, + resize_mask_to, +) + +try: + import torch +except Exception as exc: # pragma: no cover - optional runtime dependency + torch = None + _TORCH_ERR = exc +else: # pragma: no cover - exercised in ComfyUI runtime/tests when torch exists + _TORCH_ERR = None + +try: + nodes_qwen = importlib.import_module("comfy_extras.nodes_qwen") +except Exception as exc: # pragma: no cover - runtime-only import + nodes_qwen = None + _QWEN_ERR = exc +else: # pragma: no cover - runtime-only import + _QWEN_ERR = None + +try: + nodes_flux = importlib.import_module("comfy_extras.nodes_flux") +except Exception as exc: # pragma: no cover - runtime-only import + nodes_flux = None + _FLUX_ERR = exc +else: # pragma: no cover - runtime-only import + _FLUX_ERR = None + + +RESIZE_MODE_CHOICES = ["longest_edge", "keep_original"] +MASK_MODE_CHOICES = ["none", "use_mask", "invert_mask"] +CUSTOM_OUTPUT_TYPE = "LLS_FLUX2KLEIN_OUTPUT" +DEFAULT_REFERENCE_LATENTS_METHOD = "index_timestep_zero" + + +class _ConstantMask: + """Fallback mask object for non-Comfy test/runtime contexts without torch.""" + + def __init__(self, shape, value: float = 1.0, label: str = "mask"): + self.shape = tuple(shape) + self.value = float(value) + self.label = label + width = int(self.shape[2]) + height = int(self.shape[1]) + if self.value <= 0.0: + self.mask_bbox = None + self.mask_area_ratio = 0.0 + else: + self.mask_bbox = (0, 0, width, height) + self.mask_area_ratio = max(0.0, min(1.0, self.value)) + + def resized(self, width, height): + return _ConstantMask((self.shape[0], int(height), int(width)), value=self.value, label=f"{self.label}:resized") + + def normalized(self): + return _ConstantMask(self.shape, value=max(0.0, min(1.0, self.value)), label=f"{self.label}:normalized") + + def inverted(self, image_size): + width, height = image_size + return _ConstantMask( + (self.shape[0], int(height), int(width)), + value=max(0.0, min(1.0, 1.0 - self.value)), + label=f"{self.label}:inverted", + ) + + +def _unwrap_first(result): + if hasattr(result, "result"): + values = result.result + if isinstance(values, tuple): + return values[0] + if isinstance(values, list): + return values[0] + return values + if isinstance(result, tuple): + return result[0] + if isinstance(result, list): + return result[0] + return result + + +def _get_image_batch_size(image) -> int: + shape = tuple(getattr(image, "shape", ())) + if len(shape) != 4: + raise RuntimeError("[LLS] image1 must have shape [batch, height, width, channels].") + batch = int(shape[0]) + if batch <= 0: + raise RuntimeError("[LLS] image1 batch size must be positive.") + return batch + + +def _scale_size_by_longest_edge(width: int, height: int, target_longest_edge: int) -> tuple[int, int]: + width = int(width) + height = int(height) + longest = max(width, height) + if width <= 0 or height <= 0 or longest <= 0: + raise RuntimeError("[LLS] IMAGE width and height must be positive.") + + scale = float(target_longest_edge) / float(longest) + target_width = max(1, int(round(width * scale))) + target_height = max(1, int(round(height * scale))) + return target_width, target_height + + +def _prepare_image(image, resize_mode: str, ref_longest_edge: int): + if image is None: + return None + if resize_mode == "keep_original": + return image + if resize_mode != "longest_edge": + raise RuntimeError(f"[LLS] Unsupported resize_mode '{resize_mode}'.") + width, height = get_image_size(image) + target_width, target_height = _scale_size_by_longest_edge(width, height, int(ref_longest_edge)) + return resize_image_to(image, target_width, target_height) + + +def _make_full_mask_like_image(image): + batch = _get_image_batch_size(image) + width, height = get_image_size(image) + + if torch is not None: + device = getattr(image, "device", "cpu") + try: + return torch.ones((batch, height, width), dtype=torch.float32, device=device) + except Exception: + return torch.ones((batch, height, width), dtype=torch.float32) + + return _ConstantMask((batch, height, width), value=1.0, label="full-mask") + + +def _broadcast_mask_batch_if_needed(mask, image): + if torch is None or not isinstance(mask, torch.Tensor) or not isinstance(image, torch.Tensor): + return mask + + image_batch = int(image.shape[0]) + mask_batch = int(mask.shape[0]) + if image_batch == mask_batch: + return mask + if mask_batch == 1 and image_batch > 1: + return mask.repeat(image_batch, 1, 1) + raise RuntimeError( + f"[LLS] MASK batch size {mask_batch} does not match image1 batch size {image_batch}." + ) + + +def _prepare_mask(mask, image1, mask_mode: str): + # mask only applies to the main edit image (image1), never to image2/image3. + if mask_mode == "none" or mask is None: + return _make_full_mask_like_image(image1) + + width, height = get_image_size(image1) + processed = resize_mask_to(mask, width, height) + processed = normalize_mask_values(processed) + if mask_mode == "invert_mask": + processed = invert_mask_values(processed, (width, height)) + elif mask_mode != "use_mask": + raise RuntimeError(f"[LLS] Unsupported mask_mode '{mask_mode}'.") + processed = clamp_mask_values(processed) + return _broadcast_mask_batch_if_needed(processed, image1) + + +def _has_native_qwen_flux_runtime() -> bool: + return ( + nodes_qwen is not None + and getattr(nodes_qwen, "TextEncodeQwenImageEditPlus", None) is not None + ) + + +def _encode_with_native_qwen_flux(clip, vae, prompt: str, image1, image2=None, image3=None): + encoder_cls = getattr(nodes_qwen, "TextEncodeQwenImageEditPlus", None) + if encoder_cls is None: + raise RuntimeError("[LLS] Native Qwen edit encoder is unavailable.") from _QWEN_ERR + + if hasattr(encoder_cls, "execute"): + conditioning = _unwrap_first( + encoder_cls.execute( + clip, + prompt, + vae=vae, + image1=image1, + image2=image2, + image3=image3, + ) + ) + else: + encoder = encoder_cls() + conditioning = _unwrap_first( + encoder.execute( + clip, + prompt, + vae=vae, + image1=image1, + image2=image2, + image3=image3, + ) + ) + + # When reference images are present, append the standard Flux Kontext method tag + # so downstream KSampler paths can preserve multi-reference intent. + if (image2 is not None or image3 is not None) and nodes_flux is not None: + method_cls = getattr(nodes_flux, "FluxKontextMultiReferenceLatentMethod", None) + if method_cls is not None: + if hasattr(method_cls, "execute"): + conditioning = _unwrap_first( + method_cls.execute(conditioning, DEFAULT_REFERENCE_LATENTS_METHOD) + ) + else: + method_node = method_cls() + conditioning = _unwrap_first( + method_node.execute(conditioning, DEFAULT_REFERENCE_LATENTS_METHOD) + ) + + return conditioning + + +def _encode_with_lls_fallback(clip, prompt: str): + try: + conditioning, _negative, _prompt_info = LLSUniversalPromptEncode().encode( + text_encoder=clip, + positive_prompt=prompt, + negative_prompt="", + clip_skip=-1, + model_info=None, + ) + return conditioning + except Exception: + tokenize = getattr(clip, "tokenize", None) + if not callable(tokenize): + raise RuntimeError( + "[LLS] No usable text encoding backend for LLS Flux2Klein Edit Text Encode. " + "Current clip object does not expose tokenize()/encode methods." + ) + + tokens = tokenize(prompt) + scheduled = getattr(clip, "encode_from_tokens_scheduled", None) + if callable(scheduled): + try: + return scheduled(tokens) + except Exception as exc: + raise RuntimeError( + "[LLS] No usable text encoding backend for LLS Flux2Klein Edit Text Encode. " + f"Scheduled CLIP encoding failed: {exc}" + ) from exc + + plain_encode = getattr(clip, "encode_from_tokens", None) + if not callable(plain_encode): + raise RuntimeError( + "[LLS] No usable text encoding backend for LLS Flux2Klein Edit Text Encode. " + "Connect a compatible CLIP/Qwen encoder or run inside a ComfyUI build with the official edit encoders." + ) + + try: + encoded = plain_encode(tokens, return_pooled=True, return_dict=True) + except TypeError: + encoded = plain_encode(tokens, return_pooled=True) + except Exception as exc: + raise RuntimeError( + "[LLS] No usable text encoding backend for LLS Flux2Klein Edit Text Encode. " + f"Plain CLIP encoding failed: {exc}" + ) from exc + + if isinstance(encoded, dict) and "cond" in encoded: + meta = dict(encoded) + cond = meta.pop("cond") + return [[cond, meta]] + if isinstance(encoded, tuple) and len(encoded) >= 2: + cond, pooled = encoded[:2] + return [[cond, {"pooled_output": pooled}]] + return [[encoded, {}]] + + +def _encode_conditioning(clip, vae, prompt: str, image1, image2=None, image3=None): + native_error = None + if _has_native_qwen_flux_runtime(): + try: + return _encode_with_native_qwen_flux( + clip, + vae, + prompt, + image1=image1, + image2=image2, + image3=image3, + ), "native_qwen_flux", None + except Exception as exc: + native_error = str(exc) + + try: + return _encode_with_lls_fallback(clip, prompt), "lls_fallback", native_error + except Exception as exc: + message = str(exc) + if native_error: + message = f"{message} Native encoder fallback reason: {native_error}" + raise RuntimeError(message) from exc + + +class LLSFlux2KleinEditTextEncode: + CATEGORY = "LLS/Flux2Klein" + FUNCTION = "encode" + RETURN_TYPES = ("CONDITIONING", "LATENT", CUSTOM_OUTPUT_TYPE, "IMAGE", "MASK") + RETURN_NAMES = ("conditioning", "latent", "custom_output", "main_image", "mask") + DESCRIPTION = ( + "Flux2Klein-style edit text encode node for LLS. Uses image1 as the main edit image, " + "image2/image3 as reference images, never concatenates them, and prepares conditioning/latent/mask outputs." + ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip": ("CLIP",), + "vae": ("VAE",), + "image1": ("IMAGE",), + "prompt": ("STRING", {"default": "", "multiline": True}), + "ref_longest_edge": ("INT", {"default": 1024, "min": 256, "max": 4096, "step": 64}), + "resize_mode": (RESIZE_MODE_CHOICES, {"default": "longest_edge"}), + "mask_mode": (MASK_MODE_CHOICES, {"default": "none"}), + }, + "optional": { + "image2": ("IMAGE",), + "image3": ("IMAGE",), + "mask": ("MASK",), + }, + } + + def encode( + self, + clip, + vae, + image1, + prompt, + ref_longest_edge, + resize_mode, + mask_mode, + image2=None, + image3=None, + mask=None, + ): + if clip is None: + raise RuntimeError("[LLS] Missing CLIP. 请连接可用的 clip 文本编码器。") + if vae is None: + raise RuntimeError("[LLS] Missing VAE. 请连接可用的 vae。") + if image1 is None: + raise RuntimeError("[LLS] Missing image1. image1 是主编辑图 / main edit image。") + + main_image = _prepare_image(image1, str(resize_mode), int(ref_longest_edge)) + reference_image_2 = _prepare_image(image2, str(resize_mode), int(ref_longest_edge)) + reference_image_3 = _prepare_image(image3, str(resize_mode), int(ref_longest_edge)) + output_mask = _prepare_mask(mask, main_image, str(mask_mode)) + + conditioning, conditioning_backend, native_error = _encode_conditioning( + clip, + vae, + str(prompt), + main_image, + image2=reference_image_2, + image3=reference_image_3, + ) + + # latent only encodes image1, because image2/image3 are reference images rather than target latents. + try: + latent_tensor = vae.encode(main_image) + except Exception as exc: + raise RuntimeError( + f"[LLS] VAE encode failed for image1 / 主编辑图编码失败: {exc}" + ) from exc + + reference_names = [] + if reference_image_2 is not None: + reference_names.append("image2") + if reference_image_3 is not None: + reference_names.append("image3") + + custom_output = { + "node_name": "LLS Flux2Klein Edit Text Encode", + "prompt": str(prompt), + "ref_longest_edge": int(ref_longest_edge), + "resize_mode": str(resize_mode), + "mask_mode": str(mask_mode), + "has_mask": mask is not None, + "main_image": "image1", + "reference_images": reference_names, + "num_reference_images": len(reference_names), + "conditioning_backend": conditioning_backend, + "native_error": native_error, + "note": "image1 is main edit image; image2 and image3 are reference images; images are not concatenated", + } + + return ( + conditioning, + {"samples": latent_tensor}, + custom_output, + main_image, + output_mask, + ) + + +NODE_CLASS_MAPPINGS = { + "LLSFlux2KleinEditTextEncode": LLSFlux2KleinEditTextEncode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LLSFlux2KleinEditTextEncode": "LLS Flux2Klein Edit Text Encode", +} diff --git a/tests/test_flux2klein_helpers.py b/tests/test_flux2klein_helpers.py new file mode 100644 index 0000000..2e7b8b2 --- /dev/null +++ b/tests/test_flux2klein_helpers.py @@ -0,0 +1,150 @@ +import importlib +import importlib.util +import pathlib +import sys + + +ROOT = pathlib.Path(__file__).resolve().parents[1] +MODULE_NAME = "lls_node_test_flux2klein" + + +def load_plugin_package(): + for name in list(sys.modules): + if name == MODULE_NAME or name.startswith(f"{MODULE_NAME}."): + sys.modules.pop(name) + + spec = importlib.util.spec_from_file_location( + MODULE_NAME, + ROOT / "__init__.py", + submodule_search_locations=[str(ROOT)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[MODULE_NAME] = module + spec.loader.exec_module(module) + return module + + +def import_plugin_submodule(plugin, dotted_name: str): + return importlib.import_module(f"{plugin.__name__}.{dotted_name}") + + +class FakeTensor: + def __init__(self, shape, label="image"): + self.shape = tuple(shape) + self.label = label + + def resized(self, width, height): + batch, _old_h, _old_w, channels = self.shape + return FakeTensor((batch, height, width, channels), label=f"{self.label}:resized[{width}x{height}]") + + +class FakeMask: + def __init__(self, shape, mask_bbox=None, mask_area_ratio=0.0, label="mask"): + self.shape = tuple(shape) + self.mask_bbox = mask_bbox + self.mask_area_ratio = float(mask_area_ratio) + self.label = label + + def resized(self, width, height): + batch, _old_h, _old_w = self.shape + return FakeMask( + (batch, height, width), + mask_bbox=self.mask_bbox, + mask_area_ratio=self.mask_area_ratio, + label=f"{self.label}:resized[{width}x{height}]", + ) + + def normalized(self): + return FakeMask( + self.shape, + mask_bbox=self.mask_bbox, + mask_area_ratio=self.mask_area_ratio, + label=f"{self.label}:normalized", + ) + + def inverted(self, image_size): + width, height = image_size + ratio = max(0.0, min(1.0, 1.0 - self.mask_area_ratio)) + bbox = None if ratio <= 0.0 else (0, 0, width, height) + return FakeMask( + (self.shape[0], height, width), + mask_bbox=bbox, + mask_area_ratio=ratio, + label=f"{self.label}:inverted", + ) + + +class FakeLatentTensor: + def __init__(self, shape): + self.shape = tuple(shape) + + +class FakeVAE: + def __init__(self, latent_channels=4, downscale_ratio=8): + self.latent_channels = int(latent_channels) + self.downscale_ratio = int(downscale_ratio) + self.encoded_shapes = [] + + def encode(self, image): + self.encoded_shapes.append(tuple(image.shape)) + batch, height, width, _channels = image.shape + latent_height = max(1, height // self.downscale_ratio) + latent_width = max(1, width // self.downscale_ratio) + return FakeLatentTensor((batch, self.latent_channels, latent_height, latent_width)) + + +class StandardClipStub: + def __init__(self): + self.tokenize_calls = [] + self.encode_calls = [] + + def tokenize(self, text): + self.tokenize_calls.append(text) + return {"text": text} + + def encode_from_tokens_scheduled(self, tokens, add_dict=None): + self.encode_calls.append({"tokens": dict(tokens), "add_dict": dict(add_dict or {})}) + return [[f"cond::{tokens['text']}", {"pooled_output": dict(add_dict or {})}]] + + +class BrokenClipStub: + pass + + +class NodeOutputStub: + def __init__(self, *args): + self.args = args + + @property + def result(self): + return self.args + + +class NativeQwenStub: + last_call = None + + class TextEncodeQwenImageEditPlus: + @staticmethod + def execute(clip, prompt, vae=None, image1=None, image2=None, image3=None): + NativeQwenStub.last_call = { + "clip": clip, + "prompt": prompt, + "vae": vae, + "image1": image1, + "image2": image2, + "image3": image3, + } + return NodeOutputStub([["native-conditioning", {"source": "native-qwen"}]]) + + +class NativeFluxStub: + last_call = None + + class FluxKontextMultiReferenceLatentMethod: + @staticmethod + def execute(conditioning, reference_latents_method): + NativeFluxStub.last_call = { + "conditioning": conditioning, + "reference_latents_method": reference_latents_method, + } + return NodeOutputStub([["native-conditioning", {"reference_latents_method": reference_latents_method}]]) diff --git a/tests/test_flux2klein_node.py b/tests/test_flux2klein_node.py new file mode 100644 index 0000000..355a4ea --- /dev/null +++ b/tests/test_flux2klein_node.py @@ -0,0 +1,129 @@ +import unittest +from unittest import mock + +try: + from .test_flux2klein_helpers import ( + BrokenClipStub, + FakeMask, + FakeVAE, + FakeTensor, + NativeFluxStub, + NativeQwenStub, + StandardClipStub, + import_plugin_submodule, + load_plugin_package, + ) +except ImportError: + from test_flux2klein_helpers import ( + BrokenClipStub, + FakeMask, + FakeVAE, + FakeTensor, + NativeFluxStub, + NativeQwenStub, + 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(conditioning, [["cond::make it cinematic", {"pooled_output": {}}]]) + self.assertEqual(latent["samples"].shape, (1, 4, 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"], "lls_fallback") + self.assertTrue(custom_output["has_mask"]) + + 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, [["cond::change outfit", {"pooled_output": {}}]]) + self.assertEqual(main_image.shape, (1, 1024, 512, 3)) + self.assertEqual(latent["samples"].shape, (1, 16, 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") + + def test_native_qwen_flux_path_is_preferred_when_available(self): + with mock.patch.object(self.module, "nodes_qwen", NativeQwenStub), mock.patch.object( + self.module, "nodes_flux", NativeFluxStub + ): + conditioning, latent, custom_output, main_image, output_mask = self.node.encode( + clip=object(), + 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=FakeTensor((1, 512, 512, 3), label="ref2"), + image3=None, + mask=None, + ) + + self.assertEqual(conditioning, [["native-conditioning", {"reference_latents_method": "index_timestep_zero"}]]) + self.assertEqual(main_image.shape, (1, 512, 512, 3)) + self.assertEqual(output_mask.shape, (1, 512, 512)) + self.assertEqual(custom_output["conditioning_backend"], "native_qwen_flux") + self.assertFalse(custom_output["has_mask"]) + self.assertEqual(NativeQwenStub.last_call["prompt"], "replace background") + self.assertEqual(NativeFluxStub.last_call["reference_latents_method"], "index_timestep_zero") + + def test_missing_clip_backend_raises_clear_error(self): + with self.assertRaisesRegex(RuntimeError, "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() diff --git a/tests/test_flux2klein_registration.py b/tests/test_flux2klein_registration.py new file mode 100644 index 0000000..567c3e8 --- /dev/null +++ b/tests/test_flux2klein_registration.py @@ -0,0 +1,54 @@ +import unittest + +try: + from .test_flux2klein_helpers import load_plugin_package +except ImportError: + from test_flux2klein_helpers import load_plugin_package + + +class TestFlux2KleinRegistration(unittest.TestCase): + def test_plugin_registers_flux2klein_node(self): + plugin = load_plugin_package() + + self.assertIn("LLSFlux2KleinEditTextEncode", plugin.NODE_CLASS_MAPPINGS) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSFlux2KleinEditTextEncode"], + "LLS Flux2Klein Edit Text Encode", + ) + + def test_flux2klein_schema_matches_contract(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSFlux2KleinEditTextEncode"] + schema = node_cls.INPUT_TYPES() + required = schema["required"] + optional = schema["optional"] + + self.assertEqual(node_cls.CATEGORY, "LLS/Flux2Klein") + self.assertEqual(node_cls.FUNCTION, "encode") + self.assertEqual( + node_cls.RETURN_TYPES, + ("CONDITIONING", "LATENT", "LLS_FLUX2KLEIN_OUTPUT", "IMAGE", "MASK"), + ) + self.assertEqual( + node_cls.RETURN_NAMES, + ("conditioning", "latent", "custom_output", "main_image", "mask"), + ) + self.assertEqual(required["clip"], ("CLIP",)) + self.assertEqual(required["vae"], ("VAE",)) + self.assertEqual(required["image1"], ("IMAGE",)) + self.assertEqual(required["prompt"][0], "STRING") + self.assertTrue(required["prompt"][1]["multiline"]) + self.assertEqual(required["ref_longest_edge"][0], "INT") + self.assertEqual(required["ref_longest_edge"][1]["default"], 1024) + self.assertEqual(required["ref_longest_edge"][1]["min"], 256) + self.assertEqual(required["ref_longest_edge"][1]["max"], 4096) + self.assertEqual(required["ref_longest_edge"][1]["step"], 64) + self.assertEqual(required["resize_mode"][0], ["longest_edge", "keep_original"]) + self.assertEqual(required["mask_mode"][0], ["none", "use_mask", "invert_mask"]) + self.assertEqual(optional["image2"], ("IMAGE",)) + self.assertEqual(optional["image3"], ("IMAGE",)) + self.assertEqual(optional["mask"], ("MASK",)) + + +if __name__ == "__main__": + unittest.main()