Add Krea-2 Reference Edit injector.

README.md

@@ -107,6 +107,7 @@ models:
   - OnomaAIResearch/Illustrious-XL-v1.0
   - OnomaAIResearch/Illustrious-XL-v1.1
   - OnomaAIResearch/Illustrious-XL-v2.0
+  - ostris/krea2_turbo_style_reference
   - Patil/Krea-2-depth-controlnet
   - RedRayz/hikari_noob_v-pred_1.2.4
   - Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0

chain_injectors/krea2_identity_edit_injector.py

@@ -115,10 +115,24 @@ def inject(assembler, chain_definition, chain_items):
     pos_text = ""
     if pos_prompt_id and pos_prompt_id in assembler.workflow:
         pos_text = assembler.workflow[pos_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        pos_text = assembler.ui_values.get('positive_prompt') or assembler.ui_values.get('prompt') or ''
+
+    if not pos_text:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict):
+                cls = node.get('class_type', '')
+                if cls in ['Krea2EditGroundedEncode', 'TextEncodeQwenImageEditPlus', 'CLIPTextEncode']:
+                    t = node.get('inputs', {}).get('prompt') or node.get('inputs', {}).get('text')
+                    if t:
+                        pos_text = t
+                        break
 
     neg_text = ""
     if neg_prompt_id and neg_prompt_id in assembler.workflow:
         neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        neg_text = assembler.ui_values.get('negative_prompt') or assembler.ui_values.get('neg_prompt') or ''
 
     pos_grounded_id = assembler._get_unique_id()
     pos_grounded_node = assembler._get_node_template("Krea2EditGroundedEncode")

chain_injectors/krea2_reference_edit_injector.py

