nodes/ausboss-ComfyUI-AusBoss
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #41 from ausboss/feat/promote-krea2-outpaint

Browse Source 
Promote the Krea 2 outpaint pair from the lab
This commit is contained in:
ausboss
2026-08-17 22:26:08 -05:00
committed by GitHub
parent bf87fd92a6 ec50d3acfe
commit 38aeec6d09
15 changed files with 1186 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+27
View File
@@ -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
README.md
+9 -1
View File
@@ -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 🆎
__init__.py
+2
View File
@@ -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",
js/docs/AUSBOSS_NODES_Krea2Encode.md
+43
View File
@@ -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.
js/docs/AUSBOSS_NODES_Krea2OutpaintModelPatch.md
+40
View File
@@ -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 🆎.
nodes/_inpaint_crop_helpers.py
+25 -2
View File
@@ -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(
nodes/_krea2_core.py
+240
View File
@@ -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)
nodes/_krea2_helpers.py
+118
View File
@@ -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",
]
nodes/node_krea2_encode.py
+168
View File
@@ -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"]
nodes/node_krea2_model_patch.py
+179
View File
@@ -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"]
nodes/node_load_image_pad.py
+30 -5
View File
@@ -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):
scripts/validate_nodes.py
+2
View File
@@ -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",
tests/test_inpaint_crop_helpers.py
+43
View File
@@ -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))
)
tests/test_krea2_helpers.py
+194
View File
@@ -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()
tests/test_pad_helpers.py
+66 -5
View File
@@ -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