Add Krea-2 Identity Edit injector.

README.md

@@ -65,6 +65,7 @@ models:
   - Comfy-Org/Wan_2.1_ComfyUI_repackaged
   - Comfy-Org/z_image
   - Comfy-Org/z_image_turbo
+  - conradlocke/krea2-identity-edit
   - cyberdelia/CyberRealisticPony
   - diffusionmodels1254ani/hassakuAnima
   - diffusionmodels1254ani/kirazuriAnima_v30AnimaBase1
@@ -111,6 +112,7 @@ models:
   - stabilityai/stable-diffusion-3.5-controlnets
   - stabilityai/stable-diffusion-xl-base-1.0
   - stable-diffusion-v1-5/stable-diffusion-v1-5
+  - tsukiyomi/krea2_raw-nvfp4
   - Wenaka/NoobAI_XL_Inpainting_ControlNet_Full
   - xinsir/anime-painter
   - xinsir/controlnet-canny-sdxl-1.0

chain_injectors/krea2_identity_edit_injector.py

@@ -0,0 +1,159 @@
+import os
+from utils.app_utils import ensure_file_downloaded
+
+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[:2]
+
+    lora_filename = "krea2_identity_edit_v1_2.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 Identity 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
+
+    latent_connection = assembler.workflow[ksampler_id]['inputs'].get('latent_image')
+    if not latent_connection:
+        print(f"Warning: KSampler node '{ksampler_name}' is missing 'latent_image' 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]
+
+    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')
+
+    lora_loader_id = assembler._get_unique_id()
+    lora_loader_node = assembler._get_node_template("LoraLoaderModelOnly")
+    lora_loader_node['inputs']['lora_name'] = lora_filename
+    lora_loader_node['inputs']['strength_model'] = 1.0
+    lora_loader_node['inputs']['model'] = current_model_connection
+    lora_loader_node['_meta']['title'] = "Load LoRA (Krea2 Identity Edit)"
+    assembler.workflow[lora_loader_id] = lora_loader_node
+
+    image_ids = []
+    vae_encode_ids = []
+
+    for i, img_filename in enumerate(valid_images):
+        load_id = assembler._get_unique_id()
+        load_node = assembler._get_node_template("LoadImage")
+        load_node['inputs']['image'] = img_filename
+        load_node['_meta']['title'] = f"Load Image (Ref {i+1})"
+        assembler.workflow[load_id] = load_node
+        image_ids.append(load_id)
+
+        vae_enc_id = assembler._get_unique_id()
+        vae_enc_node = assembler._get_node_template("VAEEncode")
+        vae_enc_node['inputs']['pixels'] = [load_id, 0]
+        if vae_connection:
+            vae_enc_node['inputs']['vae'] = vae_connection
+        vae_enc_node['_meta']['title'] = f"VAE Encode (Ref {i+1})"
+        assembler.workflow[vae_enc_id] = vae_enc_node
+        vae_encode_ids.append(vae_enc_id)
+
+    patch_id = assembler._get_unique_id()
+    patch_node = assembler._get_node_template("Krea2EditModelPatch")
+    patch_node['inputs']['ref_boost'] = 4
+    patch_node['inputs']['ref_boost_a'] = 1
+    patch_node['inputs']['fit_mode'] = "fit"
+    patch_node['inputs']['model'] = [lora_loader_id, 0]
+    patch_node['inputs']['source_latent'] = [vae_encode_ids[0], 0]
+    if vae_connection:
+        patch_node['inputs']['vae'] = vae_connection
+    patch_node['inputs']['source_image'] = [image_ids[0], 0]
+    patch_node['inputs']['target_latent'] = latent_connection
+
+    if len(valid_images) > 1:
+        patch_node['inputs']['source_latent_b'] = [vae_encode_ids[1], 0]
+        patch_node['inputs']['source_image_b'] = [image_ids[1], 0]
+
+    patch_node['_meta']['title'] = "Krea2 Edit (source patch)"
+    assembler.workflow[patch_id] = patch_node
+
+    assembler.workflow[ksampler_id]['inputs']['model'] = [patch_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', '')
+
+    neg_text = ""
+    if neg_prompt_id and neg_prompt_id in assembler.workflow:
+        neg_text = assembler.workflow[neg_prompt_id]['inputs'].get('text', '')
+
+    pos_grounded_id = assembler._get_unique_id()
+    pos_grounded_node = assembler._get_node_template("Krea2EditGroundedEncode")
+    pos_grounded_node['inputs']['prompt'] = pos_text
+    pos_grounded_node['inputs']['grounding_px'] = 768
+    pos_grounded_node['inputs']['system_prompt'] = ""
+    if clip_connection:
+        pos_grounded_node['inputs']['clip'] = clip_connection
+    pos_grounded_node['inputs']['image'] = [image_ids[0], 0]
+    if len(valid_images) > 1:
+        pos_grounded_node['inputs']['image_b'] = [image_ids[1], 0]
+    pos_grounded_node['_meta']['title'] = "Krea2 Edit (grounded encode positive)"
+    assembler.workflow[pos_grounded_id] = pos_grounded_node
+
+    assembler.workflow[ksampler_id]['inputs']['positive'] = [pos_grounded_id, 0]
+
+    neg_grounded_id = assembler._get_unique_id()
+    neg_grounded_node = assembler._get_node_template("Krea2EditGroundedEncode")
+    neg_grounded_node['inputs']['prompt'] = neg_text
+    neg_grounded_node['inputs']['grounding_px'] = 768
+    neg_grounded_node['inputs']['system_prompt'] = ""
+    if clip_connection:
+        neg_grounded_node['inputs']['clip'] = clip_connection
+    neg_grounded_node['inputs']['image'] = [image_ids[0], 0]
+    if len(valid_images) > 1:
+        neg_grounded_node['inputs']['image_b'] = [image_ids[1], 0]
+    neg_grounded_node['_meta']['title'] = "Krea2 Edit (grounded encode negative)"
+    assembler.workflow[neg_grounded_id] = neg_grounded_node
+
+    assembler.workflow[ksampler_id]['inputs']['negative'] = [neg_grounded_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 Identity Edit injector applied with {len(valid_images)} reference image(s). Original CLIPTextEncode nodes removed.")

comfy_integration/setup.py

@@ -88,6 +88,14 @@ def initialize_comfyui():
     else:
         print("✅ comfyui-krea2-controlnet extension already exists.")
 
+    # 6. comfyui-krea2edit
+    krea2edit_nodes_path = os.path.join(APP_DIR, "custom_nodes", "comfyui-krea2edit")
+    if not os.path.exists(krea2edit_nodes_path):
+        os.system(f"git clone https://github.com/lbouaraba/comfyui-krea2edit.git {krea2edit_nodes_path}")
+        print("✅ comfyui-krea2edit extension cloned.")
+    else:
+        print("✅ comfyui-krea2edit extension already exists.")
+
     print(f"✅ Current working directory is: {os.getcwd()}")
     
     import comfy.model_management

core/pipelines/pipeline_input_processor.py

@@ -357,6 +357,17 @@ def process_pipeline_inputs(ui_inputs: Dict[str, Any], progress: gr.Progress, wo
                 temp_files_to_clean.append(temp_path)
                 active_joyai_reference.append(os.path.basename(temp_path))
 
+    krea2_identity_edit_data = ui_inputs.get('krea2_identity_edit_data', [])
+    active_krea2_identity_edit = []
+    if krea2_identity_edit_data:
+        for img in krea2_identity_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_id_ref_{random.randint(1000, 9999)}.png")
+                img.save(temp_path, "PNG")
+                temp_files_to_clean.append(temp_path)
+                active_krea2_identity_edit.append(os.path.basename(temp_path))
+
     reference_image_data = ui_inputs.get('reference_image_data', [])
     active_reference_images = []
     if reference_image_data:
@@ -410,6 +421,7 @@ def process_pipeline_inputs(ui_inputs: Dict[str, Any], progress: gr.Progress, wo
         "active_reference_latents": active_reference_latents,
         "active_hidream_o1_reference": active_hidream_o1_reference,
         "active_joyai_reference": active_joyai_reference,
+        "active_krea2_identity_edit": active_krea2_identity_edit,
         "active_reference_images": active_reference_images,
         "active_conditioning": active_conditioning,
         "temp_files_to_clean": temp_files_to_clean

core/pipelines/sd_image_pipeline.py

@@ -117,6 +117,7 @@ class SdImagePipeline(BasePipeline):
             active_reference_latents = processed["active_reference_latents"]
             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_reference_images = processed.get("active_reference_images", [])
             active_conditioning = processed["active_conditioning"]
 
@@ -175,6 +176,7 @@ class SdImagePipeline(BasePipeline):
                 "reference_latent_chain": active_reference_latents,
                 "hidream_o1_reference_chain": active_hidream_o1_reference,
                 "joyai_reference_chain": active_joyai_reference,
+                "krea2_identity_edit_chain": active_krea2_identity_edit,
                 "reference_image_chain": active_reference_images,
                 "vae_chain": [ui_inputs.get('vae_name')] if ui_inputs.get('vae_name') else [],
                 "hidream_o1_smoothing_chain": hidream_o1_smoothing_data,

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

@@ -54,6 +54,14 @@ dynamic_krea2_controlnet_chains:
     ksampler_node: "ksampler"
     vae_source: "vae_loader:0"
 
+dynamic_krea2_identity_edit_chains:
+  krea2_identity_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"

ui/events/chain_handlers.py

@@ -856,4 +856,37 @@ def create_conditioning_event_handlers(prefix, ui_components):
     add_outputs = [count_state, add_button, del_button] + rows
     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)
+    del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False)
+
+
+def create_krea2_identity_edit_event_handlers(prefix, ui_components):
+    ref_rows = ui_components.get(f'krea2_identity_edit_rows_{prefix}')
+    if not ref_rows: return
+    count_state = ui_components[f'krea2_identity_edit_count_state_{prefix}']
+    add_button = ui_components[f'add_krea2_identity_edit_button_{prefix}']
+    del_button = ui_components[f'delete_krea2_identity_edit_button_{prefix}']
+    images = ui_components[f'krea2_identity_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 < 2),
+            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)

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, 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, ref_img_acc=None):
     def update_fn(*args):
         arch = args[0]
         category = args[1]
@@ -69,6 +69,7 @@ def make_update_fn(m_comp, cat_comp, cs_comp, ar_comp, width_comp, height_comp,
         if ref_latent_acc: updates[ref_latent_acc] = gr.update(visible=('reference_latent' in enabled_chains))
         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 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))
         if vae_acc: updates[vae_acc] = gr.update(visible=('vae' in enabled_chains))
@@ -144,7 +145,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, 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, ref_img_acc=None):
     def change_fn(*args):
         model_name = args[0]
         idx = 1