@@ -0,0 +1,168 @@
+import os
+from utils.app_utils import ensure_file_downloaded
+
+def create_node(assembler, class_type, title):
+    try:
+        node = assembler._get_node_template(class_type)
+    except Exception:
+        node = {
+            "inputs": {},
+            "class_type": class_type,
+            "_meta": {"title": title}
+        }
+    node['_meta']['title'] = title
+    return node
+
+def inject(assembler, chain_definition, chain_items):
+    if not chain_items:
+        return
+
+    valid_images = []
+    for item in chain_items:
+        if not item:
+            continue
+        img_path = item
+        if isinstance(item, dict):
+            img_path = item.get('image') or item.get('filename') or item.get('path')
+        if img_path:
+            valid_images.append(img_path)
+
+    if not valid_images:
+        return
+
+    valid_images = valid_images[:3]
+
+    lora_filename = "krea2_style_reference.safetensors"
+    try:
+        ensure_file_downloaded(lora_filename)
+    except Exception as e:
+        print(f"Warning: Failed to ensure '{lora_filename}' downloaded: {e}")
+
+    ksampler_name = chain_definition.get('ksampler_node', 'ksampler')
+    pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt')
+    neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt')
+    clip_loader_name = chain_definition.get('clip_loader_node', 'clip_loader')
+    vae_loader_name = chain_definition.get('vae_loader_node', 'vae_loader')
+
+    if ksampler_name not in assembler.node_map:
+        print(f"Warning: Target node '{ksampler_name}' for Krea2 Style Reference Edit chain not found. Skipping.")
+        return
+
+    ksampler_id = assembler.node_map[ksampler_name]
+
+    if 'model' not in assembler.workflow[ksampler_id]['inputs']:
+        print(f"Warning: KSampler node '{ksampler_name}' is missing 'model' input. Skipping.")
+        return
+
+    current_model_connection = assembler.workflow[ksampler_id]['inputs']['model']
+
+    vae_connection = None
+    if vae_loader_name in assembler.node_map:
+        vae_connection = [assembler.node_map[vae_loader_name], 0]
+    else:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict) and node.get('class_type') == 'VAELoader':
+                vae_connection = [node_id, 0]
+                break
+
+    clip_connection = None
+    if clip_loader_name in assembler.node_map:
+        clip_connection = [assembler.node_map[clip_loader_name], 0]
+    elif pos_prompt_name in assembler.node_map:
+        pos_id = assembler.node_map[pos_prompt_name]
+        clip_connection = assembler.workflow[pos_id]['inputs'].get('clip')
+
+    scaled_image_ids = []
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = create_node(assembler, "LoadImage", f"Load Reference Image {i+1}")
+        load_node['inputs']['image'] = img_filename
+        assembler.workflow[load_id] = load_node
+
+        scale_id = assembler._get_unique_id()
+        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
+        scale_node['inputs']['upscale_method'] = "nearest-exact"
+        scale_node['inputs']['megapixels'] = 1
+        scale_node['inputs']['resolution_steps'] = 1
+        scale_node['inputs']['image'] = [load_id, 0]
+        assembler.workflow[scale_id] = scale_node
+        scaled_image_ids.append(scale_id)
+
+    lora_loader_id = assembler._get_unique_id()
+    lora_loader_node = create_node(assembler, "LoraLoaderModelOnly", "Load LoRA (Krea2 Style Reference)")
+    lora_loader_node['inputs']['lora_name'] = lora_filename
+    lora_loader_node['inputs']['strength_model'] = 1.0
+    lora_loader_node['inputs']['model'] = current_model_connection
+    assembler.workflow[lora_loader_id] = lora_loader_node
+
+    assembler.workflow[ksampler_id]['inputs']['model'] = [lora_loader_id, 0]
+
+    pos_prompt_id = assembler.node_map.get(pos_prompt_name)
+    neg_prompt_id = assembler.node_map.get(neg_prompt_name)
+
+    pos_text = ""
+    if pos_prompt_id and pos_prompt_id in assembler.workflow:
+        pos_text = assembler.workflow[pos_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        pos_text = assembler.ui_values.get('positive_prompt') or assembler.ui_values.get('prompt') or ''
+
+    if not pos_text:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict):
+                cls = node.get('class_type', '')
+                if cls in ['Krea2EditGroundedEncode', 'TextEncodeQwenImageEditPlus', 'CLIPTextEncode']:
+                    t = node.get('inputs', {}).get('prompt') or node.get('inputs', {}).get('text')
+                    if t:
+                        pos_text = t
+                        break
+
+    neg_text = ""
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        neg_text = assembler.ui_values.get('negative_prompt') or assembler.ui_values.get('neg_prompt') or ''
+
+    pos_encode_id = assembler._get_unique_id()
+    pos_encode_node = create_node(assembler, "TextEncodeQwenImageEditPlus", "TextEncodeQwenImageEditPlus (Positive)")
+    pos_encode_node['inputs']['prompt'] = pos_text
+    if clip_connection:
+        pos_encode_node['inputs']['clip'] = clip_connection
+    if vae_connection:
+        pos_encode_node['inputs']['vae'] = vae_connection
+    for idx, s_id in enumerate(scaled_image_ids):
+        pos_encode_node['inputs'][f"image{idx+1}"] = [s_id, 0]
+    assembler.workflow[pos_encode_id] = pos_encode_node
+
+    neg_encode_id = assembler._get_unique_id()
+    neg_encode_node = create_node(assembler, "TextEncodeQwenImageEditPlus", "TextEncodeQwenImageEditPlus (Negative)")
+    neg_encode_node['inputs']['prompt'] = neg_text
+    if clip_connection:
+        neg_encode_node['inputs']['clip'] = clip_connection
+    if vae_connection:
+        neg_encode_node['inputs']['vae'] = vae_connection
+    for idx, s_id in enumerate(scaled_image_ids):
+        neg_encode_node['inputs'][f"image{idx+1}"] = [s_id, 0]
+    assembler.workflow[neg_encode_id] = neg_encode_node
+
+    pos_ref_id = assembler._get_unique_id()
+    pos_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
+    pos_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
+    pos_ref_node['inputs']['conditioning'] = [pos_encode_id, 0]
+    assembler.workflow[pos_ref_id] = pos_ref_node
+
+    neg_ref_id = assembler._get_unique_id()
+    neg_ref_node = create_node(assembler, "FluxKontextMultiReferenceLatentMethod", "Edit Model Reference Method")
+    neg_ref_node['inputs']['reference_latents_method'] = "index_timestep_zero"
+    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]
+
+    if pos_prompt_id and pos_prompt_id in assembler.workflow:
+        del assembler.workflow[pos_prompt_id]
+
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        del assembler.workflow[neg_prompt_id]
+
+    print(f"Krea2 Style Reference Edit injector applied with {len(valid_images)} reference image(s). Original CLIPTextEncode nodes replaced.")

chain_injectors/qwen_image_edit_injector.py

@@ -70,10 +70,24 @@ def inject(assembler, chain_definition, chain_items):
     pos_text = ""
     if pos_prompt_id and pos_prompt_id in assembler.workflow:
         pos_text = assembler.workflow[pos_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        pos_text = assembler.ui_values.get('positive_prompt') or assembler.ui_values.get('prompt') or ''
+
+    if not pos_text:
+        for node_id, node in assembler.workflow.items():
+            if isinstance(node, dict):
+                cls = node.get('class_type', '')
+                if cls in ['Krea2EditGroundedEncode', 'TextEncodeQwenImageEditPlus', 'CLIPTextEncode']:
+                    t = node.get('inputs', {}).get('prompt') or node.get('inputs', {}).get('text')
+                    if t:
+                        pos_text = t
+                        break
 
     neg_text = ""
     if neg_prompt_id and neg_prompt_id in assembler.workflow:
         neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
+    elif hasattr(assembler, 'ui_values') and isinstance(assembler.ui_values, dict):
+        neg_text = assembler.ui_values.get('negative_prompt') or assembler.ui_values.get('neg_prompt') or ''
 
     scaled_image_ids = []
     for i, img_filename in enumerate(valid_images):

core/pipelines/pipeline_input_processor.py

@@ -368,6 +368,17 @@ def process_pipeline_inputs(ui_inputs: Dict[str, Any], progress: gr.Progress, wo
                 temp_files_to_clean.append(temp_path)
                 active_krea2_identity_edit.append(os.path.basename(temp_path))
 
+    krea2_reference_edit_data = ui_inputs.get('krea2_reference_edit_data', [])
+    active_krea2_reference_edit = []
+    if krea2_reference_edit_data:
+        for img in krea2_reference_edit_data:
+            if img:
+                if not os.path.exists(INPUT_DIR): os.makedirs(INPUT_DIR)
+                temp_path = os.path.join(INPUT_DIR, f"temp_krea2_reference_ref_{random.randint(1000, 9999)}.png")
+                img.save(temp_path, "PNG")
+                temp_files_to_clean.append(temp_path)
+                active_krea2_reference_edit.append(os.path.basename(temp_path))
+
     qwen_image_edit_data = ui_inputs.get('qwen_image_edit_data', [])
     active_qwen_image_edit = []
     if qwen_image_edit_data:
@@ -433,6 +444,7 @@ def process_pipeline_inputs(ui_inputs: Dict[str, Any], progress: gr.Progress, wo
         "active_hidream_o1_reference": active_hidream_o1_reference,
         "active_joyai_reference": active_joyai_reference,
         "active_krea2_identity_edit": active_krea2_identity_edit,
+        "active_krea2_reference_edit": active_krea2_reference_edit,
         "active_qwen_image_edit": active_qwen_image_edit,
         "active_reference_images": active_reference_images,
         "active_conditioning": active_conditioning,

core/pipelines/sd_image_pipeline.py

@@ -35,7 +35,8 @@ class SdImagePipeline(BasePipeline):
         decoded_images_tensor = WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects)
         
         output_images = []
-        start_seed = ui_inputs['seed'] if ui_inputs['seed'] != -1 else random.randint(0, 2**64 - 1)
+        raw_seed = ui_inputs.get('seed')
+        start_seed = int(raw_seed) if (raw_seed is not None and raw_seed != -1) else random.randint(0, 2**64 - 1)
         for i in range(decoded_images_tensor.shape[0]):
             img_tensor = decoded_images_tensor[i]
             pil_image = Image.fromarray((img_tensor.cpu().numpy() * 255.0).astype("uint8"))
@@ -118,6 +119,7 @@ class SdImagePipeline(BasePipeline):
             active_hidream_o1_reference = processed["active_hidream_o1_reference"]
             active_joyai_reference = processed.get("active_joyai_reference", [])
             active_krea2_identity_edit = processed.get("active_krea2_identity_edit", [])
+            active_krea2_reference_edit = processed.get("active_krea2_reference_edit", [])
             active_qwen_image_edit = processed.get("active_qwen_image_edit", [])
             active_reference_images = processed.get("active_reference_images", [])
             active_conditioning = processed["active_conditioning"]
@@ -126,7 +128,8 @@ class SdImagePipeline(BasePipeline):
             
             progress(0.8, desc="Assembling workflow...")
             
-            if ui_inputs.get('seed') == -1:
+            seed_val = ui_inputs.get('seed')
+            if seed_val is None or seed_val == -1:
                 ui_inputs['seed'] = random.randint(0, 2**32 - 1)
             
             model_info = ALL_MODEL_MAP[model_display_name]
@@ -178,6 +181,7 @@ class SdImagePipeline(BasePipeline):
                 "hidream_o1_reference_chain": active_hidream_o1_reference,
                 "joyai_reference_chain": active_joyai_reference,
                 "krea2_identity_edit_chain": active_krea2_identity_edit,
+                "krea2_reference_edit_chain": active_krea2_reference_edit,
                 "qwen_image_edit_chain": active_qwen_image_edit,
                 "reference_image_chain": active_reference_images,
                 "vae_chain": [ui_inputs.get('vae_name')] if ui_inputs.get('vae_name') else [],

core/pipelines/workflow_recipes/_partials/conditioning/krea-2.yaml

@@ -62,6 +62,14 @@ dynamic_krea2_identity_edit_chains:
     clip_loader_node: "clip_loader"
     vae_loader_node: "vae_loader"
 
+dynamic_krea2_reference_edit_chains:
+  krea2_reference_edit_chain:
+    ksampler_node: "ksampler"
+    pos_prompt_node: "pos_prompt"
+    neg_prompt_node: "neg_prompt"
+    clip_loader_node: "clip_loader"
+    vae_loader_node: "vae_loader"
+
 dynamic_conditioning_chains:
   conditioning_chain:
     ksampler_node: "ksampler"

core/workflow_assembler.py

@@ -127,6 +127,7 @@ class WorkflowAssembler:
         return merged_recipe
 
     def assemble(self, ui_values):
+        self.ui_values = ui_values
         for name, details in self.recipe['nodes'].items():
             class_type = details['class_type']
             template = self._get_node_template(class_type)

requirements.txt

@@ -1,5 +1,5 @@
-comfyui-frontend-package==1.47.10
-comfyui-workflow-templates==0.11.19
+comfyui-frontend-package==1.47.11
+comfyui-workflow-templates==0.11.20
 comfyui-embedded-docs==0.5.9
 torch
 torchsde
@@ -22,7 +22,7 @@ alembic
 SQLAlchemy>=2.0.0
 filelock
 av>=16.0.0
-comfy-kitchen==0.2.24
+comfy-kitchen==0.2.25
 comfy-aimdo==0.4.10
 requests
 simpleeval>=1.0.0

ui/events/chain_handlers.py

@@ -900,6 +900,39 @@ def create_qwen_image_edit_event_handlers(prefix, ui_components):
     del_button = ui_components[f'delete_qwen_image_edit_button_{prefix}']
     images = ui_components[f'qwen_image_edit_images_{prefix}']
 
+    def add_ref_row(c):
+        c += 1
+        return {
+            count_state: c,
+            ref_rows[c - 1]: gr.update(visible=True),
+            add_button: gr.update(visible=c < 3),
+            del_button: gr.update(visible=True),
+        }
+
+    def del_ref_row(c):
+        c -= 1
+        return {
+            count_state: c,
+            ref_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] + ref_rows
+    del_outputs = [count_state, add_button, del_button] + ref_rows + images
+    add_button.click(fn=add_ref_row, inputs=[count_state], outputs=add_outputs, show_progress=False)
+    del_button.click(fn=del_ref_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
+
+
+def create_krea2_reference_edit_event_handlers(prefix, ui_components):
+    ref_rows = ui_components.get(f'krea2_reference_edit_rows_{prefix}')
+    if not ref_rows: return
+    count_state = ui_components[f'krea2_reference_edit_count_state_{prefix}']
+    add_button = ui_components[f'add_krea2_reference_edit_button_{prefix}']
+    del_button = ui_components[f'delete_krea2_reference_edit_button_{prefix}']
+    images = ui_components[f'krea2_reference_edit_images_{prefix}']
+
     def add_ref_row(c):
         c += 1
         return {

ui/events/change_handlers.py

@@ -17,7 +17,7 @@ from .config_loaders import (
     load_ipadapter_config
 )
 
-def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, krea2_cn_types, krea2_cn_series, krea2_cn_filepaths, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, krea2_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp, pid_acc=None, vae_acc=None, joyai_ref_acc=None, krea2_identity_edit_acc=None, qwen_image_edit_acc=None, ref_img_acc=None):
+def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, krea2_cn_types, krea2_cn_series, krea2_cn_filepaths, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, krea2_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp, pid_acc=None, vae_acc=None, joyai_ref_acc=None, krea2_identity_edit_acc=None, krea2_reference_edit_acc=None, qwen_image_edit_acc=None, ref_img_acc=None):
     def update_fn(*args):
         arch = args[0]
         category = args[1]
@@ -70,6 +70,7 @@ def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp,
         if hidream_o1_ref_acc: updates[hidream_o1_ref_acc] = gr.update(visible=('hidream_o1_reference' in enabled_chains))
         if joyai_ref_acc: updates[joyai_ref_acc] = gr.update(visible=('joyai_reference' in enabled_chains))
         if krea2_identity_edit_acc: updates[krea2_identity_edit_acc] = gr.update(visible=('krea2_identity_edit' in enabled_chains))
+        if krea2_reference_edit_acc: updates[krea2_reference_edit_acc] = gr.update(visible=('krea2_reference_edit' in enabled_chains))
         if qwen_image_edit_acc: updates[qwen_image_edit_acc] = gr.update(visible=('qwen_image_edit' in enabled_chains))
         if ref_img_acc: updates[ref_img_acc] = gr.update(visible=('reference_image' in enabled_chains))
         if pid_acc: updates[pid_acc] = gr.update(visible=('pid' in enabled_chains))
@@ -146,7 +147,7 @@ def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp,
     return update_fn
 
 
-def make_model_change_fn(cat_comp_ref, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, krea2_cn_types, krea2_cn_series, krea2_cn_filepaths, arch_comp_ref, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, krea2_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp, pid_acc=None, vae_acc=None, joyai_ref_acc=None, krea2_identity_edit_acc=None, qwen_image_edit_acc=None, ref_img_acc=None):
+def make_model_change_fn(cat_comp_ref, cs_comp, ar_comp, width_comp, height_comp, cn_types, cn_series, cn_filepaths, anima_cn_types, anima_cn_series, anima_cn_filepaths, diffsynth_cn_types, diffsynth_cn_series, diffsynth_cn_filepaths, krea2_cn_types, krea2_cn_series, krea2_cn_filepaths, arch_comp_ref, ipa_preset, lora_acc, cn_acc, anima_cn_acc, diffsynth_cn_acc, krea2_cn_acc, ipa_acc, sd3_ipa_acc, flux1_ipa_acc, style_acc, embed_acc, cond_acc, ref_latent_acc, hidream_o1_ref_acc, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp, pid_acc=None, vae_acc=None, joyai_ref_acc=None, krea2_identity_edit_acc=None, krea2_reference_edit_acc=None, qwen_image_edit_acc=None, ref_img_acc=None):
     def change_fn(*args):
         model_name = args[0]
         idx = 1
@@ -204,6 +205,7 @@ def make_model_change_fn(cat_comp_ref, cs_comp, ar_comp, width_comp, height_comp
         if hidream_o1_ref_acc: updates[hidream_o1_ref_acc] = gr.update(visible=('hidream_o1_reference' in enabled_chains))
         if joyai_ref_acc: updates[joyai_ref_acc] = gr.update(visible=('joyai_reference' in enabled_chains))
         if krea2_identity_edit_acc: updates[krea2_identity_edit_acc] = gr.update(visible=('krea2_identity_edit' in enabled_chains))
+        if krea2_reference_edit_acc: updates[krea2_reference_edit_acc] = gr.update(visible=('krea2_reference_edit' in enabled_chains))
         if qwen_image_edit_acc: updates[qwen_image_edit_acc] = gr.update(visible=('qwen_image_edit' in enabled_chains))
         if ref_img_acc: updates[ref_img_acc] = gr.update(visible=('reference_image' in enabled_chains))
         if pid_acc: updates[pid_acc] = gr.update(visible=('pid' in enabled_chains))

ui/events/main.py

@@ -15,6 +15,7 @@ from .chain_handlers import (
     create_hidream_o1_reference_event_handlers,
     create_joyai_reference_event_handlers,
     create_krea2_identity_edit_event_handlers,
+    create_krea2_reference_edit_event_handlers,
     create_qwen_image_edit_event_handlers,
     create_reference_image_event_handlers
 )
@@ -72,6 +73,7 @@ def attach_event_handlers(ui_components, demo):
         hidream_o1_ref_accordion = ui_components.get(f'hidream_o1_reference_accordion_{prefix}')
         joyai_ref_accordion = ui_components.get(f'joyai_reference_accordion_{prefix}')
         krea2_identity_edit_accordion = ui_components.get(f'krea2_identity_edit_accordion_{prefix}')
+        krea2_reference_edit_accordion = ui_components.get(f'krea2_reference_edit_accordion_{prefix}')
         qwen_image_edit_accordion = ui_components.get(f'qwen_image_edit_accordion_{prefix}')
         ref_img_accordion = ui_components.get(f'reference_image_accordion_{prefix}')
         pid_accordion = ui_components.get(f'pid_accordion_{prefix}')
@@ -113,6 +115,7 @@ def attach_event_handlers(ui_components, demo):
             if hidream_o1_ref_accordion: outputs.append(hidream_o1_ref_accordion)
             if joyai_ref_accordion: outputs.append(joyai_ref_accordion)
             if krea2_identity_edit_accordion: outputs.append(krea2_identity_edit_accordion)
+            if krea2_reference_edit_accordion: outputs.append(krea2_reference_edit_accordion)
             if qwen_image_edit_accordion: outputs.append(qwen_image_edit_accordion)
             if ref_img_accordion: outputs.append(ref_img_accordion)
             if pid_accordion: outputs.append(pid_accordion)
@@ -129,7 +132,7 @@ def attach_event_handlers(ui_components, demo):
                 krea2_cn_types_list, krea2_cn_series_list, krea2_cn_filepaths_list,
                 ipa_preset_list, lora_accordion, cn_accordion, anima_cn_accordion, diffsynth_cn_accordion, krea2_cn_accordion, ipa_accordion, sd3_ipa_accordion, flux1_ipa_accordion, style_accordion, embedding_accordion, conditioning_accordion,
                 ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp,
-                pid_acc=pid_accordion, vae_acc=vae_accordion, joyai_ref_acc=joyai_ref_accordion, krea2_identity_edit_acc=krea2_identity_edit_accordion, qwen_image_edit_acc=qwen_image_edit_accordion, ref_img_acc=ref_img_accordion
+                pid_acc=pid_accordion, vae_acc=vae_accordion, joyai_ref_acc=joyai_ref_accordion, krea2_identity_edit_acc=krea2_identity_edit_accordion, krea2_reference_edit_acc=krea2_reference_edit_accordion, qwen_image_edit_acc=qwen_image_edit_accordion, ref_img_acc=ref_img_accordion
             )
             inputs = [arch_comp, cat_comp]
             if aspect_ratio_comp:
@@ -163,6 +166,7 @@ def attach_event_handlers(ui_components, demo):
             if hidream_o1_ref_accordion: outputs2.append(hidream_o1_ref_accordion)
             if joyai_ref_accordion: outputs2.append(joyai_ref_accordion)
             if krea2_identity_edit_accordion: outputs2.append(krea2_identity_edit_accordion)
+            if krea2_reference_edit_accordion: outputs2.append(krea2_reference_edit_accordion)
             if qwen_image_edit_accordion: outputs2.append(qwen_image_edit_accordion)
             if ref_img_accordion: outputs2.append(ref_img_accordion)
             if pid_accordion: outputs2.append(pid_accordion)
@@ -184,7 +188,7 @@ def attach_event_handlers(ui_components, demo):
                     krea2_cn_types_list, krea2_cn_series_list, krea2_cn_filepaths_list,
                     arch_comp, ipa_preset_list, lora_accordion, cn_accordion, anima_cn_accordion, diffsynth_cn_accordion, krea2_cn_accordion, ipa_accordion, sd3_ipa_accordion, flux1_ipa_accordion, style_accordion, embedding_accordion, conditioning_accordion,
                     ref_latent_accordion, hidream_o1_ref_accordion, guidance_comp, prompt_comp, neg_prompt_comp, steps_comp, cfg_comp, sampler_comp, scheduler_comp,
-                    pid_acc=pid_accordion, vae_acc=vae_accordion, joyai_ref_acc=joyai_ref_accordion, krea2_identity_edit_acc=krea2_identity_edit_accordion, qwen_image_edit_acc=qwen_image_edit_accordion, ref_img_acc=ref_img_accordion
+                    pid_acc=pid_accordion, vae_acc=vae_accordion, joyai_ref_acc=joyai_ref_accordion, krea2_identity_edit_acc=krea2_identity_edit_accordion, krea2_reference_edit_acc=krea2_reference_edit_accordion, qwen_image_edit_acc=qwen_image_edit_accordion, ref_img_acc=ref_img_accordion
                 )
                 model_comp.change(fn=change_fn, inputs=inputs2, outputs=outputs2)
 
@@ -203,6 +207,7 @@ def attach_event_handlers(ui_components, demo):
         create_hidream_o1_reference_event_handlers(prefix, ui_components)
         create_joyai_reference_event_handlers(prefix, ui_components)
         create_krea2_identity_edit_event_handlers(prefix, ui_components)
+        create_krea2_reference_edit_event_handlers(prefix, ui_components)
         create_qwen_image_edit_event_handlers(prefix, ui_components)
         create_reference_image_event_handlers(prefix, ui_components)
         create_run_event(prefix, task_type, ui_components)

ui/events/run_handlers.py

@@ -54,6 +54,7 @@ def create_run_event(prefix: str, task_type: str, ui_components: dict):
     hidream_o1_reference_data_components = ui_components.get(f'all_hidream_o1_reference_components_flat_{prefix}', [])
     joyai_reference_data_components = ui_components.get(f'all_joyai_reference_components_flat_{prefix}', [])
     krea2_identity_edit_data_components = ui_components.get(f'all_krea2_identity_edit_components_flat_{prefix}', [])
+    krea2_reference_edit_data_components = ui_components.get(f'all_krea2_reference_edit_components_flat_{prefix}', [])
     qwen_image_edit_data_components = ui_components.get(f'all_qwen_image_edit_components_flat_{prefix}', [])
     reference_image_data_components = ui_components.get(f'all_reference_image_components_flat_{prefix}', [])
     
@@ -66,7 +67,7 @@ def create_run_event(prefix: str, task_type: str, ui_components: dict):
     all_chains = [
         lora_data_components, controlnet_data_components, anima_controlnet_lllite_data_components, diffsynth_controlnet_data_components, krea2_controlnet_data_components, ipadapter_data_components,
         sd3_ipadapter_data_components, flux1_ipadapter_data_components, style_data_components,
-        embedding_data_components, conditioning_data_components, reference_latent_data_components, hidream_o1_reference_data_components, joyai_reference_data_components, krea2_identity_edit_data_components, qwen_image_edit_data_components, reference_image_data_components
+        embedding_data_components, conditioning_data_components, reference_latent_data_components, hidream_o1_reference_data_components, joyai_reference_data_components, krea2_identity_edit_data_components, krea2_reference_edit_data_components, qwen_image_edit_data_components, reference_image_data_components
     ]
     for chain in all_chains:
         if chain:
@@ -98,6 +99,7 @@ def create_run_event(prefix: str, task_type: str, ui_components: dict):
         assign_chain_data('hidream_o1_reference_data', hidream_o1_reference_data_components)
         assign_chain_data('joyai_reference_data', joyai_reference_data_components)
         assign_chain_data('krea2_identity_edit_data', krea2_identity_edit_data_components)
+        assign_chain_data('krea2_reference_edit_data', krea2_reference_edit_data_components)
         assign_chain_data('qwen_image_edit_data', qwen_image_edit_data_components)
         assign_chain_data('reference_image_data', reference_image_data_components)
 

ui/shared/hires_fix_ui.py

@@ -7,7 +7,7 @@ from .ui_components import (
     create_conditioning_ui, create_vae_override_ui,
     create_model_architecture_filter_ui, create_category_filter_ui,
     create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
-    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_qwen_image_edit_ui
+    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_krea2_reference_edit_ui, create_qwen_image_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -103,6 +103,7 @@ def create_ui():
         components.update(create_hidream_o1_reference_ui(prefix))
         components.update(create_joyai_reference_ui(prefix))
         components.update(create_krea2_identity_edit_ui(prefix))
+        components.update(create_krea2_reference_edit_ui(prefix))
         components.update(create_qwen_image_edit_ui(prefix))
         components.update(create_reference_image_ui(prefix))
         components.update(create_vae_override_ui(prefix))

ui/shared/img2img_ui.py

@@ -7,7 +7,7 @@ from .ui_components import (
     create_conditioning_ui, create_vae_override_ui,
     create_model_architecture_filter_ui, create_category_filter_ui,
     create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
-    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_qwen_image_edit_ui
+    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_krea2_reference_edit_ui, create_qwen_image_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -84,6 +84,7 @@ def create_ui():
         components.update(create_hidream_o1_reference_ui(prefix))
         components.update(create_joyai_reference_ui(prefix))
         components.update(create_krea2_identity_edit_ui(prefix))
+        components.update(create_krea2_reference_edit_ui(prefix))
         components.update(create_qwen_image_edit_ui(prefix))
         components.update(create_reference_image_ui(prefix))
         components.update(create_vae_override_ui(prefix))

ui/shared/inpaint_ui.py

@@ -6,7 +6,7 @@ from .ui_components import (
     create_conditioning_ui, create_vae_override_ui,
     create_model_architecture_filter_ui, create_category_filter_ui,
     create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
-    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_qwen_image_edit_ui
+    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_krea2_reference_edit_ui, create_qwen_image_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -116,6 +116,7 @@ def create_ui():
             components.update(create_hidream_o1_reference_ui(prefix))
             components.update(create_joyai_reference_ui(prefix))
             components.update(create_krea2_identity_edit_ui(prefix))
+            components.update(create_krea2_reference_edit_ui(prefix))
             components.update(create_qwen_image_edit_ui(prefix))
             components.update(create_reference_image_ui(prefix))
             components.update(create_vae_override_ui(prefix))

ui/shared/outpaint_ui.py

@@ -7,7 +7,7 @@ from .ui_components import (
     create_conditioning_ui, create_vae_override_ui,
     create_model_architecture_filter_ui, create_category_filter_ui,
     create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
-    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_qwen_image_edit_ui
+    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_krea2_reference_edit_ui, create_qwen_image_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -99,6 +99,7 @@ def create_ui():
         components.update(create_hidream_o1_reference_ui(prefix))
         components.update(create_joyai_reference_ui(prefix))
         components.update(create_krea2_identity_edit_ui(prefix))
+        components.update(create_krea2_reference_edit_ui(prefix))
         components.update(create_qwen_image_edit_ui(prefix))
         components.update(create_reference_image_ui(prefix))
         components.update(create_vae_override_ui(prefix))

ui/shared/txt2img_ui.py

@@ -6,7 +6,7 @@ from .ui_components import (
     create_conditioning_ui, create_vae_override_ui,
     create_model_architecture_filter_ui, create_category_filter_ui,
     create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
-    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_qwen_image_edit_ui,
+    create_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_krea2_reference_edit_ui, create_qwen_image_edit_ui,
     create_pid_ui
 )
 
@@ -56,6 +56,7 @@ def create_ui():
         components.update(create_hidream_o1_reference_ui(prefix))
         components.update(create_joyai_reference_ui(prefix))
         components.update(create_krea2_identity_edit_ui(prefix))
+        components.update(create_krea2_reference_edit_ui(prefix))
         components.update(create_qwen_image_edit_ui(prefix))
         components.update(create_reference_image_ui(prefix))
         components.update(create_vae_override_ui(prefix))

ui/shared/ui_components.py

@@ -775,7 +775,7 @@ def create_krea2_identity_edit_ui(prefix: str, max_units=2):
     components = {}
     key = lambda name: f"{name}_{prefix}"
 
-    with gr.Accordion("Krea-2 Identity Edit Setting", open=False, visible=('krea2_identity_edit' in default_enabled_chains)) as ref_accordion:
+    with gr.Accordion("Krea2 Identity Edit Settings", open=False, visible=('krea2_identity_edit' in default_enabled_chains)) as ref_accordion:
         components[key('krea2_identity_edit_accordion')] = ref_accordion
         gr.Markdown("💡 **Tip:** Processed using the [lbouaraba/comfyui-krea2edit](https://github.com/lbouaraba/comfyui-krea2edit) node. (Krea-2-Turbo recommended, Krea-2-Raw need set ZeroGPU Duration (s) to 120 ) In txt2img mode, adding a single reference image performs an **Image Edit**, while adding multiple images performs an **Image Combine**.")
         
@@ -827,4 +827,33 @@ def create_qwen_image_edit_ui(prefix: str, max_units=3):
         
         components[key('all_qwen_image_edit_components_flat')] = ref_image_inputs
 
+    return components
+
+def create_krea2_reference_edit_ui(prefix: str, max_units=3):
+    components = {}
+    key = lambda name: f"{name}_{prefix}"
+
+    with gr.Accordion("Krea2 Style Reference Edit Settings", open=False, visible=('krea2_reference_edit' in default_enabled_chains)) as ref_accordion:
+        components[key('krea2_reference_edit_accordion')] = ref_accordion
+        gr.Markdown("💡 **Tip:** (Krea-2-Turbo recommended) Add style reference images to perform style reference editing.")
+        
+        ref_image_groups = []
+        ref_image_inputs = []
+        with gr.Row():
+            for i in range(max_units):
+                with gr.Column(visible=(i < 1), min_width=160) as img_col:
+                    img_comp = gr.Image(type="pil", label=f"Ref. {i+1}", sources=["upload"], height=150)
+                    ref_image_groups.append(img_col)
+                    ref_image_inputs.append(img_comp)
+        
+        components[key('krea2_reference_edit_rows')] = ref_image_groups
+        components[key('krea2_reference_edit_images')] = ref_image_inputs
+
+        with gr.Row():
+            components[key('add_krea2_reference_edit_button')] = gr.Button("✚ Add Reference Image")
+            components[key('delete_krea2_reference_edit_button')] = gr.Button("➖ Delete Reference Image", visible=False)
+        components[key('krea2_reference_edit_count_state')] = gr.State(1)
+        
+        components[key('all_krea2_reference_edit_components_flat')] = ref_image_inputs
+
     return components

yaml/file_list.yaml

@@ -414,6 +414,10 @@ file:
       source: "hf"
       repo_id: "Comfy-Org/Mage-Flow"
       repository_file_path: "diffusion_models/mage_flow_int8_convrot.safetensors"
+    - filename: "mage_flow_base_bf16.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/Mage-Flow"
+      repository_file_path: "diffusion_models/mage_flow_base_bf16.safetensors"
     - filename: "mage_flow_edit_turbo_int8_convrot.safetensors"
       source: "hf"
       repo_id: "Comfy-Org/Mage-Flow"
@@ -422,6 +426,10 @@ file:
       source: "hf"
       repo_id: "Comfy-Org/Mage-Flow"
       repository_file_path: "diffusion_models/mage_flow_edit_int8_convrot.safetensors"
+    - filename: "mage_flow_edit_base_bf16.safetensors"
+      source: "hf"
+      repo_id: "Comfy-Org/Mage-Flow"
+      repository_file_path: "diffusion_models/mage_flow_edit_base_bf16.safetensors"
     # JoyAI-Image
     - filename: "joyai_image_edit_int8_convrot.safetensors"
       source: "hf"
@@ -771,6 +779,11 @@ file:
       source: "hf"
       repo_id: "conradlocke/krea2-identity-edit"
       repository_file_path: "krea2_identity_edit_v1_2.safetensors"
+    # Krea2 Style Reference Edit
+    - filename: "krea2_style_reference.safetensors"
+      source: "hf"
+      repo_id: "ostris/krea2_turbo_style_reference"
+      repository_file_path: "krea2_style_reference.safetensors"
     # Krea2 ControlNet
     - filename: "depth-control-lora.safetensors"
       source: "hf"

yaml/image_gen_features.yaml

@@ -21,6 +21,7 @@ krea-2:
     - lora
     - krea2_controlnet
     - krea2_identity_edit
+    - krea2_reference_edit
     - pid
 
 boogu-image:

yaml/injectors.yaml

@@ -33,6 +33,8 @@ injector_definitions:
     module: "chain_injectors.joyai_reference_injector"
   dynamic_krea2_identity_edit_chains:
     module: "chain_injectors.krea2_identity_edit_injector"
+  dynamic_krea2_reference_edit_chains:
+    module: "chain_injectors.krea2_reference_edit_injector"
   dynamic_qwen_image_edit_chains:
     module: "chain_injectors.qwen_image_edit_injector"
   dynamic_reference_image_chains:
@@ -54,6 +56,7 @@ injector_order:
   - dynamic_controlnet_chains
   - dynamic_krea2_controlnet_chains
   - dynamic_krea2_identity_edit_chains
+  - dynamic_krea2_reference_edit_chains
   - dynamic_qwen_image_edit_chains
   - dynamic_anima_controlnet_lllite_chains
   - dynamic_hidream_o1_smoothing_chains

yaml/model_list.yaml

@@ -25,6 +25,11 @@ Checkpoint:
           unet: "mage_flow_int8_convrot.safetensors"
           clip: "qwen3vl_4b_fp8_scaled.safetensors"
           vae: "mage_flow_vae_bf16.safetensors"
+      - display_name: "Mage-Flow-Base"
+        components:
+          unet: "mage_flow_base_bf16.safetensors"
+          clip: "qwen3vl_4b_fp8_scaled.safetensors"
+          vae: "mage_flow_vae_bf16.safetensors"
       - display_name: "Mage-Flow-Edit-Turbo"
         components:
           unet: "mage_flow_edit_turbo_int8_convrot.safetensors"
@@ -35,6 +40,11 @@ Checkpoint:
           unet: "mage_flow_edit_int8_convrot.safetensors"
           clip: "qwen3vl_4b_fp8_scaled.safetensors"
           vae: "mage_flow_vae_bf16.safetensors"
+      - display_name: "Mage-Flow-Edit-Base"
+        components:
+          unet: "mage_flow_edit_base_bf16.safetensors"
+          clip: "qwen3vl_4b_fp8_scaled.safetensors"
+          vae: "mage_flow_vae_bf16.safetensors"
   JoyAI-Image:
     latent_type: sd3_latent
     models: