nodes/molbal-ComfyUI-Krea2-MultiRef
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

placeholder

Browse Source
This commit is contained in:
molbal
2026-07-14 10:25:11 +02:00
commit fb755f7675
8 changed files with 780 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+36
View File
@@ -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/
README.md
+31
View File
@@ -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.
__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
js/docs/Krea2ModelPatcher.md
+34
View File
@@ -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.
js/docs/Krea2ModelPatcher/en.md
+34
View File
@@ -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.
js/docs/Krea2ReferenceEncoder.md
+68
View File
@@ -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.
js/docs/Krea2ReferenceEncoder/en.md
+68
View File
@@ -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.
nodes.py
+506
View File
@@ -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: <reference_N><|vision_start|><|image_pad|><|vision_end|></reference_N>
- 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:
<reference_N><|vision_start|><|image_pad|><|vision_end|></reference_N>
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"<reference_{n}>"
f"<|vision_start|><|image_pad|><|vision_end|>"
f"</reference_{n}>"
)
# 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",
}