添加背景替换节点
This commit is contained in:
Gin3601
@@ -45,6 +45,7 @@ _SUBPACKAGES: list[str] = [
|
||||
"conditioning", # 第 2 节：文本编码与条件生成
|
||||
"sampling", # 第 3 节：图像采样与去噪
|
||||
"qwen", # 第 4 节：Qwen 一体化节点
|
||||
"flux2klein", # 第 4 节（Flux2Klein 子域）：编辑文本编码
|
||||
"latent", # 第 4 节：Latent 空间操作
|
||||
"image", # 第 5 节：图像处理与后处理
|
||||
"repair", # 第 5 节（修复子域）：图像修复
|
||||
|
||||
@@ -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.
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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}]])
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user