@@ -201,6 +202,7 @@ def make_model_change_fn(cat_comp_ref, cs_comp, ar_comp, width_comp, height_comp
         if ref_latent_acc: updates[ref_latent_acc] = gr.update(visible=('reference_latent' in enabled_chains))
         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 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))
         if vae_acc: updates[vae_acc] = gr.update(visible=('vae' in enabled_chains))

ui/events/main.py

@@ -14,6 +14,7 @@ from .chain_handlers import (
     create_reference_latent_event_handlers,
     create_hidream_o1_reference_event_handlers,
     create_joyai_reference_event_handlers,
+    create_krea2_identity_edit_event_handlers,
     create_reference_image_event_handlers
 )
 from .change_handlers import (
@@ -69,6 +70,7 @@ def attach_event_handlers(ui_components, demo):
         ref_latent_accordion = ui_components.get(f'reference_latent_accordion_{prefix}')
         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}')
         ref_img_accordion = ui_components.get(f'reference_image_accordion_{prefix}')
         pid_accordion = ui_components.get(f'pid_accordion_{prefix}')
         vae_accordion = ui_components.get(f'vae_accordion_{prefix}')
@@ -108,6 +110,7 @@ def attach_event_handlers(ui_components, demo):
             if ref_latent_accordion: outputs.append(ref_latent_accordion)
             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 ref_img_accordion: outputs.append(ref_img_accordion)
             if pid_accordion: outputs.append(pid_accordion)
             if vae_accordion: outputs.append(vae_accordion)
