Add Reference Latent Chains

chain_injectors/reference_latent_injector.py

@@ -0,0 +1,66 @@
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    flux_guidance_name = chain_definition.get('flux_guidance_node')
+    vae_node_name = chain_definition.get('vae_node', 'vae_loader')
+
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: [ReferenceLatent] KSampler node '{ksampler_name}' not found. Skipping.")
+        return
+    if vae_node_name not in assembler.node_map:
+        print(f"Warning: [ReferenceLatent] VAE loader node '{vae_node_name}' not found. Skipping.")
+        return
+
+    ksampler_id = assembler.node_map[ksampler_name]
+    vae_node_id = assembler.node_map[vae_node_name]
+
+    pos_target_node_id = None
+    pos_target_input_name = None
+    if flux_guidance_name and flux_guidance_name in assembler.node_map:
+        flux_guidance_id = assembler.node_map[flux_guidance_name]
+        if 'conditioning' in assembler.workflow[flux_guidance_id]['inputs']:
+            pos_target_node_id = flux_guidance_id
+            pos_target_input_name = 'conditioning'
+            print(f"ReferenceLatent injector targeting FluxGuidance node '{flux_guidance_name}'.")
+    
+    if not pos_target_node_id:
+        if 'positive' in assembler.workflow[ksampler_id]['inputs']:
+            pos_target_node_id = ksampler_id
+            pos_target_input_name = 'positive'
+            print(f"ReferenceLatent injector targeting KSampler node '{ksampler_name}'.")
+        else:
+            print(f"Warning: [ReferenceLatent] Could not find a valid positive injection point. Skipping.")
+            return
+
+    current_pos_conditioning = assembler.workflow[pos_target_node_id]['inputs'][pos_target_input_name]
+
+    for i, img_filename in enumerate(chain_items):
+        if not img_filename or not isinstance(img_filename, str):
+            continue
+
+        load_id = assembler._get_unique_id()
+        load_node = assembler._get_node_template("LoadImage")
+        load_node['inputs']['image'] = img_filename
+        assembler.workflow[load_id] = load_node
+        
+        vae_encode_id = assembler._get_unique_id()
+        vae_encode_node = assembler._get_node_template("VAEEncode")
+        vae_encode_node['inputs']['pixels'] = [load_id, 0]
+        vae_encode_node['inputs']['vae'] = [vae_node_id, 0]
+        assembler.workflow[vae_encode_id] = vae_encode_node
+        
+        latent_conn = [vae_encode_id, 0]
+
+        ref_latent_id = assembler._get_unique_id()
+        ref_latent_node = assembler._get_node_template("ReferenceLatent")
+        ref_latent_node['inputs']['conditioning'] = current_pos_conditioning
+        ref_latent_node['inputs']['latent'] = latent_conn
+        assembler.workflow[ref_latent_id] = ref_latent_node
+        
+        current_pos_conditioning = [ref_latent_id, 0]
+
+    assembler.workflow[pos_target_node_id]['inputs'][pos_target_input_name] = current_pos_conditioning
+    
+    print(f"ReferenceLatent injector applied. Re-routed inputs through {len(chain_items)} reference image(s).")

core/pipelines/sd_image_pipeline.py

@@ -312,6 +312,16 @@ class SdImagePipeline(BasePipeline):
                         "strength": float(strengths[i])
                     })
 
+        reference_latent_data = ui_inputs.get('reference_latent_data', [])
+        active_reference_latents = []
+        if reference_latent_data:
+            for img_pil in reference_latent_data:
+                if img_pil is not None:
+                    temp_file_path = os.path.join(INPUT_DIR, f"temp_ref_{random.randint(1000, 9999)}.png")
+                    img_pil.save(temp_file_path, "PNG")
+                    active_reference_latents.append(os.path.basename(temp_file_path))
+                    temp_files_to_clean.append(temp_file_path)
+
         loras_string = f"LoRAs: [{', '.join(active_loras_for_meta)}]" if active_loras_for_meta else ""
         
         progress(0.8, desc="Assembling workflow...")
@@ -346,6 +356,7 @@ class SdImagePipeline(BasePipeline):
             "clip_name": components['clip'],
             "vae_name": ui_inputs.get('vae_name', components['vae']),
             "conditioning_chain": active_conditioning,
+            "reference_latent_chain": active_reference_latents,
         }
         
         if task_type == 'txt2img':

core/pipelines/workflow_recipes/_partials/conditioning/flux2.yaml

@@ -43,6 +43,11 @@ dynamic_conditioning_chains:
     ksampler_node: "ksampler"
     clip_source: "clip_loader:0"
 
+dynamic_reference_latent_chains:
+  reference_latent_chain:
+    ksampler_node: "ksampler"
+    vae_node: "vae_loader"
+
 ui_map:
   unet_name: "unet_loader:unet_name"
   clip_name: "clip_loader:clip_name"

core/settings.py

@@ -111,11 +111,13 @@ try:
     MAX_EMBEDDINGS = _constants.get('MAX_EMBEDDINGS', 5)
     MAX_CONDITIONINGS = _constants.get('MAX_CONDITIONINGS', 10)
     MAX_CONTROLNETS = _constants.get('MAX_CONTROLNETS', 5)
+    MAX_REFERENCE_LATENTS = _constants.get('MAX_REFERENCE_LATENTS', 10)
     LORA_SOURCE_CHOICES = _constants.get('LORA_SOURCE_CHOICES', ["Civitai", "File"])
     RESOLUTION_MAP = _constants.get('RESOLUTION_MAP', {})
 except Exception as e:
     print(f"FATAL: Could not load constants from YAML. Error: {e}")
     MAX_LORAS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_CONTROLNETS = 5, 5, 10, 5
+    MAX_REFERENCE_LATENTS = 10
     LORA_SOURCE_CHOICES = ["Civitai", "File"]
     RESOLUTION_MAP = {}
 

ui/events.py

@@ -10,7 +10,7 @@ from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
 
 from core.pipelines.controlnet_preprocessor import CPU_ONLY_PREPROCESSORS
 from utils.app_utils import PREPROCESSOR_MODEL_MAP, PREPROCESSOR_PARAMETER_MAP, save_uploaded_file_with_hash
-from ui.shared.ui_components import RESOLUTION_MAP, MAX_CONTROLNETS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_LORAS
+from ui.shared.ui_components import RESOLUTION_MAP, MAX_CONTROLNETS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_LORAS, MAX_REFERENCE_LATENTS
 
 
 def on_model_change(model_display_name):
@@ -253,6 +253,36 @@ def attach_event_handlers(ui_components, demo):
         del_outputs = [count_state, add_button, del_button] + rows + prompts
         add_button.click(fn=add_row, inputs=[count_state], outputs=add_outputs, show_progress=False)
         del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
+
+    def create_reference_latent_event_handlers(prefix):
+        rows = ui_components[f'reference_latent_rows_{prefix}']
+        images = ui_components[f'reference_latent_images_{prefix}']
+        count_state = ui_components[f'reference_latent_count_state_{prefix}']
+        add_button = ui_components[f'add_reference_latent_button_{prefix}']
+        del_button = ui_components[f'delete_reference_latent_button_{prefix}']
+
+        def add_row(c):
+            c += 1
+            return {
+                count_state: c,
+                rows[c - 1]: gr.update(visible=True),
+                add_button: gr.update(visible=c < MAX_REFERENCE_LATENTS),
+                del_button: gr.update(visible=True),
+            }
+
+        def del_row(c):
+            c -= 1
+            return {
+                count_state: c,
+                rows[c]: gr.update(visible=False),
+                images[c]: None,
+                add_button: gr.update(visible=True),
+                del_button: gr.update(visible=c > 0),
+            }
+
+        add_outputs = [count_state, add_button, del_button] + rows
+        del_outputs = [count_state, add_button, del_button] + rows + images
+        add_button.click(fn=add_row, inputs=[count_state], outputs=add_outputs, show_progress=False)
         del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
     
     def on_vae_upload(file_obj):
@@ -318,7 +348,7 @@ def attach_event_handlers(ui_components, demo):
 
         input_keys = list(run_inputs_map.keys())
         input_list_flat = [v for v in run_inputs_map.values() if v is not None]
