From ec50d3acfedae3d8bca7808601f21c4f473a14d3 Mon Sep 17 00:00:00 2001 From: ausboss <102309304+ausboss@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:17:19 -0500 Subject: [PATCH] Promote the Krea 2 outpaint pair from the lab Krea 2 accepts reference latents, but they arrive with no position, so the model treats them as a loose style hint: it borrows the look and reinvents the content. That is the whole reason an outpaint drifts instead of continuing the photo. Krea 2 Encode attaches the reference latents and emits the negative from the same node, so a turbo graph at CFG 1.0 stops carrying a second text encode that does nothing. Krea 2 Outpaint Model Patch registers those reference tokens into the target grid at the rectangle the stitcher records, on frame axis 1. Load Image + Pad gained a `reference` output - the unpadded source, fitted to a 384px long edge and a multiple of 16, because the VAE downsamples by 8 and the DiT patchifies by 2. It is APPENDED after `stitcher`, never inserted: a workflow stores links by output slot index, so any other position silently rewires every saved graph. The stitcher now records where the source sits, as source_bbox in pixels and bbox_normalized in 0..1. Padding knows the rectangle exactly, so it rides with the canvas instead of a parallel wire that can be left unplugged. apply_stitch reads neither key, so an older stitcher stitches identically. Ported from the lab rather than copied: the lab's _krea2_helpers had its own pad_image, feather_pad_mask, plan_pad_canvas, parse_fill_color and build_canvas_stitcher. Duplicating the padding engine here would let the outpaint path and the inpaint path disagree about what a stitcher is, so the public module keeps only what is specific to Krea 2 and uses the existing helpers for the rest. Two departures from the lab version. The model patch declared a kv_cache input and never read it - the cache was always on - so it is wired up now. And it imports comfy's flux internals when you run it rather than at startup: a core release that moves one then surfaces as an error on that node, instead of deleting it from the menu and leaving saved workflows reporting it missing. --- CHANGELOG.md | 27 ++ README.md | 10 +- __init__.py | 2 + js/docs/AUSBOSS_NODES_Krea2Encode.md | 43 ++++ .../AUSBOSS_NODES_Krea2OutpaintModelPatch.md | 40 +++ nodes/_inpaint_crop_helpers.py | 27 +- nodes/_krea2_core.py | 240 ++++++++++++++++++ nodes/_krea2_helpers.py | 118 +++++++++ nodes/node_krea2_encode.py | 168 ++++++++++++ nodes/node_krea2_model_patch.py | 179 +++++++++++++ nodes/node_load_image_pad.py | 35 ++- scripts/validate_nodes.py | 2 + tests/test_inpaint_crop_helpers.py | 43 ++++ tests/test_krea2_helpers.py | 194 ++++++++++++++ tests/test_pad_helpers.py | 71 +++++- 15 files changed, 1186 insertions(+), 13 deletions(-) create mode 100644 js/docs/AUSBOSS_NODES_Krea2Encode.md create mode 100644 js/docs/AUSBOSS_NODES_Krea2OutpaintModelPatch.md create mode 100644 nodes/_krea2_core.py create mode 100644 nodes/_krea2_helpers.py create mode 100644 nodes/node_krea2_encode.py create mode 100644 nodes/node_krea2_model_patch.py create mode 100644 tests/test_krea2_helpers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index db2358d..67ab618 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,33 @@ All notable changes to ComfyUI-AusBoss are documented here. ## Unreleased +- Added **Krea 2 Encode 🆎** and **Krea 2 Outpaint Model Patch 🆎**, promoted + from the lab. Together they make Krea 2 outpaint the source instead of + painting something next to it: the encode attaches reference latents to the + positive conditioning (and emits the negative from the same node, so a turbo + graph at CFG 1.0 stops carrying a second text encode that does nothing), and + the patch registers those reference tokens into the target grid at the + rectangle the stitcher records. Reference latents with no position are only + a style hint — the model borrows the look and reinvents the content, which + is the failure this pair fixes. + + The patch reaches into comfy's flux attention layers, so it imports them when + you run it rather than at startup: a core release that moves one surfaces as + an error on that node instead of deleting it from the menu and leaving saved + workflows reporting it missing. + +- Load Image + Pad 🆎 gained a `reference` output — the unpadded source, fitted + to a 384px long edge and a multiple of 16, ready for reference conditioning. + It is **appended** after `stitcher`, not inserted, because a workflow stores + links by output slot index and anything else would silently rewire every + saved graph. + +- The stitcher now records where the source sits on the canvas, as + `source_bbox` in pixels and `bbox_normalized` in 0..1. Padding knows the + rectangle exactly, so it rides along with the canvas rather than on a + parallel wire that can be left unplugged. Stitching itself never reads + either key, so an older stitcher still stitches identically. + - Removed **Drop Shadow 🆎** (`AUSBOSS_NODES_DropShadow`). The result never looked like a real cast shadow — a mask offset, grown and blurred has no contact darkening and no perspective, so it read as a sticker halo rather diff --git a/README.md b/README.md index c906816..69ac170 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,15 @@ Harmonizes an image against a reference by transferring per-channel LAB mean and ### Load Image + Pad 🆎 -A Load Image that opens straight into an outpaint canvas: drag any edge of the final rect drawn on the node to grow that side's padding — the whole edge is the handle, per-side pixel counts ride the bands, and the badge always shows the true output size. The mask covers exactly the padding with an optional **feather** ramped inward across the seam, the canvas rounds up to a clean multiple, and a **megapixel target** rescales the source *before* padding so the mask seam stays crisp at sampler-friendly sizes. Outputs the canvas, the mask, the final width/height as INTs, and a `stitcher` for Stitch Inpaint 🆎 that restores the original pixels bit-identically after sampling. +A Load Image that opens straight into an outpaint canvas: drag any edge of the final rect drawn on the node to grow that side's padding — the whole edge is the handle, per-side pixel counts ride the bands, and the badge always shows the true output size. The mask covers exactly the padding with an optional **feather** ramped inward across the seam, the canvas rounds up to a clean multiple, and a **megapixel target** rescales the source *before* padding so the mask seam stays crisp at sampler-friendly sizes. Outputs the canvas, the mask, the final width/height as INTs, a `stitcher` for Stitch Inpaint 🆎 that restores the original pixels bit-identically after sampling, and a `reference` — the unpadded source fitted to a small multiple of 16 for reference conditioning. + +### Krea 2 Encode 🆎 + +Encodes a Krea 2 prompt pair and attaches reference latents in one node. Wire a VAE and a reference image and the reference is encoded onto the positive conditioning; leave them unwired and it is a plain two-prompt encoder. The negative comes out of the same node, so a turbo graph at CFG 1.0 is not carrying a second text encode that does nothing. References are fitted to a 384px long edge and a multiple of 16 — the VAE downsamples by 8 and the DiT patchifies by 2, so an odd edge lands on a partial patch. + +### Krea 2 Outpaint Model Patch 🆎 + +Tells Krea 2 **where** the reference sits. Reference latents normally arrive with no position, so the model treats them as a loose style hint and reinvents the content; this registers them into the target grid at the rectangle Load Image + Pad 🆎 records on the `stitcher`, which is what makes an outpaint continue the source instead of painting something adjacent to it. Place it after any LoRA loader and before the sampler. It patches comfy's flux attention internals, so it imports them when you run it rather than at startup — a core change that moves them surfaces as an error on this node instead of quietly dropping it from the menu. ### Frame Interpolate 🆎 diff --git a/__init__.py b/__init__.py index 7347f2e..6e34b55 100644 --- a/__init__.py +++ b/__init__.py @@ -25,6 +25,8 @@ NODE_MODULES = [ "node_image_crop_rotate_pad", "node_image_size", "node_inpaint_crop_stitch", + "node_krea2_encode", + "node_krea2_model_patch", "node_lama_inpaint", "node_lmstudio_chat", "node_load_image_pad", diff --git a/js/docs/AUSBOSS_NODES_Krea2Encode.md b/js/docs/AUSBOSS_NODES_Krea2Encode.md new file mode 100644 index 0000000..3ed4ee4 --- /dev/null +++ b/js/docs/AUSBOSS_NODES_Krea2Encode.md @@ -0,0 +1,43 @@ +# Krea 2 Encode + +Encodes a Krea 2 prompt pair and attaches reference latents, in one node. + +Wire a VAE and a reference image and the reference is encoded and appended to +the positive conditioning. Leave them unwired and this is a plain two-prompt +encoder. The negative comes out of the same node, so a turbo graph running at +CFG 1.0 still has something to plug in without a second text encode sitting +there doing nothing. + +## Controls + +- **clip**: The Krea 2 text encoder. +- **prompt**: Describe the whole finished canvas, not just the new area — the + model generates all of it and matches the reference where it must. +- **negative_prompt**: Ignored at CFG 1.0, which is where turbo runs. Empty is + the normal case there. +- **vae**: Needed to turn the reference images into latents. Without it the + references are skipped and only the prompts are encoded. +- **reference**: The `reference` output of Load Image + Pad 🆎 — the unpadded + source. Any image works; it is fitted to a multiple of 16 first. +- **extra_image**: A second reference, e.g. a style or character plate. +- **vlm_reference**: Also show the references to the vision tower, so the text + encoder describes them. Off for outpaint: the latents already carry the + picture, and the description tends to pull the result toward a paraphrase of + the source. + +## Outputs + +- **positive**: Prompt conditioning with the reference latents attached. +- **negative**: Negative prompt conditioning. + +## Notes and limitations + +- **References are downscaled to a 384px long edge and snapped to a multiple + of 16.** The VAE downsamples by 8 and the DiT patchifies by 2, so an edge off + a multiple of 16 lands on a partial patch and the token grid stops lining up + with the canvas grid. The cap is deliberate — a reference is there to say + "this is the picture", and the canvas latent carries the detail. +- Attaching reference latents alone does not tell the model *where* they go. + Pair this with Krea 2 Outpaint Model Patch 🆎 for outpainting. +- On a core without the Krea 2 prompt template, `vlm_reference` still works — + it falls back to tokenizing without the template. diff --git a/js/docs/AUSBOSS_NODES_Krea2OutpaintModelPatch.md b/js/docs/AUSBOSS_NODES_Krea2OutpaintModelPatch.md new file mode 100644 index 0000000..32aec4b --- /dev/null +++ b/js/docs/AUSBOSS_NODES_Krea2OutpaintModelPatch.md @@ -0,0 +1,40 @@ +# Krea 2 Outpaint Model Patch + +Tells Krea 2 **where** a reference image sits on the canvas. + +Reference latents normally arrive as extra tokens with no position, so the +model treats them as a loose style hint — it borrows the look and reinvents +the content. This patch registers those tokens into the target grid at the +rectangle the stitcher reports, which is what makes an outpaint continue the +source rather than paint something adjacent to it. + +Place it **after** any LoRA loader and **before** the sampler. + +## Controls + +- **model**: A Krea 2 model. Patch last, so a LoRA loaded afterwards does not + replace the patched forward pass. +- **stitcher**: From Load Image + Pad 🆎. Supplies the source rectangle. A + stitcher that carries no rectangle — one from Crop For Inpaint, or an older + saved graph — places the reference over the full frame, which is what an + unpatched model already assumes. +- **kv_cache**: Compute the reference's keys and values once per run instead + of once per step. The reference does not change while sampling, so this is + free speed. Turn it off only to rule the cache out when debugging. + +## Output + +- **MODEL**: The model with reference tokens registered into the canvas grid. + +## Notes and limitations + +- **Nothing happens without reference latents.** If the conditioning carries + none, the patched forward pass calls straight through to the original. Wire + Krea 2 Encode 🆎 with a VAE and a reference image. +- **It patches comfy internals.** The placement reaches into the flux + attention layers, so a ComfyUI release that moves them can break this node + specifically. It imports those internals when you run it rather than at + startup, so a break surfaces as an error on this node instead of the node + disappearing from the menu. +- Reference tokens cost attention. The reference is fitted to a short edge + before encoding for that reason — see Krea 2 Encode 🆎. diff --git a/nodes/_inpaint_crop_helpers.py b/nodes/_inpaint_crop_helpers.py index 068f952..ee4e004 100644 --- a/nodes/_inpaint_crop_helpers.py +++ b/nodes/_inpaint_crop_helpers.py @@ -437,7 +437,11 @@ def build_crop( return cropped, sampling, stitcher -def build_canvas_stitcher(canvas: torch.Tensor, blend: torch.Tensor) -> dict: +def build_canvas_stitcher( + canvas: torch.Tensor, + blend: torch.Tensor, + bbox: tuple[int, int, int, int] | None = None, +) -> dict: """A stitcher that pastes a full-frame result back over ``canvas``. The crop/stitch pair sends a *region* to the sampler. Padding sends the @@ -449,11 +453,20 @@ def build_canvas_stitcher(canvas: torch.Tensor, blend: torch.Tensor) -> dict: Sharing one stitcher shape means Stitch Inpaint 🆎 accepts either producer; no second stitch node has to exist. + + ``bbox`` is where the *source* sits inside the canvas, as + ``(x0, y0, x1, y1)`` pixels. Padding knows it exactly, and a model that + places reference tokens on the canvas grid needs it, so it rides along + here rather than on a parallel wire that can be left unplugged. It is + stored twice - ``source_bbox`` in pixels and ``bbox_normalized`` in 0..1 - + because a consumer reading normalized coordinates should not have to know + the canvas size. Omitted when unknown; :func:`apply_stitch` never reads + either key, so an older stitcher still stitches. """ canvas = _as_image(canvas) blend = _as_mask(blend, canvas) height, width = canvas.shape[1], canvas.shape[2] - return { + stitcher = { "kind": STITCHER_KIND, "version": STITCHER_VERSION, "canvas": canvas, @@ -463,6 +476,16 @@ def build_canvas_stitcher(canvas: torch.Tensor, blend: torch.Tensor) -> dict: "scale": (1.0, 1.0), "algorithm": "bilinear", } + if bbox is not None: + x0, y0, x1, y1 = (int(value) for value in bbox) + stitcher["source_bbox"] = (x0, y0, x1, y1) + stitcher["bbox_normalized"] = [ + x0 / float(width), + y0 / float(height), + x1 / float(width), + y1 / float(height), + ] + return stitcher def apply_stitch( diff --git a/nodes/_krea2_core.py b/nodes/_krea2_core.py new file mode 100644 index 0000000..93acbce --- /dev/null +++ b/nodes/_krea2_core.py @@ -0,0 +1,240 @@ +"""Registered reference attention for Krea 2. + +Krea 2's single-stream DiT reads position from RoPE. A reference latent +appended as extra tokens therefore has to be given coordinates, or the model +has no idea where on the canvas it belongs. These helpers place the reference +grid at a normalized bbox on frame axis 1 (the target occupies axis 0), then +run the blocks once with the reference keys and values precomputed, so the +reference is attended to without being denoised. + +Everything here reaches into comfy's flux layers, so it is version-sensitive +by nature - see the Krea 2 section of the README for the versions it is +tested against. +""" + +from __future__ import annotations + +import math +import torch +import torch.nn.functional as F +from einops import rearrange + +import comfy.conds +import comfy.ldm.common_dit +import comfy.utils +from comfy.ldm.flux.layers import timestep_embedding +from comfy.ldm.flux.math import apply_rope +from comfy.ldm.modules.attention import optimized_attention_masked + + +def _target_grid(x: torch.Tensor, patch: int) -> tuple[int, int]: + hs, ws = x.shape[-2], x.shape[-1] + hp = ((hs + patch - 1) // patch) * patch + wp = ((ws + patch - 1) // patch) * patch + return hp // patch, wp // patch + + +def _pack_refs( + dit, + ref_latents, + bs: int, + device: torch.device, + dtype: torch.dtype, + target_hw: tuple[int, int] | None = None, + bbox_norm: list[float] | None = None, +): + patch = dit.patch + ref_tokens = [] + ref_pos = [] + 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, rw = ref.shape[-2] // patch, 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] = i + 1.0 + if bbox_norm is None: + rid[..., 1] = torch.arange(rh, device=device, dtype=torch.float32)[:, None] + rid[..., 2] = torch.arange(rw, device=device, dtype=torch.float32)[None, :] + else: + if target_hw is None: + raise ValueError("target_hw is required for a registered reference.") + th, tw = target_hw + x0, y0, x1, y1 = (float(v) for v in bbox_norm) + ys = y0 * th + (torch.arange(rh, device=device, dtype=torch.float32) + 0.5) * ( + (y1 - y0) * th / rh + ) - 0.5 + xs = x0 * tw + (torch.arange(rw, device=device, dtype=torch.float32) + 0.5) * ( + (x1 - x0) * tw / rw + ) - 0.5 + rid[..., 1] = ys[:, None] + rid[..., 2] = xs[None, :] + ref_pos.append(rid.reshape(1, rh * rw, 3).repeat(bs, 1, 1)) + return torch.cat(ref_tokens, dim=1), torch.cat(ref_pos, dim=1) + + +def _attn_kv( + attn, + x, + freqs, + mask=None, + kv_capture=None, + kv_cache=None, + transformer_options={}, +): + q, k, v, gate = attn.wq(x), attn.wk(x), attn.wv(x), attn.gate(x) + q = rearrange(q, "B L (H D) -> B H L D", H=attn.heads) + k = rearrange(k, "B L (H D) -> B H L D", H=attn.kvheads) + v = rearrange(v, "B L (H D) -> B H L D", H=attn.kvheads) + q, k = attn.qknorm(q, k) + if freqs is not None: + q, k = apply_rope(q, k, freqs) + if kv_capture is not None: + kv_capture.append((k, v)) + if kv_cache is not None: + k = torch.cat((k, kv_cache[0].to(k.dtype)), dim=2) + v = torch.cat((v, kv_cache[1].to(v.dtype)), dim=2) + if attn.kvheads != attn.heads: + rep = attn.heads // attn.kvheads + k = k.repeat_interleave(rep, dim=1) + v = v.repeat_interleave(rep, dim=1) + out = optimized_attention_masked( + q, + k, + v, + attn.heads, + mask=mask, + skip_reshape=True, + transformer_options=transformer_options, + ) + return attn.wo(out * F.sigmoid(gate)) + + +def _block_kv_forward( + block, + x, + vec, + freqs, + kv_capture=None, + kv_cache=None, + transformer_options={}, +): + prescale, preshift, pregate, postscale, postshift, postgate = block.mod(vec) + x = x + pregate * _attn_kv( + block.attn, + (1 + prescale) * block.prenorm(x) + preshift, + freqs, + kv_capture=kv_capture, + kv_cache=kv_cache, + transformer_options=transformer_options, + ) + x = x + postgate * block.mlp((1 + postscale) * block.postnorm(x) + postshift) + return x + + +def _precompute_ref_kv( + dit, x, timesteps, ref_latents, transformer_options, bbox_norm=None +): + temporal = x.ndim == 5 + bs = x.shape[0] * (x.shape[2] if temporal else 1) + th, tw = _target_grid(x, dit.patch) + reftok, refpos = _pack_refs( + dit, ref_latents, bs, x.device, x.dtype, target_hw=(th, tw), bbox_norm=bbox_norm + ) + h = dit.first(reftok) + t0 = dit.tmlp( + timestep_embedding(torch.zeros_like(timesteps), dit.tdim).unsqueeze(1).to(h.dtype) + ) + tvec0 = dit.tproj(t0) + freqs = dit.pe_embedder(refpos) + + ref_kv = [] + for block in dit.blocks: + cap = [] + h = _block_kv_forward( + block, + h, + tvec0, + freqs, + kv_capture=cap, + transformer_options=transformer_options, + ) + ref_kv.append(cap[0]) + return ref_kv + + +def _forward_with_cached_refs( + dit, x, timesteps, context, ref_kv, 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 + 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 + device = x.device + + context = dit._unpack_context(context) + img = dit.first( + rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) + ) + + t = dit.tmlp(timestep_embedding(timesteps, dit.tdim).unsqueeze(1).to(img.dtype)) + tvec = dit.tproj(t) + + context = dit.txtfusion(context, mask=None, transformer_options=transformer_options) + context = dit.txtmlp(context) + + txtlen, imglen = context.shape[1], img.shape[1] + combined = torch.cat((context, img), dim=1) + + 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), dim=1) + + freqs = dit.pe_embedder(pos) + + for block, kv in zip(dit.blocks, ref_kv): + combined = _block_kv_forward( + block, + combined, + tvec, + freqs, + kv_cache=kv, + transformer_options=transformer_options, + ) + + final = dit.last(combined, t) + out = final[:, txtlen : txtlen + imglen, :] + 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 + + +def _ref_fingerprint(ref_latents, bs: int, bbox_norm: list[float] | None): + key = [bs, tuple(round(v, 6) for v in bbox_norm) if bbox_norm is not None else None] + for r in ref_latents: + rf = r.float() + key.append((tuple(r.shape), float(rf.sum()), float(rf.square().sum()))) + return tuple(key) diff --git a/nodes/_krea2_helpers.py b/nodes/_krea2_helpers.py new file mode 100644 index 0000000..4aab3cd --- /dev/null +++ b/nodes/_krea2_helpers.py @@ -0,0 +1,118 @@ +"""Krea 2 reference-conditioning helpers. + +Only the parts that are specific to Krea 2 live here. Padding, canvas +planning, colour parsing and the stitcher itself already exist in +``_pad_helpers`` and ``_inpaint_crop_helpers``, and are used from there +rather than reimplemented - the outpaint path and the inpaint path must +agree on what a stitcher is, or Stitch Inpaint 🆎 stops serving both. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as functional + +from ._inpaint_crop_helpers import STITCHER_KIND + +# The VAE downsamples by 8 and the DiT patchifies by 2, so a reference edge +# that is not a multiple of 16 lands on a partial patch and the token grid +# stops lining up with the canvas grid. +_PATCH_MULTIPLE = 16 + +# A reference is there to say "this is the picture", not to carry detail - +# the canvas latent carries that. Past a few hundred pixels the extra tokens +# cost attention without changing the result. +REFERENCE_MAX_EDGE = 384 + +_FULL_CANVAS = [0.0, 0.0, 1.0, 1.0] + + +def snap16(value: int) -> int: + """Nearest positive multiple of 16, never below 16.""" + return max(_PATCH_MULTIPLE, int(round(value / 16.0)) * _PATCH_MULTIPLE) + + +def reference_size(width: int, height: int, max_edge: int = REFERENCE_MAX_EDGE) -> tuple[int, int]: + """The /16 size a reference of ``width`` x ``height`` is fitted to. + + ``max_edge`` of 0 or less keeps the size and only snaps it. + """ + width = max(1, int(width)) + height = max(1, int(height)) + if max_edge > 0 and max(width, height) > max_edge: + scale = float(max_edge) / float(max(width, height)) + return snap16(int(round(width * scale))), snap16(int(round(height * scale))) + return snap16(width), snap16(height) + + +def build_reference_image( + image: torch.Tensor, max_edge: int = REFERENCE_MAX_EDGE +) -> torch.Tensor: + """Fit a BHWC image to a /16 reference, downscaling to ``max_edge``. + + Returns a copy when nothing needs resizing, so a caller can never hand + the same tensor to two consumers and have one of them mutate it. + """ + if image.ndim != 4: + raise ValueError("Krea 2 reference expected a BHWC IMAGE batch.") + height, width = int(image.shape[1]), int(image.shape[2]) + target_width, target_height = reference_size(width, height, max_edge) + if (target_height, target_width) == (height, width): + return image.clone() + + moved = image.movedim(-1, 1).contiguous() + resized = functional.interpolate( + moved, + size=(target_height, target_width), + mode="bilinear", + align_corners=False, + # Downscaling without it aliases the reference into a different image. + antialias=target_width < width or target_height < height, + ) + return resized.movedim(1, -1).contiguous() + + +def extract_bbox_norm(stitcher: dict | None) -> list[float]: + """Where the source sits in the canvas, as ``[x0, y0, x1, y1]`` in 0..1. + + Reads the bbox padding recorded on the stitcher. Falls back to the crop + rectangle when the stitcher predates that field, and finally to the whole + canvas - which is the correct answer for a stitcher whose crop *is* the + canvas, and a harmless one otherwise, since a reference placed over the + full frame is what an unpatched model already assumes. + """ + if not isinstance(stitcher, dict) or stitcher.get("kind") != STITCHER_KIND: + return list(_FULL_CANVAS) + + normalized = stitcher.get("bbox_normalized") + if normalized is not None: + return [float(value) for value in normalized] + + canvas = stitcher.get("canvas") + if canvas is None or getattr(canvas, "ndim", 0) != 4: + return list(_FULL_CANVAS) + height, width = float(canvas.shape[1]), float(canvas.shape[2]) + + pixels = stitcher.get("source_bbox") + if pixels is not None: + x0, y0, x1, y1 = (float(value) for value in pixels) + return [x0 / width, y0 / height, x1 / width, y1 / height] + + crop_x, crop_y, crop_w, crop_h = stitcher.get( + "crop_to_canvas", (0, 0, int(width), int(height)) + ) + return [ + float(crop_x) / width, + float(crop_y) / height, + float(crop_x + crop_w) / width, + float(crop_y + crop_h) / height, + ] + + +__all__ = [ + "REFERENCE_MAX_EDGE", + "build_reference_image", + "extract_bbox_norm", + "reference_size", + "snap16", +] diff --git a/nodes/node_krea2_encode.py b/nodes/node_krea2_encode.py new file mode 100644 index 0000000..a58c127 --- /dev/null +++ b/nodes/node_krea2_encode.py @@ -0,0 +1,168 @@ +"""Krea 2 Encode 🆎.""" + +from __future__ import annotations + +from ._krea2_helpers import REFERENCE_MAX_EDGE, build_reference_image + +try: # Offline tests import this module without ComfyUI. + import node_helpers +except ImportError: + node_helpers = None + +try: # Older cores have no Krea 2 template; the encode still works without it. + from comfy.text_encoders.krea2 import KREA2_TEMPLATE +except ImportError: + KREA2_TEMPLATE = None + + +class AusBossKrea2Encode: + CATEGORY = "🆎 AusBoss/Krea2" + DESCRIPTION = ( + "Encodes a Krea 2 prompt pair and attaches reference latents in one " + "node. Wire a VAE and a reference image and the reference is encoded " + "and appended to the positive conditioning; leave them unwired and " + "this is a plain two-prompt encoder. The negative comes out of the " + "same node so a turbo graph at CFG 1.0 still has something to plug in " + "without a second text encode sitting there doing nothing." + ) + SEARCH_ALIASES = [ + "krea2 encode", + "krea 2 encode", + "krea2 outpaint encode", + "reference conditioning", + "ausboss", + ] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip": ( + "CLIP", + {"tooltip": "The Krea 2 text encoder."}, + ), + "prompt": ( + "STRING", + { + "multiline": True, + "dynamicPrompts": True, + "default": "", + "tooltip": ( + "Describe the whole finished canvas, not just the " + "new area — the model is generating all of it and " + "matching the reference where it must." + ), + }, + ), + }, + "optional": { + "negative_prompt": ( + "STRING", + { + "multiline": True, + "dynamicPrompts": True, + "default": "", + "tooltip": ( + "Ignored at CFG 1.0, which is where turbo runs — " + "an empty negative is the normal case there." + ), + }, + ), + "vae": ( + "VAE", + { + "tooltip": ( + "Needed to turn the reference images into latents. " + "Without it the references are skipped." + ) + }, + ), + "reference": ( + "IMAGE", + { + "tooltip": ( + "The `reference` output of Load Image + Pad 🆎 — " + "the unpadded source. Any image works; it is " + "fitted to a multiple of 16 first." + ) + }, + ), + "extra_image": ( + "IMAGE", + {"tooltip": "A second reference, e.g. a style or character plate."}, + ), + "vlm_reference": ( + "BOOLEAN", + { + "default": False, + "tooltip": ( + "Also show the references to the vision tower, so " + "the text encoder describes them. Off for outpaint " + "— the latents already carry the picture, and the " + "description tends to pull the result toward a " + "paraphrase of the source." + ), + }, + ), + }, + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING") + RETURN_NAMES = ("positive", "negative") + OUTPUT_TOOLTIPS = ( + "Prompt conditioning with the reference latents attached.", + "Negative prompt conditioning.", + ) + FUNCTION = "encode" + + def encode( + self, + clip, + prompt, + negative_prompt="", + vae=None, + reference=None, + extra_image=None, + vlm_reference=False, + ): + references = [image for image in (reference, extra_image) if image is not None] + + images_vl = [] + image_prompt = "" + if vlm_reference and references: + for index, image in enumerate(references): + images_vl.append(image) + image_prompt += ( + f"Picture {index + 1}: <|vision_start|><|image_pad|><|vision_end|>" + ) + + if images_vl: + if KREA2_TEMPLATE is not None: + tokens = clip.tokenize( + image_prompt + prompt, + images=images_vl, + llama_template=KREA2_TEMPLATE, + ) + else: + tokens = clip.tokenize(image_prompt + prompt, images=images_vl) + else: + tokens = clip.tokenize(prompt) + positive = clip.encode_from_tokens_scheduled(tokens) + + if vae is not None and references and node_helpers is not None: + latents = [ + vae.encode(build_reference_image(image, REFERENCE_MAX_EDGE)[:, :, :, :3]) + for image in references + ] + positive = node_helpers.conditioning_set_values( + positive, {"reference_latents": latents}, append=True + ) + + negative = clip.encode_from_tokens_scheduled(clip.tokenize(negative_prompt or "")) + return (positive, negative) + + +NODE_CLASS_MAPPINGS = {"AUSBOSS_NODES_Krea2Encode": AusBossKrea2Encode} +NODE_DISPLAY_NAME_MAPPINGS = {"AUSBOSS_NODES_Krea2Encode": "Krea 2 Encode 🆎"} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/node_krea2_model_patch.py b/nodes/node_krea2_model_patch.py new file mode 100644 index 0000000..46c82c4 --- /dev/null +++ b/nodes/node_krea2_model_patch.py @@ -0,0 +1,179 @@ +"""Krea 2 Outpaint Model Patch 🆎.""" + +from __future__ import annotations + +import math + +from ._krea2_helpers import extract_bbox_norm + +# comfy and _krea2_core (which reaches into comfy's flux layers) are imported +# inside patch(), not here. A core update that moves one of those internals +# then surfaces as a clear error on the node that needs it, instead of +# deleting the node from the menu at import time and leaving a saved workflow +# reporting it as missing. + + +class AusBossKrea2OutpaintModelPatch: + CATEGORY = "🆎 AusBoss/Krea2" + DESCRIPTION = ( + "Teaches Krea 2 where a reference image belongs on the canvas. " + "Reference latents normally arrive with no position, so the model " + "treats them as a loose style hint; this registers them into the " + "target grid at the rectangle the stitcher reports, which is what " + "makes an outpaint continue the source instead of reinventing it. " + "Place it after any LoRA loader and before the sampler." + ) + SEARCH_ALIASES = [ + "krea2 outpaint patch", + "krea2 model patch", + "registered reference", + "reference placement", + "ausboss", + ] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ( + "MODEL", + {"tooltip": "A Krea 2 model. Patch last, after any LoRAs."}, + ), + "stitcher": ( + "AUSBOSS_STITCHER", + { + "tooltip": ( + "From Load Image + Pad 🆎 — supplies the rectangle " + "the source occupies on the canvas. A stitcher " + "without one places the reference over the full " + "frame." + ) + }, + ), + "kv_cache": ( + "BOOLEAN", + { + "default": True, + "tooltip": ( + "Compute the reference's keys and values once per " + "run instead of every step. The reference does not " + "change while sampling, so this is free speed; turn " + "it off only to rule the cache out when debugging." + ), + }, + ), + }, + } + + RETURN_TYPES = ("MODEL",) + OUTPUT_TOOLTIPS = ( + "The model with reference tokens registered into the canvas grid.", + ) + FUNCTION = "patch" + + def patch(self, model, stitcher, kv_cache=True): + import comfy.conds + + from . import _krea2_core as core + + bbox_norm = extract_bbox_norm(stitcher) + + patched = model.clone() + base_model = patched.model + dit = patched.get_model_object("diffusion_model") + + orig_extra_conds = base_model.extra_conds + orig_extra_conds_shapes = base_model.extra_conds_shapes + orig_forward = dit.forward + + def extra_conds(**kwargs): + out = orig_extra_conds(**kwargs) + ref_latents = kwargs.get("reference_latents", None) + if ref_latents is not None: + out["ref_latents"] = comfy.conds.CONDList( + [base_model.process_latent_in(lat) for lat in ref_latents] + ) + return out + + def extra_conds_shapes(**kwargs): + out = orig_extra_conds_shapes(**kwargs) + ref_latents = kwargs.get("reference_latents", None) + if ref_latents is not None: + out["ref_latents"] = list( + [1, 16, sum(math.prod(a.size()) for a in ref_latents) // 16] + ) + return out + + state = {"last_sigma": None, "caches": {}} + + def forward( + x, + timesteps, + context, + attention_mask=None, + transformer_options={}, + ref_latents=None, + **kwargs, + ): + # Nothing to place: leave the model exactly as it was. + 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, + ) + + sigma = float(timesteps.max()) + sample_sigmas = transformer_options.get("sample_sigmas", None) + # Sigma climbing back up means a new run started, so the cached + # reference belongs to the previous one. + new_run = state["last_sigma"] is None or sigma > state["last_sigma"] + if ( + sample_sigmas is not None + and sigma == float(sample_sigmas[0]) + and sigma != state["last_sigma"] + ): + new_run = True + if new_run: + state["caches"].clear() + state["last_sigma"] = sigma + + batch = x.shape[0] * (x.shape[2] if x.ndim == 5 else 1) + ref_kv = None + key = None + if kv_cache: + key = core._ref_fingerprint(ref_latents, batch, bbox_norm) + ref_kv = state["caches"].get(key) + if ref_kv is None: + ref_kv = core._precompute_ref_kv( + dit, + x, + timesteps, + ref_latents, + transformer_options, + bbox_norm=bbox_norm, + ) + if kv_cache: + state["caches"][key] = ref_kv + return core._forward_with_cached_refs( + dit, x, timesteps, context, ref_kv, 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,) + + +NODE_CLASS_MAPPINGS = { + "AUSBOSS_NODES_Krea2OutpaintModelPatch": AusBossKrea2OutpaintModelPatch, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AUSBOSS_NODES_Krea2OutpaintModelPatch": "Krea 2 Outpaint Model Patch 🆎", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/node_load_image_pad.py b/nodes/node_load_image_pad.py index ec450b4..6ef7f12 100644 --- a/nodes/node_load_image_pad.py +++ b/nodes/node_load_image_pad.py @@ -6,6 +6,7 @@ import numpy as np import torch from ._inpaint_crop_helpers import build_canvas_stitcher +from ._krea2_helpers import build_reference_image from ._media_helpers import list_input_images, load_image_frames, resolve_input_path from ._pad_helpers import ( PAD_MODES, @@ -182,8 +183,11 @@ class AusBossLoadImagePad: } } - RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "AUSBOSS_STITCHER") - RETURN_NAMES = ("image", "mask", "width", "height", "stitcher") + # `reference` is appended, never inserted: a workflow stores links by + # output slot index, so putting it anywhere else would silently rewire + # every saved graph that already reads width/height/stitcher. + RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "AUSBOSS_STITCHER", "IMAGE") + RETURN_NAMES = ("image", "mask", "width", "height", "stitcher", "reference") OUTPUT_TOOLTIPS = ( "The padded image; the original pixels are untouched (resized only " "when a megapixel target is set).", @@ -192,7 +196,12 @@ class AusBossLoadImagePad: "Final canvas width after multiple/megapixel rounding.", "Final canvas height after multiple/megapixel rounding.", "Hand to Stitch Inpaint 🆎 with the sampled result to keep only the " - "new padding and restore the original pixels bit-identically.", + "new padding and restore the original pixels bit-identically. Also " + "carries where the source sits on the canvas, which Krea 2 Outpaint " + "Model Patch 🆎 reads to place the reference.", + "The source alone, no padding, fitted to a small multiple of 16 — " + "the reference image for Krea 2 Encode 🆎 and other reference " + "conditioning. Wire it nowhere and nothing is computed.", ) FUNCTION = "load_pad" @@ -236,10 +245,26 @@ class AusBossLoadImagePad: mask = feather_pad_mask( mask, plan["left"], plan["top"], plan["right"], plan["bottom"], int(feather) ) + # Padding knows exactly where the source landed, so the bbox rides on + # the stitcher rather than on a wire that can be left unplugged. + bbox = ( + plan["left"], + plan["top"], + plan["left"] + plan["source_width"], + plan["top"] + plan["source_height"], + ) # The padded canvas is the stitch base, so whatever the sampler does # outside the feathered band is discarded and the source survives. - stitcher = build_canvas_stitcher(output, mask) - return output, mask, int(plan["width"]), int(plan["height"]), stitcher + stitcher = build_canvas_stitcher(output, mask, bbox=bbox) + reference = build_reference_image(frames) + return ( + output, + mask, + int(plan["width"]), + int(plan["height"]), + stitcher, + reference, + ) @classmethod def VALIDATE_INPUTS(cls, image, **_values): diff --git a/scripts/validate_nodes.py b/scripts/validate_nodes.py index fc28e23..8dcc9ae 100644 --- a/scripts/validate_nodes.py +++ b/scripts/validate_nodes.py @@ -51,6 +51,8 @@ PUBLIC_NODE_IDS = { "AUSBOSS_NODES_FrameInterpolate", "AUSBOSS_NODES_ImageCropRotatePad", "AUSBOSS_NODES_ImageSize", + "AUSBOSS_NODES_Krea2Encode", + "AUSBOSS_NODES_Krea2OutpaintModelPatch", "AUSBOSS_NODES_LaMaInpaint", "AUSBOSS_NODES_LmStudioChat", "AUSBOSS_NODES_LoadImagePad", diff --git a/tests/test_inpaint_crop_helpers.py b/tests/test_inpaint_crop_helpers.py index 59d3860..b1c8853 100644 --- a/tests/test_inpaint_crop_helpers.py +++ b/tests/test_inpaint_crop_helpers.py @@ -801,3 +801,46 @@ class CanvasStitcherTests(unittest.TestCase): self.assertEqual(stitcher["kind"], inpaint_helpers.STITCHER_KIND) self.assertEqual(stitcher["crop_to_canvas"], (0, 0, 40, 48)) self.assertEqual(stitcher["canvas_to_original"], (0, 0, 40, 48)) + + +class CanvasStitcherBboxTests(unittest.TestCase): + """Where the source sits inside the canvas rides on the stitcher, so a + model that places reference tokens on the canvas grid cannot be handed a + canvas without its placement.""" + + def setUp(self): + torch.manual_seed(11) + self.canvas = torch.rand(1, 48, 40, 3) + self.mask = torch.zeros(1, 48, 40) + self.mask[:, :8, :] = 1.0 + + def test_bbox_is_absent_unless_given(self): + stitcher = build_canvas_stitcher(self.canvas, self.mask) + self.assertNotIn("source_bbox", stitcher) + self.assertNotIn("bbox_normalized", stitcher) + + def test_bbox_is_stored_in_pixels_and_normalized(self): + stitcher = build_canvas_stitcher(self.canvas, self.mask, bbox=(8, 8, 32, 40)) + self.assertEqual(stitcher["source_bbox"], (8, 8, 32, 40)) + # Canvas is 40 wide, 48 tall: x over width, y over height. + self.assertEqual( + stitcher["bbox_normalized"], [8 / 40, 8 / 48, 32 / 40, 40 / 48] + ) + + def test_a_full_canvas_bbox_normalizes_to_the_unit_square(self): + stitcher = build_canvas_stitcher(self.canvas, self.mask, bbox=(0, 0, 40, 48)) + self.assertEqual(stitcher["bbox_normalized"], [0.0, 0.0, 1.0, 1.0]) + + def test_floats_are_taken_as_pixels(self): + stitcher = build_canvas_stitcher(self.canvas, self.mask, bbox=(8.0, 8.0, 32.0, 40.0)) + self.assertEqual(stitcher["source_bbox"], (8, 8, 32, 40)) + + def test_stitching_ignores_the_bbox_entirely(self): + # apply_stitch never reads it, so an older stitcher still stitches and + # a newer one stitches identically. + plain = build_canvas_stitcher(self.canvas, self.mask) + with_bbox = build_canvas_stitcher(self.canvas, self.mask, bbox=(8, 8, 32, 40)) + sampled = torch.zeros_like(self.canvas) + self.assertTrue( + torch.equal(apply_stitch(plain, sampled), apply_stitch(with_bbox, sampled)) + ) diff --git a/tests/test_krea2_helpers.py b/tests/test_krea2_helpers.py new file mode 100644 index 0000000..084f333 --- /dev/null +++ b/tests/test_krea2_helpers.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +from pathlib import Path +import sys +import unittest + +import torch + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) +if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"): + del sys.modules["nodes"] + +from nodes._inpaint_crop_helpers import build_canvas_stitcher +from nodes._krea2_helpers import ( + REFERENCE_MAX_EDGE, + build_reference_image, + extract_bbox_norm, + reference_size, + snap16, +) + + +class Snap16Tests(unittest.TestCase): + """The VAE downsamples by 8 and the DiT patchifies by 2, so a reference + edge off a multiple of 16 lands on a partial patch.""" + + def test_rounds_to_the_nearest_multiple(self): + self.assertEqual(snap16(16), 16) + self.assertEqual(snap16(23), 16) + self.assertEqual(snap16(24), 32) + self.assertEqual(snap16(384), 384) + + def test_never_returns_zero(self): + # A zero-size reference would encode to an empty latent. + self.assertEqual(snap16(0), 16) + self.assertEqual(snap16(1), 16) + self.assertEqual(snap16(-40), 16) + + +class ReferenceSizeTests(unittest.TestCase): + def test_a_small_reference_is_only_snapped(self): + self.assertEqual(reference_size(64, 48), (64, 48)) + self.assertEqual(reference_size(70, 50), (64, 48)) + + def test_the_long_edge_is_capped_and_aspect_roughly_kept(self): + width, height = reference_size(1000, 500) + self.assertLessEqual(max(width, height), REFERENCE_MAX_EDGE) + self.assertEqual(width % 16, 0) + self.assertEqual(height % 16, 0) + self.assertAlmostEqual(width / height, 2.0, delta=0.15) + + def test_a_tall_source_caps_on_height(self): + width, height = reference_size(500, 1000) + self.assertEqual(height, 384) + self.assertLess(width, height) + + def test_zero_max_edge_means_snap_only(self): + # 1000/16 is exactly 62.5, and round() breaks ties to even: 62 * 16. + self.assertEqual(reference_size(1000, 500, 0), (992, 496)) + + +class BuildReferenceImageTests(unittest.TestCase): + def test_an_already_fitting_image_comes_back_unchanged(self): + image = torch.rand(1, 48, 64, 3) + out = build_reference_image(image) + self.assertTrue(torch.equal(out, image)) + + def test_it_returns_a_copy_not_the_input(self): + # Two consumers must not be able to mutate each other's reference. + image = torch.rand(1, 48, 64, 3) + out = build_reference_image(image) + out[0, 0, 0, 0] = 0.5 + self.assertNotEqual(float(image[0, 0, 0, 0]), 0.5) + + def test_a_large_image_is_downscaled_to_multiples_of_16(self): + out = build_reference_image(torch.rand(1, 500, 1000, 3)) + _batch, height, width, channels = out.shape + self.assertLessEqual(max(width, height), REFERENCE_MAX_EDGE) + self.assertEqual((height % 16, width % 16), (0, 0)) + self.assertEqual(channels, 3) + + def test_the_batch_survives(self): + out = build_reference_image(torch.rand(4, 500, 1000, 3)) + self.assertEqual(out.shape[0], 4) + + def test_a_non_bhwc_tensor_is_rejected_by_name(self): + with self.assertRaises(ValueError) as caught: + build_reference_image(torch.rand(48, 64, 3)) + self.assertIn("BHWC", str(caught.exception)) + + +class ExtractBboxNormTests(unittest.TestCase): + def setUp(self): + self.canvas = torch.zeros(1, 1024, 2048, 3) + self.mask = torch.zeros(1, 1024, 2048) + + def test_it_reads_the_recorded_bbox(self): + stitcher = build_canvas_stitcher( + self.canvas, self.mask, bbox=(0, 128, 1024, 896) + ) + self.assertEqual( + extract_bbox_norm(stitcher), [0.0, 128 / 1024, 0.5, 896 / 1024] + ) + + def test_it_falls_back_to_the_crop_rectangle(self): + # A stitcher from before the bbox existed, or one from Crop For + # Inpaint, still says where its region sits. + stitcher = build_canvas_stitcher(self.canvas, self.mask) + stitcher["crop_to_canvas"] = (256, 128, 1024, 768) + self.assertEqual( + extract_bbox_norm(stitcher), + [256 / 2048, 128 / 1024, 1280 / 2048, 896 / 1024], + ) + + def test_source_bbox_alone_is_enough(self): + stitcher = build_canvas_stitcher(self.canvas, self.mask) + stitcher["source_bbox"] = (256, 128, 1280, 896) + self.assertEqual( + extract_bbox_norm(stitcher), + [256 / 2048, 128 / 1024, 1280 / 2048, 896 / 1024], + ) + + def test_a_full_canvas_stitcher_reports_the_unit_square(self): + self.assertEqual( + extract_bbox_norm(build_canvas_stitcher(self.canvas, self.mask)), + [0.0, 0.0, 1.0, 1.0], + ) + + def test_junk_falls_back_to_the_full_canvas(self): + # The full frame is what an unpatched model already assumes, so a + # bad stitcher degrades instead of raising mid-sample. + for value in (None, {}, {"kind": "something else"}, "not a stitcher", 42): + self.assertEqual(extract_bbox_norm(value), [0.0, 0.0, 1.0, 1.0]) + + +class Krea2NodeContractTests(unittest.TestCase): + def test_model_patch_contract(self): + from nodes.node_krea2_model_patch import ( + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + ) + + key = "AUSBOSS_NODES_Krea2OutpaintModelPatch" + self.assertIn(key, NODE_CLASS_MAPPINGS) + self.assertEqual( + NODE_DISPLAY_NAME_MAPPINGS[key], "Krea 2 Outpaint Model Patch 🆎" + ) + cls = NODE_CLASS_MAPPINGS[key] + self.assertEqual(cls.CATEGORY, "🆎 AusBoss/Krea2") + self.assertEqual(cls.RETURN_TYPES, ("MODEL",)) + required = cls.INPUT_TYPES()["required"] + self.assertEqual(required["model"][0], "MODEL") + self.assertEqual(required["stitcher"][0], "AUSBOSS_STITCHER") + self.assertTrue(required["kv_cache"][1]["default"]) + + def test_encode_contract(self): + from nodes.node_krea2_encode import ( + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + ) + + key = "AUSBOSS_NODES_Krea2Encode" + self.assertIn(key, NODE_CLASS_MAPPINGS) + self.assertEqual(NODE_DISPLAY_NAME_MAPPINGS[key], "Krea 2 Encode 🆎") + cls = NODE_CLASS_MAPPINGS[key] + self.assertEqual(cls.CATEGORY, "🆎 AusBoss/Krea2") + self.assertEqual(cls.RETURN_TYPES, ("CONDITIONING", "CONDITIONING")) + self.assertEqual(cls.RETURN_NAMES, ("positive", "negative")) + # Everything but clip and prompt is optional, so the node drops into a + # plain two-prompt graph with no VAE and no reference wired. + self.assertEqual(set(cls.INPUT_TYPES()["required"]), {"clip", "prompt"}) + + def test_encode_without_a_vae_skips_the_reference(self): + from nodes.node_krea2_encode import NODE_CLASS_MAPPINGS + + class StubClip: + def tokenize(self, text, **_kwargs): + return {"text": text} + + def encode_from_tokens_scheduled(self, tokens): + return [[tokens["text"], {}]] + + cls = NODE_CLASS_MAPPINGS["AUSBOSS_NODES_Krea2Encode"] + positive, negative = cls().encode( + StubClip(), "a house", "blurry", reference=torch.rand(1, 48, 64, 3) + ) + self.assertEqual(positive[0][0], "a house") + self.assertEqual(negative[0][0], "blurry") + self.assertNotIn("reference_latents", positive[0][1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pad_helpers.py b/tests/test_pad_helpers.py index 0df685f..eb577e2 100644 --- a/tests/test_pad_helpers.py +++ b/tests/test_pad_helpers.py @@ -327,14 +327,16 @@ class LoadImagePadNodeTests(unittest.TestCase): cls = self.make_node() self.assertIn("AusBoss/Image", cls.CATEGORY) self.assertEqual( - cls.RETURN_TYPES, ("IMAGE", "MASK", "INT", "INT", "AUSBOSS_STITCHER") + cls.RETURN_TYPES, + ("IMAGE", "MASK", "INT", "INT", "AUSBOSS_STITCHER", "IMAGE"), ) self.assertEqual( - cls.RETURN_NAMES, ("image", "mask", "width", "height", "stitcher") + cls.RETURN_NAMES, + ("image", "mask", "width", "height", "stitcher", "reference"), ) with tempfile.TemporaryDirectory() as tmp: path = self.write_image(tmp) - (image, mask, width, height, _), _ = self.run_node(cls, path) + (image, mask, width, height, _, _), _ = self.run_node(cls, path) self.assertEqual((width, height), (72, 56)) # 70x56 ceiled to 8 self.assertEqual(tuple(image.shape), (1, 56, 72, 3)) self.assertEqual(tuple(mask.shape), (1, 56, 72)) @@ -348,7 +350,7 @@ class LoadImagePadNodeTests(unittest.TestCase): cls = self.make_node() with tempfile.TemporaryDirectory() as tmp: path = self.write_image(tmp) - (image, mask, width, height, _), values = self.run_node( + (image, mask, width, height, _, _), values = self.run_node( cls, path, target_megapixels=0.05 ) plan = plan_pad_canvas(64, 48, 2, 3, 4, 5, 8, 0.05) @@ -375,11 +377,70 @@ class LoadImagePadNodeTests(unittest.TestCase): cls = self.make_node() with tempfile.TemporaryDirectory() as tmp: path = self.write_image(tmp) - (_, mask, _, _, _), _ = self.run_node(cls, path, feather=6) + (_, mask, _, _, _, _), _ = self.run_node(cls, path, feather=6) values = torch.unique(mask).tolist() self.assertTrue(any(0.0 < value < 1.0 for value in values)) self.assertEqual(float(mask[0, 0, 0]), 1.0) # padding stays solid + def test_reference_is_the_unpadded_source_snapped_to_16(self): + import tempfile + + cls = self.make_node() + with tempfile.TemporaryDirectory() as tmp: + path = self.write_image(tmp) + (_, _, _, _, _, reference), _ = self.run_node(cls, path) + # 64x48 is already under the 384 cap and already /16. + self.assertEqual(tuple(reference.shape), (1, 48, 64, 3)) + # It is the SOURCE, not the canvas: no padding colour in it. + self.assertGreater(float(reference.min()), 0.0) + + def test_reference_downscales_a_large_source_to_a_multiple_of_16(self): + import tempfile + + cls = self.make_node() + with tempfile.TemporaryDirectory() as tmp: + path = self.write_image(tmp, width=1000, height=750) + (_, _, _, _, _, reference), _ = self.run_node(cls, path) + _batch, height, width, _channels = reference.shape + self.assertLessEqual(max(width, height), 384) + self.assertEqual(width % 16, 0) + self.assertEqual(height % 16, 0) + + def test_stitcher_carries_where_the_source_sits(self): + import tempfile + + cls = self.make_node() + with tempfile.TemporaryDirectory() as tmp: + path = self.write_image(tmp) + (_, mask, width, height, stitcher, _), _ = self.run_node(cls, path) + # pad_left=2, pad_top=3 on a 64x48 source. + self.assertEqual(stitcher["source_bbox"], (2, 3, 66, 51)) + x0, y0, x1, y1 = stitcher["bbox_normalized"] + self.assertAlmostEqual(x0, 2 / width) + self.assertAlmostEqual(y0, 3 / height) + self.assertAlmostEqual(x1, 66 / width) + self.assertAlmostEqual(y1, 51 / height) + # The bbox is exactly the region the mask protects. + inner = mask[0, 3:51, 2:66] + self.assertEqual(float(inner.sum()), 0.0) + self.assertEqual(float(mask.sum()), float(mask.numel() - inner.numel())) + + def test_bbox_tracks_the_megapixel_resize(self): + import tempfile + + cls = self.make_node() + with tempfile.TemporaryDirectory() as tmp: + path = self.write_image(tmp) + (_, mask, _, _, stitcher, _), _ = self.run_node( + cls, path, target_megapixels=0.05 + ) + plan = plan_pad_canvas(64, 48, 2, 3, 4, 5, 8, 0.05) + x0, y0, x1, y1 = stitcher["source_bbox"] + self.assertEqual((x0, y0), (plan["left"], plan["top"])) + self.assertEqual(x1 - x0, plan["source_width"]) + self.assertEqual(y1 - y0, plan["source_height"]) + self.assertEqual(float(mask[0, y0:y1, x0:x1].sum()), 0.0) + def test_validation_and_fingerprint(self): import tempfile