@@ -123,7 +126,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, 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, ref_img_acc=ref_img_accordion
             )
             inputs = [arch_comp, cat_comp]
             if aspect_ratio_comp:
@@ -156,6 +159,7 @@ def attach_event_handlers(ui_components, demo):
             if ref_latent_accordion: outputs2.append(ref_latent_accordion)
             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 ref_img_accordion: outputs2.append(ref_img_accordion)
             if pid_accordion: outputs2.append(pid_accordion)
             if vae_accordion: outputs2.append(vae_accordion)
@@ -176,7 +180,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, 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, ref_img_acc=ref_img_accordion
                 )
                 model_comp.change(fn=change_fn, inputs=inputs2, outputs=outputs2)
 
@@ -194,6 +198,7 @@ def attach_event_handlers(ui_components, demo):
         create_reference_latent_event_handlers(prefix, ui_components)
         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_reference_image_event_handlers(prefix, ui_components)
         create_run_event(prefix, task_type, ui_components)
 

ui/events/run_handlers.py

@@ -53,6 +53,7 @@ def create_run_event(prefix: str, task_type: str, ui_components: dict):
     reference_latent_data_components = ui_components.get(f'all_reference_latent_components_flat_{prefix}', [])
     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}', [])
     reference_image_data_components = ui_components.get(f'all_reference_image_components_flat_{prefix}', [])
     
     run_inputs_map['vae_source'] = ui_components.get(f'vae_source_{prefix}')
@@ -64,7 +65,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, 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, reference_image_data_components
     ]
     for chain in all_chains:
         if chain:
@@ -95,6 +96,7 @@ def create_run_event(prefix: str, task_type: str, ui_components: dict):
         assign_chain_data('reference_latent_data', reference_latent_data_components)
         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('reference_image_data', reference_image_data_components)
 
         return ui_dict

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_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -102,6 +102,7 @@ def create_ui():
         components.update(create_reference_latent_ui(prefix))
         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_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_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -83,6 +83,7 @@ def create_ui():
         components.update(create_reference_latent_ui(prefix))
         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_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_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -115,6 +115,7 @@ def create_ui():
             components.update(create_reference_latent_ui(prefix))
             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_reference_image_ui(prefix))
             components.update(create_vae_override_ui(prefix))
         components[f'accordion_wrapper_{prefix}'] = accordion_wrapper

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_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui
 )
 
 default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
@@ -98,6 +98,7 @@ def create_ui():
         components.update(create_reference_latent_ui(prefix))
         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_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_reference_latent_ui, create_hidream_o1_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui,
     create_pid_ui
 )
 
@@ -55,6 +55,7 @@ def create_ui():
         components.update(create_reference_latent_ui(prefix))
         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_reference_image_ui(prefix))
         components.update(create_vae_override_ui(prefix))
         components.update(create_pid_ui(prefix))

ui/shared/ui_components.py

@@ -769,4 +769,33 @@ def create_pid_ui(prefix: str):
                 interactive=True
             )
             
+    return components
+
+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:
+        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**.")
+        
+        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_identity_edit_rows')] = ref_image_groups
+        components[key('krea2_identity_edit_images')] = ref_image_inputs
+
+        with gr.Row():
+            components[key('add_krea2_identity_edit_button')] = gr.Button("✚ Add Reference Image")
+            components[key('delete_krea2_identity_edit_button')] = gr.Button("➖ Delete Reference Image", visible=False)
+        components[key('krea2_identity_edit_count_state')] = gr.State(1)
+        
+        components[key('all_krea2_identity_edit_components_flat')] = ref_image_inputs
+
     return components

utils/app_utils.py

@@ -498,6 +498,72 @@ def ensure_controlnet_model_downloaded(filename: str, progress):
                 pass
         raise gr.Error(f"Failed to download ControlNet model '{filename}': {e}")
 
