Krea-2 Identity Edit injector: Add image scaling before VAE encode.

chain_injectors/krea2_identity_edit_injector.py

@@ -1,6 +1,18 @@
 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
@@ -52,6 +64,11 @@ def inject(assembler, chain_definition, chain_items):
     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:
@@ -61,11 +78,10 @@ def inject(assembler, chain_definition, chain_items):
         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 = create_node(assembler, "LoraLoaderModelOnly", "Load LoRA (Krea2 Identity Edit)")
     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 = []
@@ -73,23 +89,29 @@ def inject(assembler, chain_definition, chain_items):
 
     for i, img_filename in enumerate(valid_images):
         load_id = assembler._get_unique_id()
-        load_node = assembler._get_node_template("LoadImage")
+        load_node = create_node(assembler, "LoadImage", f"Load Image (Ref {i+1})")
         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)
+
+        scale_id = assembler._get_unique_id()
+        scale_node = create_node(assembler, "ImageScaleToTotalPixels", f"Scale Reference {i+1}")
+        scale_node['inputs']['upscale_method'] = "lanczos"
+        scale_node['inputs']['megapixels'] = 1.0
+        scale_node['inputs']['resolution_steps'] = 1
+        scale_node['inputs']['image'] = [load_id, 0]
+        assembler.workflow[scale_id] = scale_node
+        image_ids.append(scale_id)
 
         vae_enc_id = assembler._get_unique_id()
-        vae_enc_node = assembler._get_node_template("VAEEncode")
-        vae_enc_node['inputs']['pixels'] = [load_id, 0]
+        vae_enc_node = create_node(assembler, "VAEEncode", f"VAE Encode (Ref {i+1})")
+        vae_enc_node['inputs']['pixels'] = [scale_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 = create_node(assembler, "Krea2EditModelPatch", "Krea2 Edit (source patch)")
     patch_node['inputs']['ref_boost'] = 4
     patch_node['inputs']['ref_boost_a'] = 1
     patch_node['inputs']['fit_mode'] = "fit"
@@ -104,7 +126,6 @@ def inject(assembler, chain_definition, chain_items):
         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]

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.51.9
-comfyui-workflow-templates==0.11.50
+comfyui-workflow-templates==0.11.54
 comfyui-embedded-docs==0.5.10
 torch
 torchsde