-        input_list_flat += lora_data_components + embedding_data_components + conditioning_data_components 
+        input_list_flat += lora_data_components + embedding_data_components + conditioning_data_components + reference_latent_components
 
         def create_ui_inputs_dict(*args):
             valid_keys = [k for k in input_keys if run_inputs_map[k] is not None]
@@ -331,6 +361,8 @@ def attach_event_handlers(ui_components, demo):
             arg_idx += len(embedding_data_components)
             ui_dict['conditioning_data'] = list(args[arg_idx : arg_idx + len(conditioning_data_components)])
             arg_idx += len(conditioning_data_components)
+            ui_dict['reference_latent_data'] = list(args[arg_idx : arg_idx + len(reference_latent_components)])
+
 
             return ui_dict
 
@@ -383,7 +415,9 @@ def attach_event_handlers(ui_components, demo):
                         outputs=[emb_ids[i], emb_sources[i], emb_files[i]],
                         show_progress=False
                     )
-        if f'add_conditioning_button_{prefix}' in ui_components: create_conditioning_event_handlers(prefix)        if f'vae_source_{prefix}' in ui_components:
+        if f'add_conditioning_button_{prefix}' in ui_components: create_conditioning_event_handlers(prefix)
+        if f'add_reference_latent_button_{prefix}' in ui_components: create_reference_latent_event_handlers(prefix)
+        if f'vae_source_{prefix}' in ui_components:
             upload_button = ui_components.get(f'vae_upload_button_{prefix}')
             if upload_button:
                  upload_button.upload(

ui/shared/hires_fix_ui.py

@@ -4,7 +4,8 @@ from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
 from .ui_components import (
     create_lora_settings_ui,
     create_embedding_ui,
-    create_conditioning_ui, create_vae_override_ui, create_api_key_ui
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
 )
 
 def create_ui():
@@ -68,7 +69,7 @@ def create_ui():
         # components.update(create_diffsynth_controlnet_ui(prefix))
         # components.update(create_controlnet_ui(prefix))
         # components.update(create_embedding_ui(prefix))
-        components.update((prefix))
+        components.update(create_reference_latent_ui(prefix))
         components.update(create_conditioning_ui(prefix))
         # components.update(create_vae_override_ui(prefix))
             

ui/shared/img2img_ui.py

@@ -4,7 +4,8 @@ from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
 from .ui_components import (
     create_lora_settings_ui,
     create_embedding_ui,
-    create_conditioning_ui, create_vae_override_ui, create_api_key_ui
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
 )
 
 def create_ui():
@@ -51,7 +52,7 @@ def create_ui():
         # components.update(create_diffsynth_controlnet_ui(prefix))
         # components.update(create_controlnet_ui(prefix))
         # components.update(create_embedding_ui(prefix))
-        components.update((prefix))
+        components.update(create_reference_latent_ui(prefix))
         components.update(create_conditioning_ui(prefix))
         # components.update(create_vae_override_ui(prefix))
             

ui/shared/inpaint_ui.py

@@ -3,7 +3,8 @@ from core.settings import MODEL_MAP_CHECKPOINT
 from .ui_components import (
     create_base_parameter_ui, create_lora_settings_ui,
     create_embedding_ui,
-    create_conditioning_ui, create_vae_override_ui, create_api_key_ui
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
 )
 
 def create_ui():
@@ -74,7 +75,7 @@ def create_ui():
             # components.update(create_diffsynth_controlnet_ui(prefix))
             # components.update(create_controlnet_ui(prefix))
             # components.update(create_embedding_ui(prefix))
-            components.update((prefix))
+            components.update(create_reference_latent_ui(prefix))
             components.update(create_conditioning_ui(prefix))
             # components.update(create_vae_override_ui(prefix))
         components[f'accordion_wrapper_{prefix}'] = accordion_wrapper

ui/shared/outpaint_ui.py

@@ -4,7 +4,8 @@ from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
 from .ui_components import (
     create_lora_settings_ui,
     create_embedding_ui,
-    create_conditioning_ui, create_vae_override_ui, create_api_key_ui
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
 )
 
 def create_ui():
@@ -62,7 +63,7 @@ def create_ui():
         # components.update(create_diffsynth_controlnet_ui(prefix))
         # components.update(create_controlnet_ui(prefix))
         # components.update(create_embedding_ui(prefix))
-        components.update((prefix))
+        components.update(create_reference_latent_ui(prefix))
         components.update(create_conditioning_ui(prefix))
         # components.update(create_vae_override_ui(prefix))
             

ui/shared/txt2img_ui.py

@@ -3,7 +3,8 @@ from core.settings import MODEL_MAP_CHECKPOINT
 from .ui_components import (
     create_base_parameter_ui, create_lora_settings_ui,
     create_embedding_ui,
-    create_conditioning_ui, create_vae_override_ui, create_api_key_ui
+    create_conditioning_ui, create_vae_override_ui, create_api_key_ui,
+    create_reference_latent_ui
 )
 
 def create_ui():
@@ -32,7 +33,7 @@ def create_ui():
         # components.update(create_diffsynth_controlnet_ui(prefix))
         # components.update(create_controlnet_ui(prefix))
         # components.update(create_embedding_ui(prefix))
-        components.update((prefix))
+        components.update(create_reference_latent_ui(prefix))
         components.update(create_conditioning_ui(prefix))
         # components.update(create_vae_override_ui(prefix))
         

ui/shared/ui_components.py

@@ -2,7 +2,7 @@ import gradio as gr
 from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
 from core.settings import (
     MAX_LORAS, LORA_SOURCE_CHOICES, MAX_EMBEDDINGS, MAX_CONDITIONINGS,
-    MAX_CONTROLNETS, RESOLUTION_MAP
+    MAX_CONTROLNETS, RESOLUTION_MAP, MAX_REFERENCE_LATENTS
 )
 import yaml
 import os
@@ -172,6 +172,35 @@ def create_conditioning_ui(prefix: str):
         components[key('all_conditioning_components_flat')] = all_cond_components_flat
     
     return components
+    
+def create_reference_latent_ui(prefix: str):
+    components = {}
+    key = lambda name: f"{name}_{prefix}"
+
+    with gr.Accordion("Reference Edit", open=False) as accordion:
+        components[key('reference_latent_accordion')] = accordion
+        gr.Markdown("💡 **Tip:** For multimodal models (like FLUX.2), this feature enables powerful editing and combining capabilities. In txt2img mode, adding a single reference image performs an **Image Edit**, while adding multiple images performs an **Image Combine**.")
+        
+        ref_rows, ref_images = [], []
+        components.update({
+            key('reference_latent_rows'): ref_rows,
+            key('reference_latent_images'): ref_images,
+        })
+        
+        with gr.Row():
+            for i in range(MAX_REFERENCE_LATENTS):
+                with gr.Column(visible=(i < 1), min_width=160) as row_wrapper:
+                    ref_images.append(gr.Image(type="pil", label=f"Reference {i+1}", sources=["upload"], height=150))
+                    ref_rows.append(row_wrapper)
+
+        with gr.Row():
+            components[key('add_reference_latent_button')] = gr.Button("✚ Add Reference Image")
+            components[key('delete_reference_latent_button')] = gr.Button("➖ Delete Reference Image", visible=False)
+        components[key('reference_latent_count_state')] = gr.State(1)
+
+        components[key('all_reference_latent_components_flat')] = ref_images
+
+    return components
 
 def create_vae_override_ui(prefix: str):
     components = {}

yaml/constants.yaml

@@ -2,6 +2,7 @@ MAX_LORAS: 5
 MAX_CONTROLNETS: 5
 MAX_EMBEDDINGS: 5
 MAX_CONDITIONINGS: 10
+MAX_REFERENCE_LATENTS: 10
 LORA_SOURCE_CHOICES: ["Civitai", "File"]
 
 RESOLUTION_MAP:

yaml/injectors.yaml

@@ -1,6 +1,9 @@
 injector_definitions:
   dynamic_conditioning_chains:
     module: "chain_injectors.conditioning_injector"
+  dynamic_reference_latent_chains:
+    module: "chain_injectors.reference_latent_injector"
 
 injector_order:
+  - dynamic_reference_latent_chains
   - dynamic_conditioning_chains