Upload krea2_style_reference_injector.py

chain_injectors/krea2_style_reference_injector.py

@@ -156,8 +156,41 @@ def inject(assembler, chain_definition, chain_items):
     neg_ref_node['inputs']['conditioning'] = [neg_encode_id, 0]
     assembler.workflow[neg_ref_id] = neg_ref_node
 
-    assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_ref_id, 0]
-    assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_ref_id, 0]
+    existing_pos = assembler.workflow[ksampler_id]['inputs'].get('positive')
+    existing_neg = assembler.workflow[ksampler_id]['inputs'].get('negative')
+
+    has_krea2_edit = False
+    if existing_pos and isinstance(existing_pos, (list, tuple)) and len(existing_pos) > 0:
+        pos_node_id = existing_pos[0]
+        if pos_node_id in assembler.workflow:
+            pos_node = assembler.workflow[pos_node_id]
+            if isinstance(pos_node, dict) and pos_node.get('class_type') == 'Krea2EditGroundedEncode':
+                has_krea2_edit = True
+
+    if not has_krea2_edit:
+        for node in assembler.workflow.values():
+            if isinstance(node, dict) and node.get('class_type') in ['Krea2EditModelPatch', 'Krea2EditGroundedEncode']:
+                has_krea2_edit = True
+                break
+
+    if has_krea2_edit and existing_pos and existing_neg:
+        combine_pos_id = assembler._get_unique_id()
+        combine_pos_node = create_node(assembler, "ConditioningCombine", "Conditioning (Combine)")
+        combine_pos_node['inputs']['conditioning_1'] = existing_pos
+        combine_pos_node['inputs']['conditioning_2'] = [pos_ref_id, 0]
+        assembler.workflow[combine_pos_id] = combine_pos_node
+
+        combine_neg_id = assembler._get_unique_id()
+        combine_neg_node = create_node(assembler, "ConditioningCombine", "Conditioning (Combine)")
+        combine_neg_node['inputs']['conditioning_1'] = existing_neg
+        combine_neg_node['inputs']['conditioning_2'] = [neg_ref_id, 0]
+        assembler.workflow[combine_neg_id] = combine_neg_node
+
+        assembler.workflow[ksampler_id]['inputs']['positive'] = [combine_pos_id, 0]
+        assembler.workflow[ksampler_id]['inputs']['negative'] = [combine_neg_id, 0]
+    else:
+        assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_ref_id, 0]
+        assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_ref_id, 0]
 
     if pos_prompt_id and pos_prompt_id in assembler.workflow:
         del assembler.workflow[pos_prompt_id]