+def ensure_file_downloaded(filename: str, progress=None):
+    if not filename or filename == "None":
+        return
+
+    download_info = ALL_FILE_DOWNLOAD_MAP.get(filename)
+    if not download_info:
+        print(f"⚠️ Warning: File '{filename}' not found in configuration (file_list.yaml). Cannot download.")
+        return
+
+    category = download_info.get("category", "loras")
+    dest_dir = CATEGORY_TO_DIR_MAP.get(category, LORA_DIR)
+    dest_path = os.path.join(dest_dir, filename)
+
+    if os.path.lexists(dest_path):
+        if not os.path.exists(dest_path):
+            print(f"⚠️ Found and removed broken symlink: {dest_path}")
+            os.remove(dest_path)
+        else:
+            return
+
+    source = download_info.get("source")
+    try:
+        if source == "hf":
+            repo_id = download_info.get("repo_id")
+            repo_filename = download_info.get("repository_file_path", filename)
+            if not repo_id:
+                raise ValueError("repo_id is missing for Hugging Face download.")
+
+            if progress and callable(progress):
+                progress(0, desc=f"Downloading: {filename}")
+            cached_path = hf_hub_download(repo_id=repo_id, filename=repo_filename, token=os.environ.get("HF_TOKEN"))
+            os.makedirs(dest_dir, exist_ok=True)
+            os.symlink(cached_path, dest_path)
+            print(f"✅ Symlinked '{cached_path}' to '{dest_path}'")
+            if progress and callable(progress):
+                progress(1.0, desc=f"Downloaded: {filename}")
+
+        elif source == "civitai":
+            model_version_id = download_info.get("model_version_id")
+            if not model_version_id:
+                raise ValueError("model_version_id is missing for Civitai download.")
+
+            file_info = get_civitai_file_info(model_version_id)
+            if not file_info or not file_info.get('downloadUrl'):
+                raise ConnectionError(f"Could not get download URL for Civitai model version ID {model_version_id}")
+
+            status = download_file(
+                file_info['downloadUrl'],
+                dest_path,
+                api_key=os.environ.get("CIVITAI_API_KEY", ""),
+                progress=progress,
+                desc=f"Downloading: {filename}"
+            )
+            if "Failed" in status:
+                raise ConnectionError(status)
+        else:
+            raise NotImplementedError(f"Download source '{source}' is not implemented for '{filename}'.")
+
+    except Exception as e:
+        if os.path.lexists(dest_path):
+            try:
+                os.remove(dest_path)
+            except OSError:
+                pass
+        raise gr.Error(f"Failed to download file '{filename}': {e}")
+
 def load_ipadapter_presets():
     global IPADAPTER_PRESETS
     if IPADAPTER_PRESETS is not None:

yaml/file_list.yaml

@@ -401,10 +401,10 @@ file:
       source: "hf"
       repo_id: "Comfy-Org/Krea-2"
       repository_file_path: "diffusion_models/krea2_turbo_nvfp4.safetensors"
-    - filename: "krea2_raw_fp8_scaled.safetensors"
+    - filename: "krea2_raw-nvfp4.safetensors"
       source: "hf"
-      repo_id: "Comfy-Org/Krea-2"
-      repository_file_path: "diffusion_models/krea2_raw_fp8_scaled.safetensors"
+      repo_id: "tsukiyomi/krea2_raw-nvfp4"
+      repository_file_path: "krea2_raw-nvfp4.safetensors"
     # Mage-Flow
     - filename: "mage_flow_turbo_int8_convrot.safetensors"
       source: "hf"
@@ -745,6 +745,11 @@ file:
       repo_id: "black-forest-labs/FLUX.1-Redux-dev"
       repository_file_path: "flux1-redux-dev.safetensors"
   loras: 
+    # Krea2 Identity Edit
+    - filename: "krea2_identity_edit_v1_2.safetensors"
+      source: "hf"
+      repo_id: "conradlocke/krea2-identity-edit"
+      repository_file_path: "krea2_identity_edit_v1_2.safetensors"
     # Krea2 ControlNet
     - filename: "depth-control-lora.safetensors"
       source: "hf"

yaml/image_gen_features.yaml

@@ -20,7 +20,7 @@ krea-2:
   enabled_chains:
     - lora
     - krea2_controlnet
-    - conditioning
+    - krea2_identity_edit
     - pid
 
 boogu-image:

yaml/injectors.yaml

@@ -31,6 +31,8 @@ injector_definitions:
     module: "chain_injectors.hidream_o1_reference_injector"
   dynamic_joyai_reference_chains:
     module: "chain_injectors.joyai_reference_injector"
+  dynamic_krea2_identity_edit_chains:
+    module: "chain_injectors.krea2_identity_edit_injector"
   dynamic_reference_image_chains:
     module: "chain_injectors.reference_image_injector"
   dynamic_pid_chains:
@@ -49,6 +51,7 @@ injector_order:
   - dynamic_style_chains
   - dynamic_controlnet_chains
   - dynamic_krea2_controlnet_chains
+  - dynamic_krea2_identity_edit_chains
   - dynamic_anima_controlnet_lllite_chains
   - dynamic_hidream_o1_smoothing_chains
   - dynamic_hidream_o1_reference_chains

yaml/model_list.yaml

@@ -9,7 +9,7 @@ Checkpoint:
           vae: "qwen_image_vae.safetensors"
       - display_name: "Krea-2-Raw"
         components:
-          unet: "krea2_raw_fp8_scaled.safetensors"
+          unet: "krea2_raw-nvfp4.safetensors"
           clip: "qwen3vl_4b_fp8_scaled.safetensors"
           vae: "qwen_image_vae.safetensors"
   Mage-Flow: