| 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.") |
|
|
