From fb755f76755b25d710fb633b1bcc2d70d2f8612d Mon Sep 17 00:00:00 2001 From: molbal Date: Tue, 14 Jul 2026 10:25:11 +0200 Subject: [PATCH] placeholder --- .gitignore | 36 ++ README.md | 31 ++ __init__.py | 3 + js/docs/Krea2ModelPatcher.md | 34 ++ js/docs/Krea2ModelPatcher/en.md | 34 ++ js/docs/Krea2ReferenceEncoder.md | 68 ++++ js/docs/Krea2ReferenceEncoder/en.md | 68 ++++ nodes.py | 506 ++++++++++++++++++++++++++++ 8 files changed, 780 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 js/docs/Krea2ModelPatcher.md create mode 100644 js/docs/Krea2ModelPatcher/en.md create mode 100644 js/docs/Krea2ReferenceEncoder.md create mode 100644 js/docs/Krea2ReferenceEncoder/en.md create mode 100644 nodes.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e2a936d --- /dev/null +++ b/.gitignore @@ -0,0 +1,36 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +*.egg-info/ +dist/ +build/ + +# Virtual environments +venv/ +.env/ +.env.local +.env.*.local + +# IDE / editor +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Logs / temp +*.log +*.tmp +temp/ +tmp/ + +# ComfyUI / model outputs +output/ +models/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..0313abf --- /dev/null +++ b/README.md @@ -0,0 +1,31 @@ +# ComfyUI-Krea2-MultiRef (temporary README) + +ComfyUI custom nodes for Krea2-style multi-reference image editing. + +## Nodes + +- **Krea2 Model Patcher** — patches a Flux diffusion model to consume reference latents during sampling. +- **Krea2 Reference Text Encode** — encodes a prompt together with up to 4 reference images through the Krea2 Qwen3-VL text encoder. Connect a VAE to also attach reference latents. +- **Krea2 Reference Encoder (latent only)** — attaches pre-encoded reference latents to conditioning without re-encoding text. +- **Krea2 Diagnostic** — prints diffusion model internals for debugging. + +## Quick start + +1. Place this folder under `ComfyUI/custom_nodes/`. +2. Restart ComfyUI. +3. In a workflow: + - Load your Flux model. + - Add `Krea2 Model Patcher` and connect the model. + - Add `Krea2 Reference Text Encode`, connect CLIP + VAE + reference images, and write your prompt. + - Feed the patched model and conditioning into a KSampler as usual. + +## Requirements + +- ComfyUI with Krea2/Qwen3-VL text encoder support. +- `einops` (usually already installed with ComfyUI). + +## Notes + +- Reference images are downscaled to a max of 512×512 pixels for the VLM. +- Reference latents are capped at 1024×1024 pixels and snapped to multiples of 16. +- This README is a placeholder and will be expanded. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d721463 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/js/docs/Krea2ModelPatcher.md b/js/docs/Krea2ModelPatcher.md new file mode 100644 index 0000000..ec85610 --- /dev/null +++ b/js/docs/Krea2ModelPatcher.md @@ -0,0 +1,34 @@ +# Krea2 Model Patcher + +Patches the **Krea2** native text-to-image model with a trained image-adapter so it can accept reference-image conditioning alongside a text prompt. + +## Background + +Krea2 is a diffusion model released by Krea AI. This adapter was trained on semi-synthetic data — organic target images paired with synthetic reference images and prompts — to teach the model to follow reference imagery while still obeying text instructions. The adapter injects reference information via **RoPE encoding**, enabling the model to attend to up to 4 reference images at once. + +## Inputs + +| Name | Type | Description | +| --- | --- | --- | +| `model` | MODEL | The loaded Krea2 base model. | +| `patch_path` | STRING | Path to the adapter safetensors file (e.g. `krea2_image_adapter_v1_000001500.safetensors`). Relative paths resolve from the ComfyUI root. | + +## Outputs + +| Name | Type | Description | +| --- | --- | --- | +| `model` | MODEL | The patched Krea2 model, ready to consume reference conditioning. | + +## Usage + +1. Load the Krea2 checkpoint with a **Load Diffusion Model** node. +2. Connect its `MODEL` output to this node. +3. Set `patch_path` to the location of the adapter weights. +4. Wire the patched `MODEL` into a sampler (e.g. `KSampler`). +5. Pair with **Krea2 Reference Encoder** to supply the actual reference images. + +## Notes + +- This node must be used with the Krea2 base model. Using it with an unrelated checkpoint will produce undefined behaviour. +- The adapter is applied at full strength (`1.0`). Strength control may be added in a future version. +- Without a **Krea2 Reference Encoder** connected downstream, the patched model behaves identically to the unpatched base model. diff --git a/js/docs/Krea2ModelPatcher/en.md b/js/docs/Krea2ModelPatcher/en.md new file mode 100644 index 0000000..ec85610 --- /dev/null +++ b/js/docs/Krea2ModelPatcher/en.md @@ -0,0 +1,34 @@ +# Krea2 Model Patcher + +Patches the **Krea2** native text-to-image model with a trained image-adapter so it can accept reference-image conditioning alongside a text prompt. + +## Background + +Krea2 is a diffusion model released by Krea AI. This adapter was trained on semi-synthetic data — organic target images paired with synthetic reference images and prompts — to teach the model to follow reference imagery while still obeying text instructions. The adapter injects reference information via **RoPE encoding**, enabling the model to attend to up to 4 reference images at once. + +## Inputs + +| Name | Type | Description | +| --- | --- | --- | +| `model` | MODEL | The loaded Krea2 base model. | +| `patch_path` | STRING | Path to the adapter safetensors file (e.g. `krea2_image_adapter_v1_000001500.safetensors`). Relative paths resolve from the ComfyUI root. | + +## Outputs + +| Name | Type | Description | +| --- | --- | --- | +| `model` | MODEL | The patched Krea2 model, ready to consume reference conditioning. | + +## Usage + +1. Load the Krea2 checkpoint with a **Load Diffusion Model** node. +2. Connect its `MODEL` output to this node. +3. Set `patch_path` to the location of the adapter weights. +4. Wire the patched `MODEL` into a sampler (e.g. `KSampler`). +5. Pair with **Krea2 Reference Encoder** to supply the actual reference images. + +## Notes + +- This node must be used with the Krea2 base model. Using it with an unrelated checkpoint will produce undefined behaviour. +- The adapter is applied at full strength (`1.0`). Strength control may be added in a future version. +- Without a **Krea2 Reference Encoder** connected downstream, the patched model behaves identically to the unpatched base model. diff --git a/js/docs/Krea2ReferenceEncoder.md b/js/docs/Krea2ReferenceEncoder.md new file mode 100644 index 0000000..73c7abc --- /dev/null +++ b/js/docs/Krea2ReferenceEncoder.md @@ -0,0 +1,68 @@ +# Krea2 Reference Encoder + +Prepares the reference latents for a Krea2-patched model and appends them to your existing text conditioning. The adapter uses RoPE (Rotary Position Embedding) to attend to the references during diffusion without overwriting your text prompt. + +## How the slots work + +The adapter was trained with a fixed slot convention: + +| Slot | Input | Resolution | Role | +| --- | --- | --- | --- | +| 1 | `base_latent` | **Must match output** | The image being edited, inpainted, or outpainted | +| 2 | `reference_latent_2` | Any | Style / content reference (optional) | +| 3 | `reference_latent_3` | Any | Style / content reference (optional) | +| 4 | `reference_latent_4` | Any | Style / content reference (optional) | + +Slot 1 is always the **base image** — the canvas the model will modify. Slots 2–4 are optional visual references that guide style or content. + +## Inputs + +| Name | Type | Required | Description | +| --- | --- | --- | --- | +| `positive` | CONDITIONING | ✅ | Positive text conditioning from a CLIP Text Encode node. The reference payload is **appended** — your prompt is never overwritten. | +| `negative` | CONDITIONING | ✅ | Negative text conditioning. Receives the same reference payload as `positive`. | +| `base_latent` | LATENT | ✅ | The image to edit / inpaint, VAE-encoded. **Must be the same resolution as the generation output.** | +| `strength` | FLOAT | ✅ | Scales all reference latents before storing them. Range: `0.0 – 2.0`, default `1.0`. | +| `reference_latent_2` | LATENT | ➖ | Optional style / content reference, VAE-encoded. | +| `reference_latent_3` | LATENT | ➖ | Optional style / content reference, VAE-encoded. | +| `reference_latent_4` | LATENT | ➖ | Optional style / content reference, VAE-encoded. | + +## Outputs + +| Name | Type | Description | +| --- | --- | --- | +| `positive` | CONDITIONING | Positive conditioning with all reference latents embedded under `reference_latents`. | +| `negative` | CONDITIONING | Negative conditioning with the same reference payload. | + +## Usage + +### Instruct-style editing + +Edit an existing image by describing only the change you want. + +1. **Load Image** → **VAE Encode** → `base_latent` +2. Write a prompt describing the change (e.g. *"replace the sky with a stormy sunset"*) +3. Connect `positive` / `negative` to a **KSampler** whose model comes from **Krea2 Model Patcher** + +### Inpainting / outpainting + +Fill a masked or extended region while keeping the rest of the image intact. + +1. Prepare the source image with the unknown region masked (neutral grey or noise) +2. **VAE Encode** the masked image → `base_latent` +3. Prompt for the content that should fill the masked area + +### Reference-guided generation + +Generate a new image that draws style or content from reference images. + +1. Provide a blank or noise latent (same size as the desired output) → `base_latent` +2. Load style / content reference images → **VAE Encode** → `reference_latent_2` (and optionally `_3`, `_4`) +3. Write a prompt describing the desired output + +## Notes + +- `base_latent` **must match the output resolution** — this is a hard constraint from the training data. Mismatched resolutions will produce degraded results. +- `strength` above `1.0` amplifies the reference signal and can cause the model to over-fit to the references at the expense of prompt adherence. +- Chaining two **Krea2 Reference Encoder** nodes is supported — the second call accumulates its references onto the first via `torch.cat`. +- The downstream sampler **must** use a model patched with **Krea2 Model Patcher**. An unpatched model will silently ignore the `reference_latents` key in the conditioning metadata. diff --git a/js/docs/Krea2ReferenceEncoder/en.md b/js/docs/Krea2ReferenceEncoder/en.md new file mode 100644 index 0000000..73c7abc --- /dev/null +++ b/js/docs/Krea2ReferenceEncoder/en.md @@ -0,0 +1,68 @@ +# Krea2 Reference Encoder + +Prepares the reference latents for a Krea2-patched model and appends them to your existing text conditioning. The adapter uses RoPE (Rotary Position Embedding) to attend to the references during diffusion without overwriting your text prompt. + +## How the slots work + +The adapter was trained with a fixed slot convention: + +| Slot | Input | Resolution | Role | +| --- | --- | --- | --- | +| 1 | `base_latent` | **Must match output** | The image being edited, inpainted, or outpainted | +| 2 | `reference_latent_2` | Any | Style / content reference (optional) | +| 3 | `reference_latent_3` | Any | Style / content reference (optional) | +| 4 | `reference_latent_4` | Any | Style / content reference (optional) | + +Slot 1 is always the **base image** — the canvas the model will modify. Slots 2–4 are optional visual references that guide style or content. + +## Inputs + +| Name | Type | Required | Description | +| --- | --- | --- | --- | +| `positive` | CONDITIONING | ✅ | Positive text conditioning from a CLIP Text Encode node. The reference payload is **appended** — your prompt is never overwritten. | +| `negative` | CONDITIONING | ✅ | Negative text conditioning. Receives the same reference payload as `positive`. | +| `base_latent` | LATENT | ✅ | The image to edit / inpaint, VAE-encoded. **Must be the same resolution as the generation output.** | +| `strength` | FLOAT | ✅ | Scales all reference latents before storing them. Range: `0.0 – 2.0`, default `1.0`. | +| `reference_latent_2` | LATENT | ➖ | Optional style / content reference, VAE-encoded. | +| `reference_latent_3` | LATENT | ➖ | Optional style / content reference, VAE-encoded. | +| `reference_latent_4` | LATENT | ➖ | Optional style / content reference, VAE-encoded. | + +## Outputs + +| Name | Type | Description | +| --- | --- | --- | +| `positive` | CONDITIONING | Positive conditioning with all reference latents embedded under `reference_latents`. | +| `negative` | CONDITIONING | Negative conditioning with the same reference payload. | + +## Usage + +### Instruct-style editing + +Edit an existing image by describing only the change you want. + +1. **Load Image** → **VAE Encode** → `base_latent` +2. Write a prompt describing the change (e.g. *"replace the sky with a stormy sunset"*) +3. Connect `positive` / `negative` to a **KSampler** whose model comes from **Krea2 Model Patcher** + +### Inpainting / outpainting + +Fill a masked or extended region while keeping the rest of the image intact. + +1. Prepare the source image with the unknown region masked (neutral grey or noise) +2. **VAE Encode** the masked image → `base_latent` +3. Prompt for the content that should fill the masked area + +### Reference-guided generation + +Generate a new image that draws style or content from reference images. + +1. Provide a blank or noise latent (same size as the desired output) → `base_latent` +2. Load style / content reference images → **VAE Encode** → `reference_latent_2` (and optionally `_3`, `_4`) +3. Write a prompt describing the desired output + +## Notes + +- `base_latent` **must match the output resolution** — this is a hard constraint from the training data. Mismatched resolutions will produce degraded results. +- `strength` above `1.0` amplifies the reference signal and can cause the model to over-fit to the references at the expense of prompt adherence. +- Chaining two **Krea2 Reference Encoder** nodes is supported — the second call accumulates its references onto the first via `torch.cat`. +- The downstream sampler **must** use a model patched with **Krea2 Model Patcher**. An unpatched model will silently ignore the `reference_latents` key in the conditioning metadata. diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..bc75790 --- /dev/null +++ b/nodes.py @@ -0,0 +1,506 @@ +""" +Krea2 MultiRef custom nodes for ComfyUI. +Adapted from Ostris's implementation, with RoPE coordinate fix and +training-specific adjustments. + +Training-specific details: + - Reference tag format: <|vision_start|><|image_pad|><|vision_end|> + - VLM max pixels: 512*512 = 262144 (patched from 384*384 during training) + - VAE ref latent max pixels: 1024*1024, snapped to /16 +""" + +import math + +import torch +from einops import rearrange + +import comfy.conds +import comfy.ldm.common_dit +import comfy.utils +import node_helpers +from comfy.ldm.flux.layers import timestep_embedding + +# Training-patched values +VLM_MAX_PIXELS = 512 * 512 # patched from 384*384 +REF_LATENT_MAX_PIXELS = 1024 * 1024 +REF_SNAP = 16 # VAE f8 * patch 2 + + +# --------------------------------------------------------------------------- +# Image resizing helpers +# --------------------------------------------------------------------------- + +def _fit_area(samples, max_pixels, snap=1): + """ + Downscale (never upscale) a (B, C, H, W) tensor to fit max_pixels, + keeping aspect ratio, snapping dims to snap. + """ + h, w = samples.shape[2], samples.shape[3] + scale = min(1.0, math.sqrt(max_pixels / (w * h))) + nw = max(round(w * scale / snap) * snap, snap) + nh = max(round(h * scale / snap) * snap, snap) + if (nh, nw) == (h, w): + return samples + return comfy.utils.common_upscale(samples, nw, nh, "area", "disabled") + + +# --------------------------------------------------------------------------- +# Text + vision encode node +# --------------------------------------------------------------------------- + +class Krea2ReferenceTextEncode: + """ + Encodes prompt + reference images through the Krea2 Qwen3-VL text encoder. + + Reference images are injected into the prompt using the tag format + that was used during training: + <|vision_start|><|image_pad|><|vision_end|> + + The VLM sees a downscaled version (max 512*512 px) for semantic understanding. + When a VAE is connected, full-resolution reference latents (max 1MP, snapped + to /16) are also attached for the model patch to use. + + This node replaces the standard CLIPTextEncode when using Krea2 edit LoRAs. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip": ("CLIP",), + "prompt": ("STRING", { + "multiline": True, + "dynamicPrompts": True, + }), + }, + "optional": { + "vae": ("VAE",), + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "image3": ("IMAGE",), + "image4": ("IMAGE",), + }, + } + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "encode" + CATEGORY = "krea2/editing" + DESCRIPTION = ( + "Encode prompt + reference images for a Krea2 edit LoRA. " + "Images are fed to the Qwen3-VL text encoder via training-matching " + "reference tags. Connect a VAE to also attach reference latents for " + "the Krea2ModelPatcher node." + ) + + def encode(self, clip, prompt, + vae=None, image1=None, image2=None, image3=None, image4=None): + images_vl = [] + ref_latents = [] + image_prompt = "" + + for image in (image1, image2, image3, image4): + if image is None: + continue + + n = len(images_vl) + 1 + + # Downscale for VLM (semantic understanding) + samples_hw = image.movedim(-1, 1) # (B, C, H, W) + vl_img = _fit_area(samples_hw, VLM_MAX_PIXELS).movedim(1, -1) + images_vl.append(vl_img) + + # Build reference tag matching training format + image_prompt += ( + f"" + f"<|vision_start|><|image_pad|><|vision_end|>" + f"" + ) + + # Encode reference latent for model patcher + if vae is not None: + s = _fit_area(samples_hw, REF_LATENT_MAX_PIXELS, snap=REF_SNAP) + ref_latents.append( + vae.encode(s.movedim(1, -1)[:, :, :, :3]) + ) + + # Krea2 uses its own template (KREA2_TEMPLATE) — try to import it, + # fall back to passing None which uses the model's default. + try: + from comfy.text_encoders.krea2 import KREA2_TEMPLATE + llama_template = KREA2_TEMPLATE + except ImportError: + llama_template = None + + tokens = clip.tokenize( + image_prompt + prompt, + images=images_vl if images_vl else None, + llama_template=llama_template, + ) + conditioning = clip.encode_from_tokens_scheduled(tokens) + + if ref_latents: + conditioning = node_helpers.conditioning_set_values( + conditioning, + {"reference_latents": ref_latents}, + append=True, + ) + + return (conditioning,) + + +# --------------------------------------------------------------------------- +# Conditioning helper (for the simple latent-only encoder) +# --------------------------------------------------------------------------- + +def _append_ref_to_cond(conditioning, ref_tensor, strength): + out = [] + scaled = ref_tensor * strength + for embedding, meta in conditioning: + meta = dict(meta) + refs = list(meta.get("reference_latents") or []) + refs.append(scaled) + meta["reference_latents"] = refs + out.append((embedding, meta)) + return out + + +# --------------------------------------------------------------------------- +# Reference packing — Kontext-style RoPE layout +# --------------------------------------------------------------------------- + +def _pack_refs(dit, ref_latents, bs, device, dtype, target_w): + """ + Patchify reference latents into tokens + RoPE positions. + + RoPE layout (same axis-0=0 as target, offset in x): + target: axis-0=0, y in [0, h_), x in [0, target_w) + ref 0: axis-0=0, y in [0, rh0), x in [target_w, target_w+rw0) + ref 1: axis-0=0, y in [0, rh1), x in [target_w+rw0, ...) + """ + patch = dit.patch + ref_tokens = [] + ref_pos = [] + x_offset = target_w + + for i, ref in enumerate(ref_latents): + if ref.ndim == 5: + rb, rc, rt, rh5, rw5 = ref.shape + ref = ref.reshape(rb * rt, rc, rh5, rw5) + + ref = comfy.ldm.common_dit.pad_to_patch_size( + ref.to(device, dtype), (patch, patch) + ) + ref = comfy.utils.repeat_to_batch_size(ref, bs) + + rh = ref.shape[-2] // patch + rw = ref.shape[-1] // patch + + ref_tokens.append( + rearrange(ref, "b c (h ph) (w pw) -> b (h w) (c ph pw)", + ph=patch, pw=patch) + ) + + rid = torch.zeros(rh, rw, 3, device=device, dtype=torch.float32) + rid[..., 0] = 0.0 + rid[..., 1] = torch.arange(rh, device=device, + dtype=torch.float32)[:, None] + rid[..., 2] = ( + torch.arange(rw, device=device, dtype=torch.float32)[None, :] + + x_offset + ) + ref_pos.append(rid.reshape(1, rh * rw, 3).repeat(bs, 1, 1)) + x_offset += rw + + return torch.cat(ref_tokens, dim=1), torch.cat(ref_pos, dim=1) + + +# --------------------------------------------------------------------------- +# Per-block forward with split modulation +# --------------------------------------------------------------------------- + +def _block_ref_forward(block, x, vec, refvec, split, freqs, transformer_options): + """ + SingleStreamBlock forward with per-span AdaLN: + [:split] text + noisy image -> vec (real timestep) + [split:] clean references -> refvec (timestep=0) + """ + m = block.mod(vec) + r = block.mod(refvec) + + def mod(h, scale, shift): + return torch.cat( + ( + (1 + m[scale]) * h[:, :split] + m[shift], + (1 + r[scale]) * h[:, split:] + r[shift], + ), + dim=1, + ) + + def gate(h, g): + return torch.cat( + (m[g] * h[:, :split], r[g] * h[:, split:]), + dim=1, + ) + + x = x + gate( + block.attn( + mod(block.prenorm(x), 0, 1), + freqs, None, + transformer_options=transformer_options, + ), + 2, + ) + x = x + gate(block.mlp(mod(block.postnorm(x), 3, 4)), 5) + return x + + +# --------------------------------------------------------------------------- +# Full forward with reference latents +# --------------------------------------------------------------------------- + +def _forward_with_refs(dit, x, timesteps, context, ref_latents, + transformer_options): + temporal = x.ndim == 5 + if temporal: + b5, c5, t5, h5, w5 = x.shape + x = x.reshape(b5 * t5, c5, h5, w5) + + bs, c, H_orig, W_orig = x.shape + patch = dit.patch + device = x.device + + x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch)) + H, W = x.shape[-2], x.shape[-1] + h_, w_ = H // patch, W // patch + + context = dit._unpack_context(context) + + img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", + ph=patch, pw=patch) + + reftok, refpos = _pack_refs(dit, ref_latents, bs, device, x.dtype, + target_w=w_) + reflen = reftok.shape[1] + + img = dit.first(torch.cat((img, reftok), dim=1)) + + t = dit.tmlp(timestep_embedding(timesteps, dit.tdim) + .unsqueeze(1).to(img.dtype)) + tvec = dit.tproj(t) + t0 = dit.tmlp(timestep_embedding(torch.zeros_like(timesteps), dit.tdim) + .unsqueeze(1).to(img.dtype)) + tvec0 = dit.tproj(t0) + + context = dit.txtfusion(context, mask=None, + transformer_options=transformer_options) + context = dit.txtmlp(context) + + txtlen = context.shape[1] + imglen = img.shape[1] + + combined = torch.cat((context, img), dim=1) + split = txtlen + imglen - reflen + + txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32) + imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32) + imgids[..., 1] = torch.arange(h_, device=device, + dtype=torch.float32)[:, None] + imgids[..., 2] = torch.arange(w_, device=device, + dtype=torch.float32)[None, :] + imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1) + + pos = torch.cat((txtpos, imgpos, refpos), dim=1) + freqs = dit.pe_embedder(pos) + + for block in dit.blocks: + combined = _block_ref_forward( + block, combined, tvec, tvec0, split, freqs, transformer_options + ) + + final = dit.last(combined, t) + out = final[:, txtlen:split, :] + out = rearrange( + out, + "b (h w) (c ph pw) -> b c (h ph) (w pw)", + h=h_, w=w_, ph=patch, pw=patch, c=dit.channels, + ) + out = out[:, :, :H_orig, :W_orig] + + if temporal: + out = out.reshape(b5, t5, dit.channels, H_orig, W_orig).movedim(1, 2) + return out + + +# --------------------------------------------------------------------------- +# Model patcher +# --------------------------------------------------------------------------- + +class Krea2ModelPatcher: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"model": ("MODEL",)}} + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + CATEGORY = "krea2/editing" + FUNCTION = "patch" + + def patch(self, model): + patched = model.clone() + base_model = patched.model + dit = patched.get_model_object("diffusion_model") + + orig_forward = dit.forward + orig_extra_conds = base_model.extra_conds + orig_extra_conds_shapes = base_model.extra_conds_shapes + + def extra_conds(**kwargs): + out = orig_extra_conds(**kwargs) + refs = kwargs.get("reference_latents") + if refs is not None: + out["ref_latents"] = comfy.conds.CONDList( + [base_model.process_latent_in(lat) for lat in refs] + ) + return out + + def extra_conds_shapes(**kwargs): + out = orig_extra_conds_shapes(**kwargs) + refs = kwargs.get("reference_latents") + if refs is not None: + out["ref_latents"] = [ + 1, 16, + sum(math.prod(r.size()) for r in refs) // 16, + ] + return out + + def forward( + x, + timesteps, + context, + attention_mask=None, + transformer_options={}, + ref_latents=None, + **kwargs, + ): + if ref_latents is None or len(ref_latents) == 0: + return orig_forward( + x, timesteps, context, + attention_mask=attention_mask, + transformer_options=transformer_options, + **kwargs, + ) + return _forward_with_refs( + dit, x, timesteps, context, ref_latents, transformer_options, + ) + + patched.add_object_patch("extra_conds", extra_conds) + patched.add_object_patch("extra_conds_shapes", extra_conds_shapes) + patched.add_object_patch("diffusion_model.forward", forward) + return (patched,) + + +# --------------------------------------------------------------------------- +# Simple latent-only reference encoder (no VLM — use when you already have +# a text conditioning from a normal CLIPTextEncode and just want to attach +# latent refs without re-encoding text) +# --------------------------------------------------------------------------- + +class Krea2ReferenceEncoder: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive": ("CONDITIONING",), + "base_latent": ("LATENT", { + "tooltip": "VAE-encoded image to edit. Reference slot 1." + }), + "strength": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, + }), + }, + "optional": { + "reference_latent_2": ("LATENT",), + "reference_latent_3": ("LATENT",), + "reference_latent_4": ("LATENT",), + }, + } + + RETURN_TYPES = ("CONDITIONING",) + CATEGORY = "krea2/editing" + FUNCTION = "encode" + + def encode( + self, + positive, + base_latent, + strength = 1.0, + reference_latent_2 = None, + reference_latent_3 = None, + reference_latent_4 = None, + ): + slots = [base_latent, reference_latent_2, + reference_latent_3, reference_latent_4] + tensors = [s["samples"] for s in slots if s is not None] + + new_positive = positive + for ref in tensors: + new_positive = _append_ref_to_cond(new_positive, ref, strength) + return (new_positive,) + + +# --------------------------------------------------------------------------- +# Diagnostic +# --------------------------------------------------------------------------- + +class Krea2DiagnosticNode: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"model": ("MODEL",)}} + + RETURN_TYPES = () + OUTPUT_NODE = True + CATEGORY = "krea2/editing" + FUNCTION = "run" + + def run(self, model): + import inspect + dit = model.get_model_object("diffusion_model") + print("\n=== DIT TYPE ===") + print(type(dit).__name__, type(dit).__module__) + print("\n=== DIT NAMED CHILDREN ===") + for name, mod in dit.named_children(): + print(f" {name}: {type(mod).__name__}") + print("\n=== DIT NON-MODULE ATTRS ===") + for k, v in vars(dit).items(): + if not k.startswith("_") and not isinstance(v, torch.nn.Module): + print(f" {k}: {type(v).__name__} = {repr(v)[:120]}") + print("\n=== DIT FORWARD SIGNATURE ===") + print(inspect.signature(dit.forward)) + print("\n=== FIRST BLOCK ===") + if hasattr(dit, "blocks"): + block = dit.blocks[0] + print(f" {type(block).__name__}") + for n, m in block.named_children(): + print(f" {n}: {type(m).__name__}") + print(f" sig: {inspect.signature(block.forward)}") + print(f"\n has _unpack_context: {hasattr(dit, '_unpack_context')}") + return () + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + +NODE_CLASS_MAPPINGS = { + "Krea2ModelPatcher": Krea2ModelPatcher, + "Krea2ReferenceEncoder": Krea2ReferenceEncoder, + "Krea2ReferenceTextEncode": Krea2ReferenceTextEncode, + "Krea2DiagnosticNode": Krea2DiagnosticNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Krea2ModelPatcher": "Krea2 Model Patcher", + "Krea2ReferenceEncoder": "Krea2 Reference Encoder (latent only)", + "Krea2ReferenceTextEncode": "Krea2 Reference Text Encode", + "Krea2DiagnosticNode": "Krea2 Diagnostic", +} \ No newline at end of file