修改替换背景节点逻辑

tests/test_example_workflows.py

@@ -16,6 +16,12 @@ class TestExampleWorkflows(unittest.TestCase):
                 return node
         self.fail(f"Missing node type {node_type!r} in workflow")
 
+    def _find_input(self, node: dict, name: str) -> dict:
+        for item in node.get("inputs", []):
+            if item.get("name") == name:
+                return item
+        self.fail(f"Missing input {name!r} in node {node.get('type')!r}")
+
     def test_universal_txt2img_ksampler_keeps_seed_control_widget(self):
         workflow = self._load("LLS-universal-txt2img-workflow.json")
         node = self._find_node(workflow, "LLSSimpleKSampler")
@@ -36,6 +42,33 @@ class TestExampleWorkflows(unittest.TestCase):
         self.assertEqual(widgets[2], "randomize")
         self.assertEqual(widgets[3], 20)
 
+    def test_flux2klein_background_replace_workflow_uses_lls_edit_and_flux2_stack(self):
+        workflow = self._load("LLS-Flux2Klein-电商换背景工作流.json")
+
+        encode = self._find_node(workflow, "LLSFlux2KleinEditTextEncode")
+        unet = self._find_node(workflow, "UNETLoader")
+        clip = self._find_node(workflow, "CLIPLoader")
+        vae = self._find_node(workflow, "VAELoader")
+        self._find_node(workflow, "RMBG")
+        self._find_node(workflow, "InvertMask")
+        self._find_node(workflow, "ConditioningZeroOut")
+        self._find_node(workflow, "KSampler")
+        self._find_node(workflow, "VAEDecode")
+
+        self.assertEqual(unet.get("widgets_values", [None])[0], "flux-2-klein-9b.safetensors")
+        self.assertEqual(clip.get("widgets_values", [None, None])[1], "flux2")
+        self.assertEqual(vae.get("widgets_values", [None])[0], "flux2-vae.safetensors")
+
+        widgets = encode.get("widgets_values", [])
+        self.assertGreaterEqual(len(widgets), 4)
+        self.assertEqual(widgets[2], "longest_edge")
+        self.assertEqual(widgets[3], "use_mask")
+
+        self.assertIsNotNone(self._find_input(encode, "image1").get("link"))
+        self.assertIsNotNone(self._find_input(encode, "image2").get("link"))
+        self.assertIsNone(self._find_input(encode, "image3").get("link"))
+        self.assertIsNotNone(self._find_input(encode, "mask").get("link"))
+
 
 if __name__ == "__main__":
     unittest.main()