nodes/shootthesound-ComfyUI-Angelo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
d9327e8632b4e5a9182f8a5146ba06cb01d3d8a8
shootthesound-ComfyUI-Angelo/angelo_nodes.py
T
shootthesound d9327e8632 feat: Large Image Refine — chess-pattern Xtra-Fine boxes over the canvas 
New ▦ button beside ⬆ 2x Pixel. The canvas divides into ~1MP boxes
that exact-tile it (no overlap); each box runs the Xtra-Fine pipeline
(ref 0.4, denoise 0.55, 176px ctx pad, hard feather-0 masks) under
'restore the image. make it clear and sharp.' — in CHESS order, so
the second colour's boxes refine with already-refined neighbours
visible in their context pad and match them. Compositing is
Xtra-Fine's bit-exact latent blend per box — no pixel feathering
anywhere in the path. Whole pass = ONE history entry; per-box seeds
derive from the Seed widget. The MP-target nudge forces the crop
path so the whole-canvas shortcut can never fire. Box size is one
constant (_LIR_BOX_MP) for later tweaking.
2026-06-11 00:01:30 +01:00

3066 lines
158 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""AngeloRefine — single-node click-to-refine sampler.
How it works:
- On first run (no clicks yet), the node decodes the incoming latent
via VAE and shows the image in its own preview area. The latent is
cached server-side keyed by node id.
- The JS frontend attaches click handling to that preview. When the
user clicks a region they want improved, the JS updates this node's
click_x / click_y / click_seq widgets and re-queues the workflow.
- On re-run, the node detects the new click_seq, builds a feathered
circular mask centred on (click_x, click_y) with radius click_radius,
re-noises the cached latent partially (`denoise`), and re-samples
`steps` steps with the mask applied. ComfyUI's noise_mask handling
re-stitches the original outside the mask at each step (the noise-
injection inpaint pattern), so the unclicked region stays preserved.
- The refined latent is cached for the next click, so successive clicks
keep refining the SAME working latent rather than starting over.
- Toggle the `reset` widget (or press the Reset button the JS adds) to
throw away the cache and start fresh from the incoming latent.
The cache is in-process state only — restart of ComfyUI clears it.
"""
from __future__ import annotations
import copy
import hashlib
import json
import math
import os
import torch
import comfy.model_management
import comfy.sample
import comfy.samplers
import comfy.utils
import latent_preview
import folder_paths
import node_helpers
# Reuse PreviewImage's save machinery for the preview output.
import nodes as comfy_nodes
# Per-node state cache:
# unique_id -> {
# "history": list[tuple[Tensor, Tensor | None]], # stack of (latent, pixels) (oldest first, current = [-1])
# "click_seq": int, # last processed click_seq from JS
# "undo_seq": int, # last processed undo_seq from JS
# "redo_seq": int, # last processed redo_seq from JS
# "redo_stack": list, # entries Undo popped, awaiting Redo (cleared on new edit)
# "source_latent": Tensor, # session base latent, for the source_image output
# "source_pixels": Tensor|None,# decoded source base (lazy, cached)
# "fingerprint": str, # hash of incoming latent; mismatch = upstream changed
# }
_STATE: dict[str, dict] = {}
# ---- Custom-sampler helpers (#8) -----------------------------------------
# These three helpers + the _do_sample wrapper below let users wire a
# GUIDER + SAMPLER + SIGMAS bundle through Angelo — Overrides into Angelo
# in place of the toolbar's steps / cfg / sampler_name / scheduler. The
# helpers themselves are taken verbatim from @KursatAs's customSampler
# branch (https://github.com/KursatAs/ComfyUI-Angelo/tree/customSampler);
# the integration shape (optional via Overrides rather than required on
# the main node) is the Angelo adaptation. Full credit to him for the
# NAG-Extended fix encoded in _guider_sample.
def _guider_sample(
temp_g,
noise: torch.Tensor,
latent: torch.Tensor,
sampler,
sigmas: torch.Tensor,
*,
denoise_mask: torch.Tensor | None = None,
callback=None,
disable_pbar: bool = False,
seed: int | None = None,
) -> torch.Tensor:
"""Device-safe wrapper around guider.sample(). (From @KursatAs.)
ComfyUI's built-in CFGGuider.sample() moves noise, latent, and
denoise_mask to the model's load device before sampling. Some
third-party extensions (e.g. ComfyUI-NAG-Extended) override
inner_sample() without repeating that movement, so CPU tensors
survive into the k-sampler's inpaint path where they collide with
GPU tensors and raise a device-mismatch RuntimeError. Moving
everything to load_device here is a safe no-op when the built-in
already does it, and fixes the crash for extensions that don't."""
device = temp_g.model_patcher.load_device
noise = noise.to(device)
latent = latent.to(device)
if denoise_mask is not None:
denoise_mask = denoise_mask.to(device)
samples = temp_g.sample(
noise, latent, sampler, sigmas,
denoise_mask=denoise_mask,
callback=callback,
disable_pbar=disable_pbar,
seed=seed,
)
# Mirror comfy.sample.sample()'s exit move so the rest of Angelo's
# pipeline (VAE decode, pixel composite, latent blend in Fine Upscale)
# sees the same intermediate device + dtype it does on the default
# path. Without this, the guider path returns the sampled latent on
# the model's load_device (cuda) while cached_pixels / mask are on
# intermediate_device (typically CPU), causing a device-mismatch at
# the composite step in _refine_with_fine_upscaling.
return samples.to(
device=comfy.model_management.intermediate_device(),
dtype=comfy.model_management.intermediate_dtype(),
)
def _guider_with_conds(guider, positive, negative):
"""Copy a wired GUIDER and apply Angelo's per-call positive/negative
conds to it. Handles both CFGGuider (takes both conds) and
BasicGuider (positive only). (From @KursatAs.) Lets the user wire
a generic guider once and Angelo keeps using its dynamic conds
(Refine vs Area Prompt vs Smart Inpaint reference_latents etc.)
for each sample call."""
g = copy.copy(guider)
try:
g.set_conds(positive, negative) # comfy.samplers.CFGGuider
except TypeError:
g.set_conds(positive) # BasicGuider / BaseGuider
return g
def _truncate_sigmas_for_denoise(sigmas: torch.Tensor, denoise: float) -> torch.Tensor:
"""Tail-slice the wired SIGMAS tensor by Angelo's per-call denoise,
matching ComfyUI's SplitSigmasDenoise convention. (From @KursatAs.)
A wired sigmas tensor comes pre-baked from a scheduler node at
denoise=1.0; Angelo applies its own refine denoise here so the
refine slider keeps meaning what it did before."""
if denoise >= 1.0:
return sigmas
if denoise <= 0.0:
return sigmas[-1:].new_zeros(2)
n_total = len(sigmas) - 1
n_refine = max(1, round(n_total * denoise))
return sigmas[-(n_refine + 1):]
def _do_sample(
*,
guider,
sampler,
sigmas,
model,
noise: torch.Tensor,
steps: int,
cfg: float,
sampler_name: str,
scheduler: str,
positive,
negative,
source_latent: torch.Tensor,
denoise: float,
callback,
disable_pbar: bool,
seed: int,
noise_mask: torch.Tensor | None = None,
) -> torch.Tensor:
"""Single sample dispatch for every Angelo sample call. If a custom
guider + sampler + sigmas trio is wired through Overrides, take the
custom path (via _guider_sample); otherwise fall through to the
standard comfy.sample.sample(...) path that's existed since v1.0.
All-or-nothing on the trio: partial wiring (e.g. sampler without
sigmas) silently falls through to the default path. Users see the
custom-sampler kwargs only by wiring the full bundle from a proper
GUIDER + SAMPLER + SIGMAS chain in their workflow."""
if guider is not None and sampler is not None and sigmas is not None:
g = _guider_with_conds(guider, positive, negative)
s = _truncate_sigmas_for_denoise(sigmas, denoise)
return _guider_sample(
g, noise, source_latent, sampler, s,
denoise_mask=noise_mask,
callback=callback,
disable_pbar=disable_pbar,
seed=seed,
)
return comfy.sample.sample(
model, noise, steps, cfg, sampler_name, scheduler,
positive, negative, source_latent,
denoise=denoise,
noise_mask=noise_mask,
callback=callback,
disable_pbar=disable_pbar,
seed=seed,
)
# Max number of latents to keep in the undo stack per node. Each FLUX 2
# latent at 832x1776 is ~180 KB (bf16); 10 = ~1.8 MB per node. Cheap.
_HISTORY_CAP: int = 10
# Valid resize methods for the Fine Upscaling crop. All routed through
# comfy.utils.common_upscale which accepts both 4D image tensors and 4D
# latent tensors. lanczos is image-quality; bislerp is latent-aware.
# Default nearest-exact preserves exact sample values (good for
# latents); for the pixel-space path lanczos / bicubic typically look
# better. The user picks one for both ops in the Fine Upscale flow.
_FINE_UPSCALE_RESIZE_METHODS = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp", "lanczos"]
# Smart Guided Inpaint: maps a location dropdown LABEL (set by the JS
# location selector) to a natural-language prefix that's prepended to
# the Area Prompt text before CLIP encoding. The labels here are the
# single source of truth — the JS dropdown lists exactly these keys.
# Order is preserved (py3.7+ dicts) so the JS can build its dropdown
# from the same list via the widget's tooltip / a mirrored array.
_GUIDED_LOCATION_PREFIXES = {
"(none)": "",
"Whole image": "Across the whole image, ",
"Top left": "In the top left of the image, ",
"Top middle": "In the top middle of the image, ",
"Top right": "In the top right of the image, ",
"Middle left": "On the left middle of the image, ",
"Center": "In the center of the image, ",
"Middle right": "On the right middle of the image, ",
"Bottom left": "In the bottom left of the image, ",
"Bottom middle": "In the bottom middle of the image, ",
"Bottom right": "In the bottom right of the image, ",
"Left edge": "Along the left edge of the image, ",
"Right edge": "Along the right edge of the image, ",
"Top edge": "Along the top edge of the image, ",
"Bottom edge": "Along the bottom edge of the image, ",
"Top half": "In the top half of the image, ",
"Bottom half": "In the bottom half of the image, ",
"Left half": "In the left half of the image, ",
"Right half": "In the right half of the image, ",
"Top of the image": "At the top of the image, ",
"Bottom of the image": "At the bottom of the image, ",
}
def _latent_fingerprint(latent: torch.Tensor) -> str:
"""Quick non-cryptographic fingerprint of an incoming latent.
Used to detect when the upstream KSampler has produced a fresh
latent (e.g. user changed the prompt + re-queued), so we can
automatically reset the cache instead of layering refinements
onto a now-irrelevant base.
"""
flat = latent.detach().to(torch.float32).flatten()
n = min(flat.numel(), 1024)
sample = flat[::max(1, flat.numel() // n)][:n]
h = hashlib.sha1()
h.update(str(tuple(latent.shape)).encode())
h.update(sample.cpu().numpy().tobytes())
return h.hexdigest()
def _parse_stroke_points(raw: str) -> list[tuple[float, float]]:
"""Parse the JS-set stroke_points widget. Empty / malformed → []."""
raw = (raw or "").strip()
if not raw:
return []
try:
data = json.loads(raw)
except Exception:
return []
out = []
if not isinstance(data, list):
return []
for item in data:
if isinstance(item, (list, tuple)) and len(item) >= 2:
try:
out.append((float(item[0]), float(item[1])))
except (TypeError, ValueError):
continue
return out
def _stroke_mask_latent(
latent_h: int,
latent_w: int,
stroke_points_pixel: list[tuple[float, float]],
r_latent: float,
scale_x: float,
scale_y: float,
device: torch.device,
) -> torch.Tensor:
"""Vectorised union of circles in latent space, one circle per
point in stroke_points_pixel. Points come in image-pixel coords;
we scale them to latent space here using the per-axis ratios.
Memory: one [N_points, H, W] float tensor briefly. For N=200, a
typical FLUX 2 latent (~52x111), that's about 4 MB peak — fine.
"""
if not stroke_points_pixel:
return torch.zeros((1, latent_h, latent_w), device=device, dtype=torch.float32)
pts = torch.tensor(stroke_points_pixel, device=device, dtype=torch.float32)
cx = pts[:, 0] * scale_x # [N]
cy = pts[:, 1] * scale_y # [N]
ys = torch.arange(latent_h, device=device, dtype=torch.float32).view(1, -1, 1)
xs = torch.arange(latent_w, device=device, dtype=torch.float32).view(1, 1, -1)
cxv = cx.view(-1, 1, 1)
cyv = cy.view(-1, 1, 1)
dist2 = (xs - cxv) ** 2 + (ys - cyv) ** 2
circles = (dist2 <= r_latent * r_latent).to(torch.float32) # [N, H, W]
mask = circles.amax(dim=0) # union via max
return mask.unsqueeze(0) # [1, H, W]
def _parse_rect_points(raw: str) -> tuple[float, float, float, float] | None:
"""Parse the JS-set rect_points widget — JSON list of one or more
[x1, y1, x2, y2] entries in image-pixel coords. We use only the
LAST rectangle (the most recent drag); earlier entries are kept
in the widget history for the undo stack to consume but don't
affect mask building. Returns None for empty / malformed input.
"""
raw = (raw or "").strip()
if not raw:
return None
try:
data = json.loads(raw)
except Exception:
return None
if not isinstance(data, list) or not data:
return None
last = data[-1]
if not isinstance(last, (list, tuple)) or len(last) < 4:
return None
try:
return (float(last[0]), float(last[1]), float(last[2]), float(last[3]))
except (TypeError, ValueError):
return None
def _rect_mask_latent(
latent_h: int,
latent_w: int,
rect_pixel: tuple[float, float, float, float],
scale_x: float,
scale_y: float,
device: torch.device,
) -> torch.Tensor:
"""Build a [1, H, W] filled-rectangle mask in latent space.
rect_pixel is (x1, y1, x2, y2) in image-pixel coords; corners
aren't required to be ordered (the user may drag in any
direction). Output is clamped to the latent bounds.
"""
x1, y1, x2, y2 = rect_pixel
x_lo_p, x_hi_p = min(x1, x2), max(x1, x2)
y_lo_p, y_hi_p = min(y1, y2), max(y1, y2)
xlat_lo = max(0, min(latent_w, int(round(x_lo_p * scale_x))))
xlat_hi = max(0, min(latent_w, int(round(x_hi_p * scale_x))))
ylat_lo = max(0, min(latent_h, int(round(y_lo_p * scale_y))))
ylat_hi = max(0, min(latent_h, int(round(y_hi_p * scale_y))))
mask = torch.zeros((1, latent_h, latent_w), device=device, dtype=torch.float32)
if xlat_hi > xlat_lo and ylat_hi > ylat_lo:
mask[0, ylat_lo:ylat_hi, xlat_lo:xlat_hi] = 1.0
return mask
def _parse_seg_polygons(raw: str):
"""Parse the seg_polygon widget — a JSON list of polygons, each a flat
[x0,y0,x1,y1,...] coord list in image-pixel space (a SAM 3 / YOLO
detection's silhouette). Returns the list, or None if empty/invalid."""
if not raw or not str(raw).strip():
return None
try:
polys = json.loads(raw)
except (ValueError, TypeError):
return None
if not isinstance(polys, list) or not polys:
return None
return polys
def _polygons_mask_latent(
latent_h: int,
latent_w: int,
polygons_pixel,
scale_x: float,
scale_y: float,
device: torch.device,
) -> torch.Tensor:
"""Rasterise one or more silhouette polygons (image-pixel coords) into
a filled [1, H, W] latent-space mask (their union)."""
import numpy as np
from PIL import Image, ImageDraw
img = Image.new("L", (latent_w, latent_h), 0)
draw = ImageDraw.Draw(img)
for poly in (polygons_pixel or []):
if not poly or len(poly) < 6:
continue
pts = [
(float(poly[i]) * scale_x, float(poly[i + 1]) * scale_y)
for i in range(0, len(poly) - 1, 2)
]
draw.polygon(pts, fill=255)
arr = np.array(img, dtype=np.float32) / 255.0
return torch.from_numpy(arr)[None, ...].to(device)
def _raster_mask_latent(latent_h, latent_w, png_b64, device):
"""Decode a base64-PNG touch-up mask (image-pixel resolution, white =
masked) into a [1, H, W] latent-space mask. The Detect Shift/Alt brush
produces this — a raster handles brushed holes / unions that a polygon
silhouette can't. Resized straight to the latent grid (no scale args)."""
import base64
import io
import numpy as np
from PIL import Image
raw = base64.b64decode(png_b64)
img = Image.open(io.BytesIO(raw)).convert("L")
img = img.resize((latent_w, latent_h), Image.BILINEAR)
arr = np.array(img, dtype=np.float32) / 255.0
return torch.from_numpy(arr)[None, ...].to(device)
def _mask_bbox_latent(mask: torch.Tensor) -> tuple[int, int, int, int] | None:
"""Tight latent-space bbox of non-zero mask values. Returns
(y_min, y_max, x_min, x_max) or None if the mask is empty.
Threshold of 0.01 includes the feathered edge — bbox covers the
full soft-edge region not just the binary interior.
"""
m = mask[0] if mask.dim() == 3 else mask
nz = m > 0.01
if not nz.any():
return None
rows = nz.any(dim=1)
cols = nz.any(dim=0)
ridx = rows.nonzero(as_tuple=False).squeeze(-1)
cidx = cols.nonzero(as_tuple=False).squeeze(-1)
return (
int(ridx[0].item()),
int(ridx[-1].item()) + 1,
int(cidx[0].item()),
int(cidx[-1].item()) + 1,
)
def _fine_upscale_factor(
bbox_w_latent: int,
bbox_h_latent: int,
scale_x: float,
scale_y: float,
target_mp: float,
max_linear: float,
) -> float:
"""Linear scale factor to apply to the cropped latent so that the
crop is processed at ≥ target_mp (in image-pixel-equivalent terms),
clamped to max_linear. Returns 1.0 when the crop already meets
target — no upscale needed."""
if scale_x <= 0 or scale_y <= 0:
return 1.0
bbox_w_pix = bbox_w_latent / scale_x
bbox_h_pix = bbox_h_latent / scale_y
current_mp = bbox_w_pix * bbox_h_pix / 1_000_000.0
if current_mp <= 0 or current_mp >= target_mp:
return 1.0
needed = math.sqrt(target_mp / current_mp)
return min(needed, max_linear)
def _resize_latent(t: torch.Tensor, target_h: int, target_w: int, method: str) -> torch.Tensor:
"""Resize the spatial dims of a latent or mask tensor using one of
ComfyUI's standard latent-resize methods. Accepts [C,H,W], [B,C,H,W],
or [1,H,W] (mask). Returns the same rank as input.
`method` is one of _FINE_UPSCALE_RESIZE_METHODS. Routes through
comfy.utils.common_upscale so bislerp + lanczos custom paths work."""
method = method if method in _FINE_UPSCALE_RESIZE_METHODS else "nearest-exact"
if t.dim() == 3:
t4 = t.unsqueeze(0)
out = comfy.utils.common_upscale(t4, target_w, target_h, method, "disabled")
return out.squeeze(0)
if t.dim() == 4:
return comfy.utils.common_upscale(t, target_w, target_h, method, "disabled")
raise ValueError(f"_resize_latent: unexpected ndim {t.dim()}")
# ----- VAE boundary -----
# Every latent->pixel and pixel->latent conversion in the node routes
# through these two functions. Centralising them means model-family
# quirks live in ONE place instead of at ~7 scattered call sites. The
# motivating case is Qwen Image Edit / Wan-derived VAEs, whose latents
# carry an extra temporal axis ([B, C, T, H, W] with T=1) that the rest
# of the node — and PIL's previewer — expect to be absent ([B, C, H, W]).
# Both helpers are thin pass-throughs today; this is the seam where that
# 5D normalisation will land without disturbing any caller.
def _vae_decode(vae, latent: torch.Tensor) -> torch.Tensor:
"""Decode a latent to pixels. Single decode chokepoint — see the
VAE-boundary note above. Always returns a 4D image batch
(B, H, W, C) float in [0, 1].
Temporal/video VAEs (Qwen Image Edit, Wan) keep a frame axis: their
latents are 5D ([B, C, T, H, W]) and `vae.decode` accordingly returns
a 5D frame stack ([B, T, H, W, C] — ComfyUI moves channels last). The
rest of the node, and ComfyUI's PreviewImage/PIL path, only understand
4D image batches, so fold the frame axis into the batch dim. For image
editing T is 1, so this is just dropping the singleton frame axis; if a
future model ever produces T>1 the frames surface as extra batch items
rather than crashing. The latent is passed through to `vae.decode`
untouched — the video VAE wants its native 5D input — we only normalise
the *pixels* it returns."""
image = vae.decode(latent)
if image.ndim == 5:
b, t, h, w, c = image.shape
image = image.reshape(b * t, h, w, c)
return image
def _vae_encode(vae, pixels: torch.Tensor) -> torch.Tensor:
"""Encode pixels to latent samples. Single encode chokepoint —
counterpart to _vae_decode. See the VAE-boundary note above.
Deliberately returns the VAE's native latent shape WITHOUT collapsing
it: a temporal/video VAE (Qwen, Wan) returns a 5D latent
([B, C, T, H, W]) and the sampler + model require that 5D shape to flow
through unchanged (comfy.sample.sample is ndim-agnostic and prepare_noise
matches the latent's shape exactly). Squeezing the frame axis here would
break Qwen sampling — do not add a squeeze."""
return vae.encode(pixels)
def _refine_with_fine_upscaling(
*,
model,
vae,
current: torch.Tensor, # [B, C, H_lat, W_lat] cached full-res latent
current_pixels: torch.Tensor | None, # [B, H_pix, W_pix, C] cached full-res pixels to avoid redundant VAE decode
mask: torch.Tensor, # [1, H_lat, W_lat] feathered mask, latent res
scale_x: float,
scale_y: float,
target_mp: float,
max_linear: float,
resize_method: str,
context_pad_pixel: int,
inpainting_mode: str,
reference_strength: float = 0.0, # Refine + Ref value: anchor on the
# (pre-refine) crop for this fraction
# of the schedule (_apply_reference) —
# WITHOUT Smart Inpaint's latent-zeroing
seed: int = 0,
steps: int,
cfg: float,
sampler_name: str,
scheduler: str,
positive,
negative,
denoise: float,
callback,
disable_pbar: bool,
# #8 custom-sampler trio — None = use the default comfy.sample.sample
# path; all three set = dispatch via _do_sample to the guider path.
ov_guider=None,
ov_sampler=None,
ov_sigmas=None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Pixel-space crop + upscale + VAE encode + refine + VAE decode +
downscale + composite + VAE encode. The latent-space crop+upscale
approach smears bilinearly-interpolated latents into a low-freq
starting state that the model can't recover detail from. Going
through pixel space (where there's an image-upscale toolkit that's
been tuned for natural images) and re-encoding gives the model a
"natural" latent at the higher resolution to denoise from.
Returns a tuple of (new_latent, new_pixels). Returns (current, current_pixels)
if the mask bbox is empty / degenerate.
"""
bbox = _mask_bbox_latent(mask)
if bbox is None:
return current, current_pixels
y0_tight, y1_tight, x0_tight, x1_tight = bbox
# Apply context padding: grow the bbox outward by context_pad_pixel
# in every direction (clamped to the latent boundaries). This is
# the area the model SEES during refine. The painted-shape mask
# stays unchanged — areas inside the padded bbox but outside the
# painted shape have mask=0 in the cropped tensor, so the noise-
# injection inpaint preserves them as context (the model uses them
# to inform what to draw inside the mask, but doesn't overwrite
# them). All downstream code uses the PADDED bbox.
H_lat = current.shape[-2]
W_lat = current.shape[-1]
pad_lat_y = max(0, round(context_pad_pixel * scale_y))
pad_lat_x = max(0, round(context_pad_pixel * scale_x))
y0 = max(0, y0_tight - pad_lat_y)
y1 = min(H_lat, y1_tight + pad_lat_y)
x0 = max(0, x0_tight - pad_lat_x)
x1 = min(W_lat, x1_tight + pad_lat_x)
bbox_h_lat = y1 - y0
bbox_w_lat = x1 - x0
if bbox_h_lat <= 0 or bbox_w_lat <= 0:
return current, current_pixels
scale = _fine_upscale_factor(bbox_w_lat, bbox_h_lat, scale_x, scale_y, target_mp, max_linear)
if scale <= 1.0 and inpainting_mode != "Smart Inpaint":
# Refine with no upscale needed — fall back to the standard latent-space
# noise-injection inpaint. Avoids unnecessary VAE round-trips when the
# painted region already meets the MP target.
#
# Smart Inpaint must NOT take this shortcut. It needs the crop +
# reference_latents + masked-zero treatment below regardless of rect
# size; skipping it made a LARGE rectangle (already at/above the MP
# target, so scale<=1.0 — roughly >1024px on FLUX 2) degrade to a
# whole-latent edit with NO crop reference, so the model worked on the
# whole image instead of the selected rect.
print(f"[Angelo fine-upscale] scale=1.0 — using latent-space path (no VAE round-trip)")
if reference_strength > 0.0:
# Refine + Reference on the no-upscale path: anchor on the whole
# current image (there's no crop on this path).
positive = _apply_reference(positive, current.clone(), reference_strength)
noise = comfy.sample.prepare_noise(current, seed, None)
new_latent = _do_sample(
guider=ov_guider, sampler=ov_sampler, sigmas=ov_sigmas,
model=model, noise=noise,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=positive, negative=negative,
source_latent=current,
denoise=denoise,
noise_mask=mask,
callback=callback,
disable_pbar=disable_pbar,
seed=seed,
)
# Return None for pixels because the latent was modified directly;
# this forces a fresh VAE decode for the preview in the main run() method.
return new_latent, None
# Smart Inpaint with a large rectangle still crops + references the selected
# region; it just doesn't upscale (and must never downscale) — clamp the
# factor to identity so the crop is taken at native resolution.
if inpainting_mode == "Smart Inpaint":
scale = max(1.0, scale)
# ----- VAE decode the full cached latent → cached pixels -----
# Optimization: Reuse cached pixels if available to prevent VAE degradation
# (loss of high-frequency details) across multiple consecutive edits.
if current_pixels is not None:
cached_pixels = current_pixels
else:
cached_pixels = _vae_decode(vae, current) # (B, H_pix, W_pix, C) float [0,1]
H_pix = cached_pixels.shape[1]
W_pix = cached_pixels.shape[2]
# Pixel-per-latent ratio per axis (16 for FLUX 2, 8 for SDXL/SD1.5).
# round(), not floor-divide: a non-integer true ratio (exotic VAEs)
# floor-divided gives e.g. 15 for 15.8, drifting the pixel-space bbox
# ~1px against the latent bbox and leaving a seam in the composite.
# (#28, from @KursatAs.)
px_per_lat_y = max(1, round(H_pix / current.shape[-2]))
px_per_lat_x = max(1, round(W_pix / current.shape[-1]))
# Pixel-space bbox derived from the latent-space bbox.
y0_p = y0 * px_per_lat_y
y1_p = y1 * px_per_lat_y
x0_p = x0 * px_per_lat_x
x1_p = x1 * px_per_lat_x
bbox_h_p = y1_p - y0_p
bbox_w_p = x1_p - x0_p
# Upscaled target dims in pixel space. Snap to multiples of the
# VAE downscale (16 for FLUX 2) so the subsequent VAE encode
# produces a clean integer-dim latent.
vae_snap = max(px_per_lat_y, px_per_lat_x)
target_h_p = max(vae_snap, math.ceil(bbox_h_p * scale / vae_snap) * vae_snap)
target_w_p = max(vae_snap, math.ceil(bbox_w_p * scale / vae_snap) * vae_snap)
print(f"[Angelo fine-upscale] bbox_lat=(h={bbox_h_lat}, w={bbox_w_lat}) "
f"bbox_px=(h={bbox_h_p}, w={bbox_w_p}) scale={scale:.2f} "
f"target_px=(h={target_h_p}, w={target_w_p}) "
f"resize={resize_method} max_linear={max_linear} "
f"vae_ratio=(x={px_per_lat_x}, y={px_per_lat_y})")
# ----- Crop pixel image + upscale in pixel space -----
pixel_crop = cached_pixels[:, y0_p:y1_p, x0_p:x1_p, :] # (B, h, w, C)
# common_upscale expects (B, C, H, W) — permute, upscale, permute back.
pixel_crop_chw = pixel_crop.movedim(-1, 1)
pixel_crop_up_chw = comfy.utils.common_upscale(
pixel_crop_chw, target_w_p, target_h_p, resize_method, "disabled",
)
pixel_crop_up = pixel_crop_up_chw.movedim(1, -1) # back to (B, H, W, C)
# ----- VAE encode the upscaled pixel crop → latent at high res -----
latent_up = _vae_encode(vae, pixel_crop_up)
target_h_lat = latent_up.shape[-2]
target_w_lat = latent_up.shape[-1]
# ----- Build mask at the upscaled latent resolution -----
# Mask resizing always uses bilinear regardless of the user's choice.
# The user's resize_method is for the IMAGE content upscale (where
# lanczos / bicubic / etc. have real quality differences). The mask
# is a 1-channel feathered alpha where we just want smooth values;
# lanczos's grayscale-branch returns a transposed 3D tensor (PIL
# quirk) and bislerp's spherical-vector math is semantically wrong
# on a single channel.
mask_crop = mask[..., y0:y1, x0:x1].contiguous()
mask_crop_up = _resize_latent(mask_crop, target_h_lat, target_w_lat, "bilinear").clamp(0.0, 1.0)
# ===== Smart Inpaint pre-processing on the upscaled patch =====
# Klein 9B's edit branch only activates when reference_latents is
# present on the conditioning. We then zero the masked area so the
# sampler regenerates that region from full noise at sigma_max
# (the denoise=1.0 lock makes this clean: every pixel in the
# painted rect is brand-new content, with the surrounding context
# band restored each step by the noise_mask compositing). The
# reference uses the PRE-ZERO upscaled patch so Klein still sees
# what was there before we blanked it.
# POSITIVE ONLY — putting reference_latents on negative would tell
# CFG>1 samplers to steer AWAY from the reference scene. Non-edit
# models ignore the field, so this is harmless on any checkpoint.
#
# append=False (REPLACE, not append): the reference must be ONLY this
# upscaled crop. When the Area Prompt is empty, refine_positive falls back
# to the node's `positive` input, which in a Klein edit workflow already
# carries reference_latents = the WHOLE source image (from an upstream
# ReferenceLatent node). append=True stacked the crop onto that whole-image
# reference, and the whole-image one dominated — so the patch reproduced
# the entire original scene instead of editing the selected region.
# Replacing guarantees Klein sees the crop and nothing else.
# Mask used for the SAMPLING step (both the masked-zero and the noise_mask).
# Defaults to the feathered mask; hardened to binary for Smart Inpaint on
# 5D temporal latents — see the note in the Smart Inpaint block below.
sample_mask = mask_crop_up
if inpainting_mode == "Smart Inpaint":
reference_latent = latent_up.clone()
positive = node_helpers.conditioning_set_values(
positive, {"reference_latents": [reference_latent]}, append=False,
)
# Feather goes in the COMPOSITE, not the denoise mask — for Qwen/Wan.
#
# FLUX/Klein (4D, ~zero-mean latents): a soft mask works directly as
# both the masked-zero and the noise_mask; their inpaint blend handles a
# soft boundary cleanly, so keep the feathered mask.
#
# Qwen Image Edit / Wan (5D temporal latents): sample with a HARD
# (binary) mask. These models have no clean soft-mask-during-denoise
# behaviour — the community-standard Qwen-edit inpaint recipe is "blank
# the region, regenerate, then composite back with a feathered blend",
# NOT feathering the denoise mask. A soft noise_mask at denoise=1.0 on a
# non-zero-mean latent space distorts exactly the feather band (the
# symptom: artifacts only where the feathering happens). The smooth
# visible edge is still produced downstream by the feathered PIXEL
# composite + final latent blend (both use the full-res feathered
# `mask`), so the feather is preserved without the sampling artifacts.
if latent_up.ndim == 5:
sample_mask = (mask_crop_up >= 0.5).to(mask_crop_up.dtype)
latent_up = (1.0 - sample_mask.unsqueeze(0)) * latent_up
elif reference_strength > 0.0:
# Refine + Reference value: the upscaled (pre-refine) crop anchors
# identity/content through the edit branch for the first
# reference_strength fraction of the schedule, so a high denoise
# can fully re-render texture without losing the subject — the
# photo-restoration recipe. Unlike Smart Inpaint, the latent is NOT
# zeroed: the existing content remains the starting state.
positive = _apply_reference(positive, latent_up.clone(), reference_strength)
# ----- Refine via noise-injection inpaint on the upscaled latent -----
noise = comfy.sample.prepare_noise(latent_up, seed, None)
refined_latent_up = _do_sample(
guider=ov_guider, sampler=ov_sampler, sigmas=ov_sigmas,
model=model, noise=noise,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=positive, negative=negative,
source_latent=latent_up,
denoise=denoise,
noise_mask=sample_mask,
callback=callback,
disable_pbar=disable_pbar,
seed=seed,
)
# ----- VAE decode refined latent → high-res pixel patch -----
refined_pixel_up = _vae_decode(vae, refined_latent_up) # (B, target_h_p, target_w_p, C)
# ----- Downscale refined patch back to original bbox pixel size -----
refined_pixel_up_chw = refined_pixel_up.movedim(-1, 1)
refined_pixel_chw = comfy.utils.common_upscale(
refined_pixel_up_chw, bbox_w_p, bbox_h_p, resize_method, "disabled",
)
refined_pixel = refined_pixel_chw.movedim(1, -1) # (B, bbox_h_p, bbox_w_p, C)
# ----- Composite refined patch into the cached pixel image -----
# Build a pixel-space alpha by resizing the latent feathered mask to
# full pixel resolution, cropping to the bbox. Always bilinear for
# the same reasons as the mask upscale above — lanczos's grayscale
# path is broken, bislerp doesn't apply to 1-channel.
mask_4d = mask.unsqueeze(0) # [1, 1, H_lat, W_lat]
pixel_mask = comfy.utils.common_upscale(
mask_4d, W_pix, H_pix, "bilinear", "disabled",
).clamp(0.0, 1.0) # [1, 1, H_pix, W_pix]
pixel_alpha_crop = pixel_mask[0, 0, y0_p:y1_p, x0_p:x1_p] # [bbox_h_p, bbox_w_p]
pixel_alpha_crop = pixel_alpha_crop.unsqueeze(0).unsqueeze(-1) # [1, h, w, 1]
new_pixels = cached_pixels.clone()
pixel_orig_crop = cached_pixels[:, y0_p:y1_p, x0_p:x1_p, :]
composited = refined_pixel * pixel_alpha_crop + pixel_orig_crop * (1.0 - pixel_alpha_crop)
new_pixels[:, y0_p:y1_p, x0_p:x1_p, :] = composited
# ----- VAE encode the composited full image → encoded latent -----
encoded_latent = _vae_encode(vae, new_pixels)
# ----- Blend in LATENT space using the feathered mask as alpha -----
# The VAE encode is lossy, so naively returning encoded_latent would
# mean the *unaltered* regions of the image drift slightly with every
# Fine Upscale click. Avoidable: keep the original cached latent
# outside the mask, take the encoded latent inside the mask. Mask is
# already feathered, so the transition is smooth. Now unaltered
# regions stay bit-exact across successive clicks; only the masked
# area accumulates any VAE-roundtrip cost (and it gets a fresh
# refine each click anyway, so any drift there is overwritten).
alpha_lat = mask.unsqueeze(0) # [1, 1, H_lat, W_lat]
new_latent = encoded_latent * alpha_lat + current * (1.0 - alpha_lat)
return new_latent, new_pixels
# Quick Photo Refine: the one-click restoration recipe — a MAGIC BUTTON
# with its own fixed settings, deliberately independent of every toolbar
# box. The whole canvas is re-rendered with this prompt, anchored to the
# current image via reference_latents — identity from the reference,
# texture from the re-render. Tested values; tuned here in ONE place.
# The winning prompt is an INSTRUCTION — edit models are instruction-
# trained, so speaking their vocabulary beat every descriptive/
# constraint phrasing tried ("high quality photo", keep-the-colours
# variants — see git history). At ref 1.0 the blend collapses to a
# single cond, so this recipe also costs nothing extra.
#
# Per-model phrasing: Qwen-Image-Edit responds better to a fuller,
# gentler instruction than Klein does. Selected at run time by latent
# dimensionality — Qwen/Wan-family latents are 5D [B,C,T,H,W], the
# same load-bearing check the inpaint hard-mask logic uses.
_QUICK_REFINE_PROMPT = "restore the photo"
_QUICK_REFINE_PROMPT_QWEN = ("lightly restore this old photo, remove dust and scratches, "
"improve sharpness and contrast, preserve original feel")
_QUICK_REFINE_DENOISE = 1.0
_QUICK_REFINE_REF = 1.0
# Large Image Refine (the ▦ button): the canvas is divided into ~_LIR_BOX_MP
# boxes that EXACT-TILE it (no overlap), each refined through the Xtra-Fine
# pipeline (bit-exact latent compositing — no pixel feathering anywhere) in
# CHESS-PATTERN order: one colour first, then the other, so second-pass
# boxes refine with already-refined neighbours visible in their context
# pad and match them. Hard box masks (feather 0); seam defences are the
# shared source at this denoise, the ctx pad's cross-border visibility,
# and the chess ordering.
_LIR_PROMPT = "restore the image. make it clear and sharp."
_LIR_REF = 0.4
_LIR_DENOISE = 0.55
_LIR_CTX_PAD = 176
_LIR_BOX_MP = 1.0 # target box area — the tweakable knob
# Tiled restore engine (2× Restore Upscale + big-canvas Quick Refine).
# Working tile size + overlap in PIXELS: tiles are sampled at ~1MP no
# matter how large the canvas is, so the latent fed to the model never
# outgrows the resolution it renders well at. Quick Photo Refine
# auto-routes through the tiled engine above the MP threshold for the
# same reason.
_TILE_PX = 1400
_TILE_OVERLAP_PX = 128
def _tile_min_overlap_px(canvas_long_edge: int) -> int:
"""Canvas-proportional MINIMUM tile overlap. Bigger canvases mean each
tile sees a smaller slice of the scene, so neighbours need more shared
context to agree. Linear in the long edge, calibrated on tested
points: 3872px -> 172, 7744px -> 256; clamps to the 128px floor below
~1840px."""
return max(_TILE_OVERLAP_PX, 88 + round(canvas_long_edge / 46.0))
_QUICK_REFINE_TILE_THRESHOLD_MP = 1.6
def _apply_reference(positive, ref_latent: torch.Tensor, strength: float):
"""Attach ref_latent as reference_latents at a TRUE fractional strength
via dual-conditioning prediction blending.
There is no native scalar weight on reference_latents — the reference
becomes extra image tokens, take-it-or-leave-it. But ComfyUI's cond
`strength` field is a RELATIVE weight: each cond's prediction is
multiplied by its strength, accumulated, then normalised by the total.
So two full-coverage conds — [with-reference, strength s] + [plain,
strength 1-s] — yield exactly s·pred_ref + (1-s)·pred_plain at every
step: a genuine prediction-space lerp between anchored and free.
1.0 → one cond, fully anchored (no extra cost)
0 < s <1 → the dual-cond blend above. COSTS A SECOND positive-side
model evaluation per step (like CFG's negative does) —
the price of a real interpolation.
0.0 → untouched (no reference, no extra cost)
(A timestep-range variant — strength as DURATION of anchoring — was
tried first and rejected by testing; see git history.)
"""
s = max(0.0, min(1.0, float(strength)))
if s <= 0.0:
return positive
with_ref = node_helpers.conditioning_set_values(
positive, {"reference_latents": [ref_latent]}, append=False,
)
if s >= 1.0:
return with_ref
head = node_helpers.conditioning_set_values(with_ref, {"strength": s})
tail = node_helpers.conditioning_set_values(positive, {"strength": 1.0 - s})
return head + tail
def _tile_positions(total: int, tile: int, overlap: int) -> list[int]:
"""Start offsets covering [0, total) with `tile`-sized windows,
UNIFORMLY distributed so every seam gets the same overlap (>= the
requested minimum). The old flush-to-end layout produced wildly
inconsistent seams (128px here, 900px there) — uniform spacing means
every boundary shares at least the minimum context, evenly."""
if total <= tile:
return [0]
stride = max(1, tile - max(0, overlap))
n = max(2, math.ceil((total - overlap) / stride))
step = (total - tile) / (n - 1)
return [round(i * step) for i in range(n)]
def _tiled_restore_pass(
*,
model,
vae,
current: torch.Tensor,
current_pixels: torch.Tensor | None,
scale: float,
positive_base,
negative,
seed: int,
steps: int,
cfg: float,
sampler_name: str,
scheduler: str,
callback,
disable_pbar: bool,
ov_guider=None,
ov_sampler=None,
ov_sigmas=None,
tile_ref: float = _QUICK_REFINE_REF,
tile_denoise: float = _QUICK_REFINE_DENOISE,
) -> tuple[torch.Tensor, torch.Tensor, int]:
"""The tiled restore engine: the Quick Photo Refine recipe run over
overlapping ~1MP tiles, so the model never samples a latent bigger
than its happy place — regardless of canvas size.
scale=2.0 → 2× Restore Upscale: decode, lanczos-upscale the PIXELS
(latent upscaling smears; pixel space is the Xtra-Fine lesson),
re-encode, then restore tile by tile. scale=1.0 → in-place tiled
restore (how Quick Photo Refine handles big canvases).
Each tile is sampled at denoise 1.0 anchored to ITS OWN pre-refine
latent via reference_latents. The anchor is what makes tiled
full-denoise viable: tiles can't hallucinate new content (the
classic tiled-upscaler failure), and adjacent tiles agree because
they're anchored to the same underlying image — the reference does
the job tile-ControlNets do in other pipelines.
Tiles composite in pixel space under a separable ramp weight
(linear fade over the overlap), normalised by the accumulated
weight, so seams blend and un-overlapped edges keep weight 1.
Returns (final_latent, final_pixels, n_tiles).
"""
pixels = current_pixels if current_pixels is not None else _vae_decode(vae, current)
B, H, W, C = pixels.shape
if scale != 1.0:
new_W = max(16, int(round(W * scale / 16.0)) * 16)
new_H = max(16, int(round(H * scale / 16.0)) * 16)
chw = pixels.movedim(-1, 1)
chw = comfy.utils.common_upscale(chw, new_W, new_H, "lanczos", "disabled")
big_pixels = chw.movedim(1, -1).contiguous()
else:
big_pixels = pixels
new_H, new_W = big_pixels.shape[1], big_pixels.shape[2]
big_latent = _vae_encode(vae, big_pixels)
nlh, nlw = big_latent.shape[-2], big_latent.shape[-1]
ppl_x = max(1, round(new_W / nlw))
ppl_y = max(1, round(new_H / nlh))
tile_lx = max(8, _TILE_PX // ppl_x)
tile_ly = max(8, _TILE_PX // ppl_y)
min_ov_px = _tile_min_overlap_px(max(new_W, new_H))
ov_lx = max(1, min_ov_px // ppl_x)
ov_ly = max(1, min_ov_px // ppl_y)
xs = _tile_positions(nlw, tile_lx, ov_lx)
ys = _tile_positions(nlh, tile_ly, ov_ly)
n_tiles = len(xs) * len(ys)
print(f"[Angelo tiled-restore] scale={scale} canvas={new_W}x{new_H} "
f"grid={len(xs)}x{len(ys)} ({n_tiles} tiles of ~{_TILE_PX}px, "
f"min overlap {min_ov_px}px, uniform spacing)")
out = torch.zeros_like(big_pixels)
wsum = torch.zeros((1, new_H, new_W, 1), dtype=big_pixels.dtype, device=big_pixels.device)
# ONE global noise field, cropped per tile. This is the anti-ghosting
# measure: neighbouring tiles see IDENTICAL init noise (and identical
# reference content) in their shared overlap, so at denoise 1.0 their
# trajectories converge on essentially the same fine detail there —
# the blend averages two near-identical renderings instead of two
# independently-imagined ones (decorrelated per-tile seeds produced
# "double eyelash" ghosting at seams). Deterministic samplers (euler,
# the default) get the full benefit; ancestral samplers re-noise per
# step in tile-local coordinates, so some residual seam softness is
# expected there.
global_noise = comfy.sample.prepare_noise(big_latent, seed, None)
def _ramp(length_px: int, ov_px: int, device):
r = torch.ones(length_px, dtype=torch.float32, device=device)
n = min(ov_px, length_px)
fade = (torch.arange(n, dtype=torch.float32, device=device) + 1.0) / (n + 1.0)
fade = fade * fade * (3.0 - 2.0 * fade) # smoothstep — softer crossover
r[:n] = torch.minimum(r[:n], fade)
r[-n:] = torch.minimum(r[-n:], fade.flip(0))
return r
idx = 0
for y0 in ys:
for x0 in xs:
tw = min(tile_lx, nlw - x0)
th = min(tile_ly, nlh - y0)
tile_lat = big_latent[..., y0:y0 + th, x0:x0 + tw].clone().contiguous()
tile_pos = _apply_reference(positive_base, tile_lat.clone(), tile_ref)
tile_seed = int(seed)
noise = global_noise[..., y0:y0 + th, x0:x0 + tw].clone().contiguous()
refined = _do_sample(
guider=ov_guider, sampler=ov_sampler, sigmas=ov_sigmas,
model=model, noise=noise,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=tile_pos, negative=negative,
source_latent=tile_lat,
denoise=tile_denoise,
noise_mask=None,
callback=callback,
disable_pbar=disable_pbar,
seed=tile_seed,
)
tile_px = _vae_decode(vae, refined).to(big_pixels.device, big_pixels.dtype)
py0, px0 = y0 * ppl_y, x0 * ppl_x
pth, ptw = tile_px.shape[1], tile_px.shape[2]
wy = _ramp(pth, min_ov_px, big_pixels.device)
wx = _ramp(ptw, min_ov_px, big_pixels.device)
w = (wy.view(1, -1, 1, 1) * wx.view(1, 1, -1, 1)).to(big_pixels.dtype)
out[:, py0:py0 + pth, px0:px0 + ptw, :] += tile_px * w
wsum[:, py0:py0 + pth, px0:px0 + ptw, :] += w
idx += 1
final_pixels = out / wsum.clamp(min=1e-6)
final_latent = _vae_encode(vae, final_pixels)
return final_latent, final_pixels, n_tiles
# Direction-aware instruction prepended to the outpaint conditioning when a
# CLIP is wired (same pattern as Smart Guided's location prefixes). Edit
# models (Klein / Qwen) follow explicit continue-don't-repeat instructions
# well — this is the main lever against the classic outpainting failure of
# duplicating the scene's subject into the new strip.
_OUTPAINT_INSTRUCTIONS = {
"left": "Extend the image to the left, continuing the scene and background naturally. Do not repeat or add new subjects. ",
"right": "Extend the image to the right, continuing the scene and background naturally. Do not repeat or add new subjects. ",
"up": "Extend the image upward, continuing the scene and background naturally. Do not repeat or add new subjects. ",
"down": "Extend the image downward, continuing the scene and background naturally. Do not repeat or add new subjects. ",
"all": "Extend the image outward on all sides, continuing the scene and background naturally. Do not repeat or add new subjects. ",
}
def _outpaint_prepare(
*,
vae,
current: torch.Tensor,
current_pixels: torch.Tensor | None,
direction: str,
amount_px: int,
overlap_px: int,
protect_json: str = "",
) -> tuple[torch.Tensor, torch.Tensor, int, torch.Tensor]:
"""Build the init latent + mask for one outpaint pass.
Pads the decoded canvas in `direction` by `amount_px` (snapped to a
/16 multiple so both 8x and 16x VAEs land on integer latent cells),
fills the new strip with edge-replicated pixels (a colour-continuity
hint that costs nothing — at denoise 1.0 the masked region starts from
pure noise on flow models anyway), VAE-encodes the padded canvas, then
PASTES the original latent back over the old region so the existing
image stays bit-exact in latent terms (no VAE round-trip drift — the
same trick as Xtra-Fine's final blend).
The mask covers the new strip at 1.0 plus a feathered band reaching
`overlap_px` into the old image, so the seam is redrawn rather than
butted. The pure-new region is re-hardened to 1.0 after the blur.
Every extension spans the FULL current edge, which is what makes
corners a non-problem: extend right then down, and the bottom strip
spans the new wider canvas — the corner is generated by whichever
direction runs second, with both neighbours as context.
`protect_json` is the protect-brush payload — a JSON list of
[x, y, r] circles in OLD-image pixel coords. Protected pixels are
subtracted from the overlap band (clamped so the pure-new strip can
never be protected), keeping e.g. a car at the frame edge frozen
while the rest of the band still blends generously.
Also returns the REFERENCE latent for edit models: the edge-adjacent
band (~512px deep) for directional extensions — showing the model
the texture/lighting to continue WITHOUT re-showing it the scene's
subject (whole-image references invite the edit branch to reproduce
the subject into the strip — the Smart Inpaint lesson again). "all"
keeps the whole image (everything borders the new space).
Returns (latent_init, mask, snapped_amount_px, reference_latent).
"""
pixels = current_pixels if current_pixels is not None else _vae_decode(vae, current)
B, H, W, C = pixels.shape
amt = max(16, int(round(amount_px / 16.0)) * 16)
pl, pr, pt, pb = {
"left": (amt, 0, 0, 0),
"right": (0, amt, 0, 0),
"up": (0, 0, amt, 0),
"down": (0, 0, 0, amt),
"all": (amt, amt, amt, amt),
}.get(str(direction), (0, amt, 0, 0))
chw = pixels.movedim(-1, 1)
chw = torch.nn.functional.pad(chw, (pl, pr, pt, pb), mode="replicate")
new_pixels = chw.movedim(1, -1).contiguous()
new_H, new_W = new_pixels.shape[1], new_pixels.shape[2]
latent_init = _vae_encode(vae, new_pixels)
nlh, nlw = latent_init.shape[-2], latent_init.shape[-1]
px_per_lat_x = max(1, round(new_W / nlw))
px_per_lat_y = max(1, round(new_H / nlh))
lat_l = pl // px_per_lat_x
lat_t = pt // px_per_lat_y
lh, lw = current.shape[-2], current.shape[-1]
latent_init[..., lat_t:lat_t + lh, lat_l:lat_l + lw] = current.to(
device=latent_init.device, dtype=latent_init.dtype,
)
ov_lx = max(0, round(int(overlap_px) / px_per_lat_x))
ov_ly = max(0, round(int(overlap_px) / px_per_lat_y))
mask = torch.ones((1, nlh, nlw), dtype=torch.float32, device=latent_init.device)
y0 = lat_t + (ov_ly if pt > 0 else 0)
y1 = lat_t + lh - (ov_ly if pb > 0 else 0)
x0 = lat_l + (ov_lx if pl > 0 else 0)
x1 = lat_l + lw - (ov_lx if pr > 0 else 0)
if y1 > y0 and x1 > x0:
mask[..., y0:y1, x0:x1] = 0.0
# strip_only = 1.0 over ONLY the genuinely-new pixels. It's the floor
# the blur and the protect brush are both clamped against: nothing may
# reduce the strip below full regeneration.
strip_only = torch.ones_like(mask)
strip_only[..., lat_t:lat_t + lh, lat_l:lat_l + lw] = 0.0
if int(overlap_px) > 0:
sigma = max(0.5, (ov_lx + ov_ly) / 4.0)
mask = torch.maximum(_gaussian_blur_2d(mask, sigma), strip_only).clamp(0.0, 1.0)
# Protect brush: subtract painted circles from the overlap band so
# e.g. a car at the frame edge stays bit-exact frozen while the rest
# of the band still blends. Circles arrive in OLD-image pixel coords;
# offset by the pad, scale to latent cells, union, soften 1 cell.
circles = []
try:
for item in json.loads(protect_json or "[]"):
if isinstance(item, (list, tuple)) and len(item) >= 3:
circles.append((float(item[0]), float(item[1]), float(item[2])))
except (ValueError, TypeError):
circles = []
if circles:
pts = torch.tensor(circles, dtype=torch.float32, device=mask.device)
cx = (pts[:, 0] + pl) / px_per_lat_x
cy = (pts[:, 1] + pt) / px_per_lat_y
rl = pts[:, 2] / math.sqrt(px_per_lat_x * px_per_lat_y)
ys = torch.arange(nlh, device=mask.device, dtype=torch.float32).view(1, -1, 1)
xs = torch.arange(nlw, device=mask.device, dtype=torch.float32).view(1, 1, -1)
dist2 = (xs - cx.view(-1, 1, 1)) ** 2 + (ys - cy.view(-1, 1, 1)) ** 2
protect = (dist2 <= (rl.view(-1, 1, 1) ** 2)).to(torch.float32).amax(dim=0).unsqueeze(0)
protect = _gaussian_blur_2d(protect, 1.0).clamp(0.0, 1.0)
mask = torch.maximum(mask * (1.0 - protect), strip_only).clamp(0.0, 1.0)
# Reference for edit models: the edge-adjacent band (~512px deep) for
# a directional extension; the whole image for "all" (see docstring).
ref_latent = current
if direction in ("left", "right", "up", "down"):
d_x = max(8, round(512.0 / px_per_lat_x))
d_y = max(8, round(512.0 / px_per_lat_y))
if direction == "right":
ref_latent = current[..., :, max(0, lw - d_x):]
elif direction == "left":
ref_latent = current[..., :, :min(lw, d_x)]
elif direction == "down":
ref_latent = current[..., max(0, lh - d_y):, :]
else: # up
ref_latent = current[..., :min(lh, d_y), :]
ref_latent = ref_latent.clone().contiguous()
return latent_init, mask, amt, ref_latent
def _circle_mask_latent_direct(
latent_h: int,
latent_w: int,
cx_latent: float,
cy_latent: float,
r_latent: float,
device: torch.device,
) -> torch.Tensor:
"""Build a binary [1, latent_h, latent_w] mask with a filled circle
centred on (cx_latent, cy_latent) of radius `r_latent`, all in
latent-space coordinates. The caller is responsible for converting
pixel-space click coords to latent space using the correct per-axis
scale (image_dim_pixel → latent_dim) so this function stays VAE-
agnostic.
"""
ys = torch.arange(latent_h, device=device, dtype=torch.float32).view(-1, 1)
xs = torch.arange(latent_w, device=device, dtype=torch.float32).view(1, -1)
dist2 = (xs - cx_latent) ** 2 + (ys - cy_latent) ** 2
mask = (dist2 <= r_latent * r_latent).to(torch.float32)
return mask.unsqueeze(0) # [1, H, W]
def _gaussian_blur_2d(mask: torch.Tensor, sigma_latent: float) -> torch.Tensor:
"""Separable gaussian blur on a [B, H, W] or [H, W] mask tensor.
sigma_latent is in latent-space units.
"""
if sigma_latent <= 0:
return mask
# Kernel covers ~±3σ. Force odd size.
ksize = int(2 * math.ceil(3 * sigma_latent) + 1)
half = ksize // 2
x = torch.arange(ksize, device=mask.device, dtype=torch.float32) - half
k1d = torch.exp(-0.5 * (x / sigma_latent) ** 2)
k1d = k1d / k1d.sum()
# Reshape mask to [N, 1, H, W]
orig_ndim = mask.dim()
if orig_ndim == 2:
m = mask.unsqueeze(0).unsqueeze(0)
elif orig_ndim == 3:
m = mask.unsqueeze(1)
elif orig_ndim == 4:
m = mask
else:
raise ValueError(f"gaussian_blur_2d: unexpected mask ndim {orig_ndim}")
kh = k1d.view(1, 1, 1, ksize)
kv = k1d.view(1, 1, ksize, 1)
m = torch.nn.functional.pad(m, (half, half, 0, 0), mode="replicate")
m = torch.nn.functional.conv2d(m, kh)
m = torch.nn.functional.pad(m, (0, 0, half, half), mode="replicate")
m = torch.nn.functional.conv2d(m, kv)
if orig_ndim == 2:
return m.squeeze(0).squeeze(0)
if orig_ndim == 3:
return m.squeeze(1)
return m
def _encode_loaded_image(vae, ref_json: str, resize_mode: str, target_mp: float):
"""Load an image the user picked via the Load Image button and
VAE-encode it into a base latent.
`ref_json` is a JSON {name, subfolder, type} ref returned by
ComfyUI's /upload/image (falls back to treating the string as a bare
input-dir filename). resize_mode is "keep" (native res) or "mp"
(scaled to ~target_mp megapixels). In both cases dimensions are
rounded to a multiple of 16 so any supported VAE (8x or 16x) is
happy. Returns the latent samples tensor.
"""
import numpy as np
from PIL import Image, ImageOps
# Resolve the image reference to a path.
name, subfolder, type_ = ref_json, "", "input"
try:
ref = json.loads(ref_json)
name = ref.get("name") or ref.get("filename") or ""
subfolder = ref.get("subfolder", "") or ""
type_ = ref.get("type", "input") or "input"
except (ValueError, TypeError):
pass
if type_ == "output":
base_dir = folder_paths.get_output_directory()
elif type_ == "temp":
base_dir = folder_paths.get_temp_directory()
else:
base_dir = folder_paths.get_input_directory()
path = os.path.normpath(os.path.join(base_dir, subfolder, name))
if not path.startswith(os.path.normpath(base_dir)):
raise ValueError("Angelo: invalid loaded-image path")
if not os.path.exists(path):
raise ValueError(f"Angelo: loaded image not found: {name}")
img = Image.open(path)
img = ImageOps.exif_transpose(img).convert("RGB")
w, h = img.size
if resize_mode == "mp" and target_mp > 0:
cur_mp = (w * h) / 1.0e6
if cur_mp > 0:
s = math.sqrt(target_mp / cur_mp)
w = int(round(w * s))
h = int(round(h * s))
# Round down to a multiple of 16 (divisible by both 8x and 16x VAEs).
w = max(16, (w // 16) * 16)
h = max(16, (h // 16) * 16)
if (w, h) != img.size:
img = img.resize((w, h), Image.LANCZOS)
arr = np.array(img).astype(np.float32) / 255.0 # (H, W, 3)
pixels = torch.from_numpy(arr)[None, ...] # (1, H, W, 3)
samples = _vae_encode(vae, pixels[:, :, :, :3])
return samples
def _decode_to_preview(vae, latent_samples: torch.Tensor):
"""Decode the latent and save to the temp directory in the same
format PreviewImage uses, so we can return the same ui dict shape.
Returns (image_tensor, list_of_image_refs). Each image ref is a
{filename, subfolder, type} dict.
"""
image = _vae_decode(vae, latent_samples) # (B, H, W, C) float in [0, 1]
# Reuse PreviewImage's save logic via a transient instance.
previewer = comfy_nodes.PreviewImage()
ui = previewer.save_images(image, filename_prefix="Angelo_preview")
return image, ui["ui"]["images"]
class AngeloRefine:
"""Click-to-refine sampler. See module docstring."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"vae": ("VAE",),
# ===== Sampler Mode controls (top of native widget area) =====
# Visible in the node body. When `mode == "Sampler Mode"`, the
# toolbar is greyed and canvas clicks do nothing — this widget
# group is the active one, producing the base latent from the
# incoming latent via comfy.sample.sample. When mode flips to
# Refinement, sampler_seed_control is auto-forced to "fixed"
# so subsequent Queue presses don't regenerate the base.
"mode": (["Sampler Mode", "Edit Mode"], {"default": "Sampler Mode",
"tooltip": "Sampler Mode: AngeloRefine acts "
"like a KSampler — generates the "
"base latent from the incoming "
"(usually empty) latent. Toolbar "
"and canvas clicks are inert. "
"Edit Mode: click / paint / drag "
"to refine or inpaint the cached "
"base. Switching to Edit Mode auto-"
"locks sampler_seed_control to "
"'fixed' so the base stays stable."}),
"sampler_denoise": ("FLOAT", {"default": 1.0, "min": 0.05, "max": 1.0, "step": 0.05,
"tooltip": "[Sampler Mode] Denoise level for the base "
"generation. 1.0 = generate fully from noise "
"(KSampler default). Lower = img2img-style from "
"the incoming latent."}),
"sampler_seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF,
"control_after_generate": False,
"tooltip": "[Sampler Mode] Seed for the base generation. "
"Controlled by sampler_seed_control after each run."}),
"sampler_seed_control": (["fixed", "increment", "decrement", "randomize"],
{"default": "randomize",
"tooltip": "[Sampler Mode] After-generation seed behaviour. "
"Defaults to randomize so each base generation gets "
"a fresh seed. Auto-forced to 'fixed' (locked to the "
"seed that produced the base) when you switch to Edit "
"Mode so refine clicks stay stable."}),
# ===== Edit Mode controls (driven from toolbar) =====
# Hidden from the native widget area; the toolbar above the
# preview canvas drives these. seed_control follows the same
# fixed/random/increment/decrement pattern as sampler_seed.
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF,
"control_after_generate": False,
"tooltip": "[Edit Mode] Seed for the refine pass. Hidden — "
"controlled by the Seed input on the toolbar."}),
"seed_control": (["fixed", "increment", "decrement", "randomize"],
{"default": "randomize",
"tooltip": "[Edit Mode] After-click seed behaviour. Hidden — "
"controlled by the Seed Ctrl dropdown on the toolbar. "
"Defaults to randomize so each refine click produces a "
"fresh variation rather than repeating the same result."}),
"steps": ("INT", {"default": 4, "min": 1, "max": 100,
"tooltip": "Match the model's expected step count. "
"FLUX 2 Klein distilled = 4."}),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.1,
"tooltip": "FLUX 2 Klein distilled uses CFG=1 (no negative)."}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "simple"}),
"denoise": ("FLOAT", {"default": 0.5, "min": 0.05, "max": 1.0, "step": 0.05,
"tooltip": "How much noise to add back for the refinement. "
"0.3 = subtle touch-up, 0.6 = real redo, "
"0.9+ = essentially regenerate that region."}),
"click_radius": ("INT", {"default": 96, "min": 8, "max": 1024, "step": 4,
"tooltip": "Pixel-space radius of the refinement region. "
"Updated automatically by the JS click widget."}),
"feather_radius": ("INT", {"default": 24, "min": 0, "max": 256, "step": 4,
"tooltip": "Pixel-space gaussian blur applied to the mask "
"before sampling. Smooths the seam between the "
"refined region and the preserved surroundings. "
"Roughly half of click_radius is a good default."}),
# JS-driven widgets. click_x/y in pixel space; -1 = no click yet.
# click_seq increments per click to force ComfyUI to detect a change
# and re-execute this node (otherwise identical inputs would skip).
# image_w/h are the actual decoded image dimensions in pixels —
# set by the JS whenever a fresh preview is loaded into the canvas,
# so we can compute the correct pixel→latent scale without
# hardcoding the VAE downscale factor (FLUX 2 is 16, FLUX 1 / SDXL
# are 8, exotic VAEs vary).
"click_x": ("INT", {"default": -1, "min": -1, "max": 16384}),
"click_y": ("INT", {"default": -1, "min": -1, "max": 16384}),
"click_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
"image_w": ("INT", {"default": 0, "min": 0, "max": 16384}),
"image_h": ("INT", {"default": 0, "min": 0, "max": 16384}),
# Undo bookkeeping: undo_seq increments when the user clicks
# the Undo button. Python pops the last refined latent off
# the history stack on each new undo_seq value.
"undo_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
"reset": ("BOOLEAN", {"default": False,
"tooltip": "Tick + re-queue to discard the cached refined "
"latent and start over from the incoming latent."}),
# DEPRECATED as a control — always-on now. Kept declared
# (not removed) so its slot in widgets_values stays put;
# deleting it would shift every later widget's position
# and drift old saved workflows on load. Hidden in the UI;
# run() ignores the value and always decodes the preview.
"auto_decode": ("BOOLEAN", {"default": True,
"tooltip": "(Deprecated — preview always decodes now.)"}),
# STEP 2 of the Fine Upscaling re-introduction: JS toggle
# bar drives these. Both widgets are hidden from the node
# UI; the green "Fine Upscale" toggle and the small "MP"
# numeric input on the bar above the preview canvas are
# the user-facing surface. Python's run() still prints the
# received values so we can verify the JS bar correctly
# drives them. SAMPLING IS STILL UNCHANGED — neither value
# is read by any code path below the print. Step 3 wires
# the actual crop+upscale+refine.
"fine_upscaling": ("BOOLEAN", {"default": False,
"tooltip": "When ON, the painted region is cropped, "
"bilinear-upscaled in latent space to hit "
"min_megapixels, refined, downscaled, and "
"composited back. Gives the model more "
"effective resolution on small painted "
"regions. Upscale capped at 8× linear "
"internally. Set via the Fine Upscale "
"toggle above the preview."}),
"min_megapixels": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 4.0, "step": 0.1,
"tooltip": "Target megapixels for the refine pass when "
"Fine Upscaling is on. Only used in that mode. "
"Set via the MP input above the preview."}),
"max_upscale": ("FLOAT", {"default": 8.0, "min": 1.0, "max": 16.0, "step": 0.5,
"tooltip": "Hard cap on the linear upscale factor in Fine "
"Upscaling. 8.0 was the original internal default. "
"Lower values reduce smoothing artifacts from "
"extreme upscales at the cost of less effective "
"resolution gain. Set via the Max input above "
"the preview."}),
"resize_method": (_FINE_UPSCALE_RESIZE_METHODS, {"default": "lanczos",
"tooltip": "Pixel-space upscale method for "
"Fine Upscale. lanczos is the "
"default — sharpest detail recovery "
"for natural images. bilinear is "
"smooth (good for skin/faces); "
"bicubic middle ground; nearest-"
"exact preserves exact sample values; "
"area/bislerp niche."}),
"inpainting_mode": (["Refine", "Smart Inpaint", "Smart Guided Inpaint", "Outpaint"], {"default": "Refine",
"tooltip": "How the painted region is treated.\n\n"
"Refine — the painted region is partially "
"denoised from the existing content. Best "
"for refining what's already there (faces, "
"hands, textures). Paint/click as normal.\n\n"
"Smart Inpaint — adds NEW content where you "
"drag a rectangle. Click+hold one corner, "
"release at the opposite corner. Locks "
"denoise=1.0, Fine Upscale=ON, Ctx Pad=0 "
"(the right defaults for adding new subjects "
"with an edit model — Klein 9B etc.). "
"Reference_latents are auto-injected so the "
"model's edit branch activates.\n\n"
"Smart Guided Inpaint — no painting or boxes. "
"Pick a location from the dropdown above the "
"Area Prompt; it's prepended to your prompt "
"(e.g. 'In the top left of the image, ...') "
"and the edit model places the content there "
"across the whole image. Locks denoise=1.0, "
"Fine Upscale=OFF, Area Prompt=ON; press "
"'Generate Guided Edit' to run.\n\n"
"Outpaint — extend the canvas. Pick a direction "
"(arrows, or click near an edge of the preview), "
"review the result, Accept to commit. Accepting "
"installs the new canvas as a fresh session base "
"— history resets."}),
# Hidden — DOM location dropdown (Smart Guided Inpaint)
# drives this. Holds a LABEL key from
# _GUIDED_LOCATION_PREFIXES; run() maps it to the prompt
# prefix at encode time.
"guided_location": ("STRING", {"default": "(none)", "multiline": False}),
"fine_context_pad": ("INT", {"default": 64, "min": 0, "max": 512, "step": 8,
"tooltip": "[Fine Upscale] Pixel-space padding around the "
"painted-shape bbox before cropping. Gives the "
"model surrounding context (skin, hair, "
"background) so a tight face mask at high denoise "
"still produces a coherent face that matches its "
"surroundings. The painted shape is unchanged — "
"only the area the model SEES grows. Outside the "
"painted shape, the surrounding pixels are "
"preserved (not refined). Larger pad = more "
"context + less effective resolution on the "
"painted area (bump MP to compensate)."}),
"persistent_mask": ("BOOLEAN", {"default": False,
"tooltip": "When ON, the last mask used (click point or paint "
"stroke) is held in the node. Pressing the standard "
"ComfyUI Queue button re-runs that same mask + a fresh "
"seed on the LATEST result, so each press builds further "
"on that region — use it to gradually morph an area over "
"several presses while the rest of the image stays "
"unchanged. (To re-roll the same edit on the ORIGINAL "
"image instead of building on it, use the Re-roll "
"button.) Toggled via the Persistent Mask button above "
"the preview."}),
"paint_mode": ("BOOLEAN", {"default": False,
"tooltip": "When ON, hold + drag on the preview paints a "
"freeform brush stroke (each dragged point is the "
"centre of a circle of click_radius; the union is "
"the refine mask). Release to submit. Single clicks "
"in paint mode work as one-point strokes. When OFF, "
"clicks behave as before (single-circle refine)."}),
# Hidden — JS-set JSON list of [x_pixel, y_pixel] points
# captured during a paint drag. Empty when no paint stroke
# is pending (e.g. the user just single-clicked instead).
"stroke_points": ("STRING", {"default": "", "multiline": False}),
# Hidden — JS-set JSON list of [x1, y1, x2, y2] rectangles
# in image-pixel coords, captured during Smart Inpaint
# rectangle drags. The backend uses only the most recent
# entry as the active mask.
"rect_points": ("STRING", {"default": "", "multiline": False}),
"area_prompt": ("BOOLEAN", {"default": False,
"tooltip": "When ON, refinements encode the Area Prompt text "
"(typed in the box below the canvas) with the "
"connected CLIP and use it instead of the main "
"positive/negative — useful for reshaping a region "
"with a different prompt (e.g. main prompt = wide "
"shot, area prompt = \"detailed face, photorealistic "
"eyes\"). The click-and-mask behaviour is unchanged; "
"only the prompt differs. If CLIP isn't connected or "
"the area text is empty, the main positive/negative "
"are used. Forced ON in Smart Inpaint."}),
# Hidden — DOM text box below the canvas drives these.
# The Area Prompt input has a Pos/Neg toggle that decides
# which of these two the textarea is editing. Encoded
# with the connected CLIP at run time when area_prompt is
# on. Negative is usually unused (Klein / CFG=1) but kept
# for non-distilled models.
"area_text_positive": ("STRING", {"default": "", "multiline": True}),
"area_text_negative": ("STRING", {"default": "", "multiline": True}),
# Hidden — the Load Image button drives these. loaded_image
# is a JSON ref {name,subfolder,type} from /upload/image;
# loaded_image_seq bumps on each load so run() knows a new
# image arrived; resize_mode/target_mp come from the popup.
"loaded_image": ("STRING", {"default": "", "multiline": False}),
"loaded_image_seq": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
"loaded_resize_mode": (["keep", "mp"], {"default": "keep"}),
"loaded_target_mp": ("FLOAT", {"default": 1.5, "min": 0.1, "max": 8.0, "step": 0.1}),
# Hidden — the Detect (SAM 3 / YOLO) flow drives this. A JSON
# list of silhouette polygons (image-pixel coords) for the
# confirmed detection; in Refine mode it becomes the mask.
# Cleared by the JS after the confirm run.
"seg_polygon": ("STRING", {"default": "", "multiline": False}),
# Hidden — the Re-roll button drives this. Bumps on each
# press; run() pops the most recent edit to expose its
# pre-edit base, rebuilds the SAME mask from the widgets
# above, and re-samples with a fresh seed — swapping the
# new variation in place of the last attempt. Declared LAST
# so older saved workflows (which lack it) don't shift their
# positional widgets_values; it just defaults to 0.
"reroll_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
# Hidden — the Detect Shift/Alt touch-up brush drives this. A
# base64-PNG mask at image resolution (white = masked) that, in
# Refine, takes priority over seg_polygon. Lets the user add to
# or subtract from (incl. holes) a SAM mask before committing.
# Declared LAST so older saved workflows don't shift their
# positional widgets_values; defaults to "".
"seg_mask_png": ("STRING", {"default": "", "multiline": False}),
# Redo (#6): the Redo button / Ctrl-Y bumps this. Python pushes
# back onto the history stack the entry that Undo most recently
# popped off (held on a per-node redo stack). Declared LAST so
# older saved workflows don't shift their positional
# widgets_values; defaults to 0.
"redo_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
# Restore brush (#12): when ON, clicks / paint strokes / detect
# masks RESTORE the painted region back to the session's base
# image via a feathered latent blend — no sampling at all, so
# it's effectively instant. The Lightroom "erase part of an
# edit" gesture: refine a region, then brush back the bits of
# the original you wanted kept (the face inside the helmet).
# Honoured in Refine mode only; the Smart modes ignore it.
# Declared LAST so older saved workflows don't shift their
# positional widgets_values; defaults to False.
"restore_mode": ("BOOLEAN", {"default": False,
"tooltip": "When ON, clicks and paint strokes restore "
"the painted region back to the session base "
"image (a feathered latent blend — no "
"sampling, instant). Refine mode only. "
"Toggled via the Restore button on the "
"toolbar."}),
# Prompt slots (#12): JSON persistence for the Area Prompt
# slot strip ({active, slots: [{pos, neg} × 6]}). Pure UI
# state — Python never reads it; the JS mirrors the active
# slot's text into area_text_positive/negative, which remain
# the encode source of truth. Declared LAST (see above).
"area_prompt_slots": ("STRING", {"default": "", "multiline": False}),
# Vary ×4: the Vary button bumps vary_seq. run() re-runs the
# most recent edit's mask from its PRE-edit base four times
# with different seeds, stashes the candidates in state
# WITHOUT touching history, and ships one preview ref per
# candidate (Angelo_vary_previews) to the JS chooser overlay.
# Clicking a candidate sets vary_pick + bumps vary_pick_seq;
# run() then swaps the chosen candidate in for the last
# attempt — a pure restore, no sampling, mirroring Re-roll's
# replace semantics. Cancelling the chooser needs no Python
# round-trip (the stash is invalidated by the next edit).
# Declared LAST (see redo_seq note).
"vary_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
"vary_pick": ("INT", {"default": -1, "min": -1, "max": 15}),
"vary_pick_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
# Outpaint: the JS arrows / edge-click bump outpaint_seq with
# a direction + amount + overlap. run() pads the canvas, fills
# the new strip at denoise 1.0, and STASHES the result for
# review (Angelo_outpaint_preview → the JS Accept overlay) —
# history is untouched. Accepting bumps outpaint_accept_seq;
# run() then installs the new canvas as a FRESH session base
# (Load-Image semantics: history resets — a documented fact of
# outpainting, which is why nothing commits without review).
# Cancel needs no Python round-trip. Declared LAST.
"outpaint_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
"outpaint_dir": ("STRING", {"default": "right", "multiline": False}),
"outpaint_amount": ("INT", {"default": 256, "min": 16, "max": 2048, "step": 16,
"tooltip": "Outpaint: how many pixels to extend the "
"canvas by (snapped to /16). Set via the "
"Amount input on the Outpaint row."}),
"outpaint_overlap": ("INT", {"default": 64, "min": 0, "max": 512, "step": 8,
"tooltip": "Outpaint: feathered band reaching this many "
"pixels INTO the existing image, redrawn so "
"the seam blends instead of butting. Set via "
"the Overlap input on the Outpaint row."}),
"outpaint_accept_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
# Outpaint protect brush: JSON list of [x, y, r] circles
# (old-image pixel coords) painted on the canvas in Outpaint
# mode. Subtracted from the overlap band so marked content
# (a car at the frame edge) stays frozen while the rest of
# the band blends. Declared LAST.
"outpaint_protect": ("STRING", {"default": "", "multiline": False}),
# Where the outpaint instruction sits relative to the user's
# Area Prompt text: "prepend" = instruction first (default),
# "append" = the user's text first, instruction after. Some
# models weight the head of the prompt more heavily, so the
# order can change reliability — the JS shows a live preview
# of the combined prompt so there's no guessing. The exact
# composition here MUST stay in lockstep with the JS preview
# (syncOutpaintPromptPreview). Declared LAST.
"outpaint_instruction_pos": (["prepend", "append"], {"default": "prepend"}),
# Reference toggle: ON/OFF for the anchored-refine feature;
# reference_strength (declared further down) is the 0–1
# blend it applies when ON. Manual refines only — Quick
# Photo Refine runs its own fixed recipe and ignores both.
"refine_reference": ("BOOLEAN", {"default": False,
"tooltip": "Anchor refines to the current image "
"(edit models). The strength box "
"beside the toolbar button sets the "
"blend. Toggled via the Reference "
"button."}),
# Quick Photo Refine: the ✨ button bumps this. One-shot
# restoration pass — whole canvas, with the button's OWN
# fixed settings (_QUICK_REFINE_PROMPT / _DENOISE / _REF).
# A magic button: no toolbar box affects it except the
# seed. Pushes a normal history entry so Undo covers it.
# Declared LAST.
"quick_refine_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
# Reference strength (Refine + Quick Photo Refine): a TRUE
# 0–1 blend between the reference-anchored prediction and
# the plain one — see _apply_reference. 0 = no reference
# (classic refine), 1 = fully anchored. Declared LAST.
"reference_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05,
"tooltip": "How strongly the current image anchors "
"a refine — a true blend between the "
"anchored and free predictions. 0 = "
"off, 1 = fully anchored. Set via the "
"Ref box on the toolbar."}),
# ⬆ 2× Pixel: the button bumps this. A PURE pixel-space
# lanczos 2× + re-encode — no AI, deterministic — committed
# directly as a fresh session base (dimension change =
# Load-Image semantics, history resets). The AI step is the
# user's next move (e.g. ✨, which auto-tiles when large).
"upscale_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
# ▦ Large Image Refine: the button bumps this. Chess-pattern
# Xtra-Fine boxes over the whole canvas (see _LIR_* constants).
# One history entry for the whole pass. Declared LAST.
"lir_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
},
"optional": {
# CLIP / text encoder for the Area Prompt. Optional —
# without it the area text is ignored and the main
# positive/negative are used.
"clip": ("CLIP",),
# Base latent. OPTIONAL now: if the Load Image button is
# used, the loaded photo becomes the base and no latent
# input is needed. Sampler Mode still needs one (it defines
# the output dimensions for a fresh generation).
"latent": ("LATENT",),
# Optional sampler-quintet overrides for users who want a
# single source of truth for these values across the
# workflow (issue #25). Pipe-pattern: a single ANGELO_OVERRIDES
# slot, fed by the companion `Angelo — Overrides` node which
# bundles steps / cfg / sampler / scheduler into one dict.
# If unwired, the toolbar widget values flow through as
# before. Picked over per-param forceInput slots because
# forceInput on optional inputs proved unreliable across
# ComfyUI versions for delivering wired primitive values.
"overrides": ("ANGELO_OVERRIDES", {"tooltip": "Optional. Wire from an "
"'Angelo — Overrides' node to "
"drive steps / cfg / sampler / "
"scheduler from your workflow "
"instead of the toolbar."}),
},
"hidden": {
"unique_id": "UNIQUE_ID",
},
}
RETURN_TYPES = ("IMAGE", "LATENT", "IMAGE")
RETURN_NAMES = ("image", "latent", "source_image")
FUNCTION = "run"
OUTPUT_NODE = True
CATEGORY = "sampling/Angelo"
DESCRIPTION = (
"Click-to-refine sampler. Click a region in the preview to do "
"extra denoising on just that area while the rest of the latent "
"is preserved. Repeated clicks layer refinements on top of each other."
)
def run(
self,
model,
positive,
negative,
vae,
mode,
sampler_denoise,
sampler_seed,
sampler_seed_control,
seed,
seed_control,
steps,
cfg,
sampler_name,
scheduler,
denoise,
click_radius,
feather_radius,
click_x,
click_y,
click_seq,
image_w,
image_h,
undo_seq,
reset,
auto_decode,
fine_upscaling,
min_megapixels,
max_upscale,
resize_method,
inpainting_mode,
fine_context_pad,
persistent_mask,
paint_mode,
stroke_points,
rect_points,
area_prompt,
guided_location="(none)",
area_text_positive="",
area_text_negative="",
loaded_image="",
loaded_image_seq=0,
loaded_resize_mode="keep",
loaded_target_mp=1.5,
seg_polygon="",
reroll_seq=0,
seg_mask_png="",
redo_seq=0,
restore_mode=False,
area_prompt_slots="",
vary_seq=0,
vary_pick=-1,
vary_pick_seq=0,
outpaint_seq=0,
outpaint_dir="right",
outpaint_amount=256,
outpaint_overlap=64,
outpaint_accept_seq=0,
outpaint_protect="",
outpaint_instruction_pos="prepend",
refine_reference=False,
quick_refine_seq=0,
reference_strength=0.0,
upscale_seq=0,
lir_seq=0,
latent=None,
clip=None,
overrides=None,
unique_id=None,
):
# Optional upstream overrides: if an ANGELO_OVERRIDES bundle is
# wired in, its non-None entries beat the toolbar widget values.
# Kept here at the top so every downstream code path sees the
# effective value without each one having to know about the override.
# Carries sampler-config (#25), display flags (#21
# disable_live_preview), and the custom-sampler trio (#8
# guider/sampler/sigmas — all-or-nothing, applied inside _do_sample).
disable_live_preview = False
ov_guider = None
ov_sampler = None
ov_sigmas = None
if isinstance(overrides, dict):
if overrides.get("steps") is not None:
steps = overrides["steps"]
if overrides.get("cfg") is not None:
cfg = overrides["cfg"]
if overrides.get("sampler_name") is not None:
sampler_name = overrides["sampler_name"]
if overrides.get("scheduler") is not None:
scheduler = overrides["scheduler"]
disable_live_preview = bool(overrides.get("disable_live_preview"))
ov_guider = overrides.get("guider")
ov_sampler = overrides.get("sampler")
ov_sigmas = overrides.get("sigmas")
node_id = str(unique_id)
state = _STATE.get(node_id)
# ===== Base latent selection =====
# While an image is LOADED (loaded_image non-empty), it owns the
# base and the wired `latent` input is ignored — until the user
# hits Unload (which clears loaded_image). This stops the wired
# latent from re-winning after a load (which made undo snap back
# to the latent image). Priority:
# 1. a freshly loaded image (encode it, force reset)
# 2. an already-loaded image still active → keep the cached base
# 3. the wired `latent` input
# 4. the cached base (no latent, no load)
loaded_ref = str(loaded_image).strip()
loaded_active = bool(loaded_ref)
loaded_seq = int(loaded_image_seq)
new_loaded = loaded_active and (
state is None or state.get("loaded_seq") != loaded_seq
)
forced_base = None
if new_loaded:
forced_base = _encode_loaded_image(
vae, loaded_ref, str(loaded_resize_mode), float(loaded_target_mp)
)
# base_from_wired_latent: only the wired-`latent` path carries a
# meaningful fingerprint for upstream-change detection. The cache /
# loaded-image paths take incoming from history[0], which MUTATES as
# the undo stack is capped (the original base is evicted after
# _HISTORY_CAP refines), so a fingerprint comparison there is bogus
# and would spuriously reset mid-session. See the need_reset note.
base_from_wired_latent = False
if forced_base is not None:
incoming = forced_base
elif loaded_active and state is not None and state.get("history"):
# Loaded image still active — keep its base, ignore `latent`.
hist_0 = state["history"][0]
incoming = hist_0[0] if isinstance(hist_0, tuple) else hist_0
elif latent is not None:
incoming = latent["samples"]
base_from_wired_latent = True
elif state is not None and state.get("history"):
hist_0 = state["history"][0]
incoming = hist_0[0] if isinstance(hist_0, tuple) else hist_0
else:
raise ValueError(
"Angelo: no base image — connect a `latent` input or use the "
"Load Image button."
)
# Normalise the base latent to the dimensionality the MODEL expects.
# ComfyUI's stock KSampler runs the latent through
# fix_empty_latent_channels before sampling; Angelo calls
# comfy.sample.sample directly, so it must do this itself. The
# load-bearing case is video/temporal VAEs (Qwen Image Edit, Wan):
# their diffusion model wants a 5D latent [B, C, T, H, W], and a 4D
# [B, C, H, W] base (e.g. from a plain EmptyLatentImage, or any path
# that didn't go through a Qwen-aware latent node) makes process_img
# fail with "not enough values to unpack (expected 5, got 4)". This
# unsqueezes the temporal axis (and channel-pads an empty latent) for
# those models; for ordinary 4D models (FLUX/SDXL/SD) it is a no-op,
# and an already-5D latent is returned unchanged. Done once here so
# every downstream path — Sampler Mode, the Edit-Mode history seed,
# the fingerprint, and the refine round-trips — sees the right shape.
incoming = comfy.sample.fix_empty_latent_channels(model, incoming)
incoming_fp = _latent_fingerprint(incoming)
# ===== Smart Inpaint locks =====
# The whole point of Smart Inpaint as a mode is to opinionate
# the params that matter for "add new content in a rectangle"
# workflows. Override the widget values up front so every code
# path downstream sees the locked values regardless of what
# the user set the toolbar to.
# area_prompt ON makes Smart Inpaint encode the Area Prompt text
# (typed below the canvas) with the connected CLIP and use it — a
# separate "what to insert" prompt, kept independent of the main scene
# prompt. While area_prompt is on the refine uses the Area text ONLY
# (empty Area text → empty conditioning), never the main positive — see
# the area-conditioning block further down in run().
if inpainting_mode == "Smart Inpaint":
denoise = 1.0
fine_context_pad = 0
fine_upscaling = True
area_prompt = True
# feather_radius is left under user control — a soft edge can
# help blend the inserted content into the surroundings.
elif inpainting_mode == "Smart Guided Inpaint":
# No painting / boxes — the whole image is edited and the
# location dropdown supplies a prompt prefix that tells the
# edit model where to place the content. No region to crop
# (Fine Upscale OFF), no mask edge to feather, and no mask to
# persist (Persistent Mask is meaningless here).
denoise = 1.0
fine_context_pad = 0
fine_upscaling = False
area_prompt = True
feather_radius = 0
persistent_mask = False
elif inpainting_mode == "Outpaint":
# Canvas extension, not a masked edit — there's no held mask to
# re-run, so the Persistent Mask queue hook must not fire. The
# other edit params aren't read on the outpaint path (denoise is
# fixed at 1.0 inside it; the JS dims their controls).
persistent_mask = False
# ===== Sampler Mode branch =====
# Acts like a KSampler: take the incoming latent (typically empty),
# run a fresh denoise pass with sampler_seed + sampler_denoise, cache
# the result as the new base. All toolbar / canvas / refine logic is
# skipped — those are Edit Mode concerns.
if mode == "Sampler Mode":
# #21: skip the live-preview callback when the user has
# disabled it via Overrides — keeps ComfyUI's global preview
# on for other samplers but stops it squashing Angelo's editor.
callback = None if disable_live_preview else latent_preview.prepare_callback(model, steps)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
noise = comfy.sample.prepare_noise(incoming, sampler_seed, None)
# #8: _do_sample dispatches to the custom guider path when
# ov_guider/ov_sampler/ov_sigmas are all wired, otherwise
# falls back to the standard comfy.sample.sample(...) call.
new_latent = _do_sample(
guider=ov_guider, sampler=ov_sampler, sigmas=ov_sigmas,
model=model, noise=noise,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=positive, negative=negative,
source_latent=incoming,
denoise=sampler_denoise,
callback=callback,
disable_pbar=disable_pbar,
seed=sampler_seed,
)
# Replace the cache with the freshly-sampled base. Drops the
# undo history (it's irrelevant — we have a brand-new image).
#
# CRITICAL: store the fingerprint of the INCOMING latent, not
# the freshly-generated new_latent. The fingerprint is the
# "did upstream change?" signal used by Edit Mode's
# auto-reset check. If we stored new_latent's fingerprint,
# Edit Mode would always see (cached_fp != incoming_fp)
# → think upstream changed → wipe cache → refine the empty
# latent. Storing incoming_fp means Refinement only resets
# when the user actually rewires upstream.
#
# sampler_seed_at_run = the seed Python used. Sent back to JS
# via the ui message so JS can (a) apply after-gen control
# next, and (b) restore this value if the user later switches
# the control to "fixed".
_STATE[node_id] = {
"history": [(new_latent, None)],
"click_seq": click_seq,
"undo_seq": undo_seq,
# Anchor every action-seq to its current widget value so a
# stale bump can't refire its gate against the new base.
"redo_seq": redo_seq,
"reroll_seq": reroll_seq,
"vary_seq": vary_seq,
"vary_pick_seq": vary_pick_seq,
"outpaint_seq": outpaint_seq,
"outpaint_accept_seq": outpaint_accept_seq,
"quick_refine_seq": quick_refine_seq,
"upscale_seq": upscale_seq,
"lir_seq": lir_seq,
"fingerprint": incoming_fp,
"sampler_seed_at_run": int(sampler_seed),
"loaded_seq": loaded_seq,
# Source image (#3/#9): the session base, captured once so the
# source_image output survives _HISTORY_CAP eviction of history[0].
"source_latent": new_latent,
"source_pixels": None,
}
out_latent = {"samples": new_latent}
ui_msg = {
"Angelo_preview": [],
"Angelo_mode": ["Sampler Mode"],
"Angelo_sampler_seed_at_run": [int(sampler_seed)],
}
# Preview always decodes now (auto_decode deprecated).
image, image_refs = _decode_to_preview(vae, new_latent)
ui_msg["Angelo_preview"] = image_refs
# Freshly-generated base IS the source image — cache + emit it.
# The preview refs double as the source refs for the JS
# hold-to-compare key (the base and the preview are the same
# image until the first edit).
_STATE[node_id]["source_pixels"] = image
_STATE[node_id]["source_preview_refs"] = image_refs
ui_msg["Angelo_source_preview"] = image_refs
return {"ui": ui_msg, "result": (image, out_latent, image)}
# ===== Edit Mode branch (existing behaviour) =====
# Decide whether to (re)seed the cache from the incoming latent.
# Reset on explicit toggle, first run, or fingerprint change —
# but NOT on fingerprint change while persistent_mask is on, because
# the whole point of persistent_mask is to keep refining the cached
# latent across upstream re-rolls (the user's pressing Queue
# specifically because they want a variation of the held region,
# not a fresh image).
#
# The fingerprint check ONLY applies when the base is the wired
# `latent` input (where a change really does mean "upstream produced
# a fresh latent"). When the base comes from the cache / loaded
# image, `incoming` is history[0] — which mutates as the undo stack
# is capped (_HISTORY_CAP). Comparing against it there would
# spuriously reset to a mid-stage latent after enough refines (the
# "suddenly reverts to an earlier stage while painting" bug). Those
# bases only change via explicit Load / Unload / Reset, all handled
# by `reset` / `new_loaded` above — so the fingerprint isn't needed.
state = _STATE.get(node_id)
fingerprint_changed = (
base_from_wired_latent
and state is not None
and state.get("fingerprint") != incoming_fp
)
need_reset = (
reset
or state is None
or new_loaded
or (fingerprint_changed and not persistent_mask)
)
if need_reset:
# Any reset means the base just changed (load, unload, upstream
# rewire). Anchor click_seq/undo_seq to the CURRENT widget
# values so a click that was meant for the OLD base can't trip
# the new-click gate and replay a stale inpaint onto the new
# base. The user's next genuine click bumps click_seq and fires
# normally.
_STATE[node_id] = {
"history": [(incoming.clone(), None)],
"click_seq": click_seq,
"undo_seq": undo_seq,
# Anchor every action-seq (see the Sampler Mode dict note).
"redo_seq": redo_seq,
"reroll_seq": reroll_seq,
"vary_seq": vary_seq,
"vary_pick_seq": vary_pick_seq,
"outpaint_seq": outpaint_seq,
"outpaint_accept_seq": outpaint_accept_seq,
"quick_refine_seq": quick_refine_seq,
"upscale_seq": upscale_seq,
"lir_seq": lir_seq,
"fingerprint": incoming_fp,
"loaded_seq": loaded_seq,
# Source image (#3/#9): capture the base once, independent of
# history[0] (which mutates under _HISTORY_CAP eviction).
"source_latent": incoming.clone(),
"source_pixels": None,
}
state = _STATE[node_id]
# Undo: if undo_seq advanced and we have history to pop, pop it.
# We always keep at least one latent (the base / earliest refine)
# so the preview stays valid.
new_undo = undo_seq > 0 and undo_seq != state.get("undo_seq", -1)
if new_undo:
if len(state["history"]) > 1:
# Move the popped entry onto the redo stack so Redo (#6) can
# restore it. Bounded like the history stack.
popped = state["history"].pop()
redo = state.setdefault("redo_stack", [])
redo.append(popped)
if len(redo) > _HISTORY_CAP:
state["redo_stack"] = redo[-_HISTORY_CAP:]
state["undo_seq"] = undo_seq
# Undo is a PURE restore — pop the cached latent and decode it.
# It must NEVER re-sample, or it would produce a different image
# than the one being restored. Absorb the current click_seq so
# the new-click gate below stays False on this run. This is
# load-bearing because the Persistent Mask queue hook bumps
# click_seq on EVERY queue, including the Undo button's — without
# this, an undo while Persistent Mask is on would look like a new
# click and re-run the last mask with the (since-randomized) seed,
# restoring the WRONG result. (Harmless no-op when the hook didn't
# bump it.)
state["click_seq"] = click_seq
# ===== Redo (#6): restore an entry Undo moved to the redo stack =====
# A PURE restore — never re-samples — and runs BEFORE `current` is read
# below, so the restored entry becomes history[-1]. Absorbs click_seq
# for the same reason Undo does (the Persistent Mask queue hook bumps
# it on every queue, including the Redo button's). No-op if the redo
# stack is empty. A genuine new edit clears the redo stack (below).
new_redo = redo_seq > 0 and redo_seq != state.get("redo_seq", -1)
if new_redo:
redo = state.get("redo_stack") or []
if redo:
state["history"].append(redo.pop())
if len(state["history"]) > _HISTORY_CAP:
state["history"] = state["history"][-_HISTORY_CAP:]
state["redo_seq"] = redo_seq
state["click_seq"] = click_seq
# ===== Vary pick: commit a chosen variation (pure restore) =====
# The JS chooser sets vary_pick + bumps vary_pick_seq after a Vary
# run stashed candidates. Swap the chosen candidate in for the last
# attempt — replace, not stack, exactly like Re-roll — and clear the
# stash. Never samples. Absorbs click_seq for the same queue-hook
# reason as Undo/Redo.
new_vary_pick = vary_pick_seq > 0 and vary_pick_seq != state.get("vary_pick_seq", -1)
if new_vary_pick:
state["vary_pick_seq"] = vary_pick_seq
cands = state.get("vary_candidates")
if cands and 0 <= int(vary_pick) < len(cands):
state["history"][-1] = cands[int(vary_pick)]
# Picking a variation is a real edit decision — the redo
# branch (which predates it) no longer applies.
state["redo_stack"] = []
state["vary_candidates"] = None
state["click_seq"] = click_seq
# ===== Outpaint accept: commit the reviewed canvas as a NEW session =====
# The JS Accept button bumps outpaint_accept_seq after the review
# overlay. The stashed canvas becomes a fresh session base with
# Load-Image semantics — history resets, the Restore anchor and the
# hold-to-compare base move to the new canvas. This wholesale state
# replacement is what keeps every same-shape assumption in the rest
# of the node true across a dimension change.
new_op_accept = (
outpaint_accept_seq > 0
and outpaint_accept_seq != state.get("outpaint_accept_seq", -1)
)
if new_op_accept:
pend = state.get("outpaint_pending")
if pend is not None:
op_lat, op_px = pend
_STATE[node_id] = {
"history": [(op_lat, op_px)],
"click_seq": click_seq,
"undo_seq": undo_seq,
"redo_seq": redo_seq,
"reroll_seq": reroll_seq,
"vary_seq": vary_seq,
"vary_pick_seq": vary_pick_seq,
"outpaint_seq": outpaint_seq,
"outpaint_accept_seq": outpaint_accept_seq,
"quick_refine_seq": quick_refine_seq,
"upscale_seq": upscale_seq,
"lir_seq": lir_seq,
# Preserve the wired-latent fingerprint + load marker so
# the next run doesn't read the unchanged upstream as a
# fresh latent and blow away the canvas we just committed.
"fingerprint": state.get("fingerprint"),
"loaded_seq": state.get("loaded_seq"),
"sampler_seed_at_run": state.get("sampler_seed_at_run", int(sampler_seed)),
"refine_seed_at_run": (state.get("refine_seed_at_run")
if state.get("refine_seed_at_run") is not None
else int(seed)),
"source_latent": op_lat,
"source_pixels": op_px,
}
state = _STATE[node_id]
else:
state["outpaint_accept_seq"] = outpaint_accept_seq
state["click_seq"] = click_seq
# ===== Re-roll: redo the most recent edit with a fresh seed =====
# The Re-roll button bumps reroll_seq (and sets a new seed) without
# touching the mask widgets. Re-run the SAME mask on the SAME pre-
# edit base and swap the result in place of the last attempt — so
# the user can cycle seeds on one edit without reset → re-mask →
# rerun. Implemented as "pop the last refine to expose its pre-edit
# base as history[-1]"; the edit block below then re-runs from there
# and appends the fresh variation, restoring the stack depth (net
# effect: replace, not stack). No-op if there's no prior edit yet.
new_reroll = reroll_seq > 0 and reroll_seq != state.get("reroll_seq", -1)
reroll_now = (
new_reroll
and len(state["history"]) > 1
and inpainting_mode != "Outpaint"
)
if reroll_now:
state["history"].pop()
state["reroll_seq"] = reroll_seq
# ===== Vary ×4 gate =====
# Re-run the most recent edit's mask from its PRE-edit base, four
# times, WITHOUT touching history — candidates are stashed and the
# user commits one via the chooser (vary_pick above). Needs a prior
# edit (history > 1). Excluded when the Restore brush is active
# (restores are deterministic — four identical candidates would be
# noise); the JS blocks that combination too.
new_vary = vary_seq > 0 and vary_seq != state.get("vary_seq", -1)
vary_now = (
new_vary
and len(state["history"]) > 1
and inpainting_mode != "Outpaint"
and not (bool(restore_mode) and inpainting_mode == "Refine")
)
state["vary_seq"] = vary_seq
hist_last = state["history"][-1]
if isinstance(hist_last, tuple):
current, current_pixels = hist_last
else:
current = hist_last
current_pixels = None
# ===== Outpaint (generate): pad + fill, stash for review =====
# Driven by outpaint_seq from the JS arrows / edge-click, gated on
# Outpaint mode. The result is NOT committed — it's stashed with a
# preview ref so the JS can show the Accept / Try again / Cancel
# overlay. The node keeps outputting the current canvas underneath.
new_outpaint = (
inpainting_mode == "Outpaint"
and outpaint_seq > 0
and outpaint_seq != state.get("outpaint_seq", -1)
)
state["outpaint_seq"] = outpaint_seq
if new_outpaint:
latent_init, op_mask, op_amt, op_ref = _outpaint_prepare(
vae=vae, current=current, current_pixels=current_pixels,
direction=str(outpaint_dir), amount_px=int(outpaint_amount),
overlap_px=int(outpaint_overlap),
protect_json=str(outpaint_protect),
)
# Conditioning. With a CLIP wired, the outpaint encodes its OWN
# prompt: a direction-aware continue-don't-repeat instruction +
# the Area Prompt text (if armed) describing the new space. The
# main scene prompt is deliberately NOT used — conditioning the
# strip on "a red car on a road" is precisely what paints a
# second car into it. Without a CLIP we can't encode, so the
# main conditioning flows through (with its known duplication
# risk — wiring CLIP is the documented recommendation).
if clip is not None:
instr = _OUTPAINT_INSTRUCTIONS.get(
str(outpaint_dir), _OUTPAINT_INSTRUCTIONS["right"])
# Instruction order: "prepend" (default) puts the instruction
# first; "append" puts the user's Area text first. MUST stay
# in lockstep with the JS preview (syncOutpaintPromptPreview).
user_txt = str(area_text_positive).strip() if area_prompt else ""
if user_txt and str(outpaint_instruction_pos) == "append":
op_text = user_txt.rstrip(" .,") + ". " + instr.strip()
else:
op_text = instr + user_txt
tokens_p = clip.tokenize(op_text)
op_positive = clip.encode_from_tokens_scheduled(tokens_p)
if area_prompt and str(area_text_negative).strip():
tokens_n = clip.tokenize(str(area_text_negative))
op_negative = clip.encode_from_tokens_scheduled(tokens_n)
else:
op_negative = negative
else:
op_positive = positive
op_negative = negative
# Edge-band reference for edit models (Klein / Qwen): texture +
# lighting to continue, without re-showing the scene's subject
# (a whole-image reference invites the edit branch to reproduce
# it into the strip). REPLACE (append=False) per the Smart
# Inpaint lesson — non-edit models ignore the field entirely.
op_positive = node_helpers.conditioning_set_values(
op_positive, {"reference_latents": [op_ref]}, append=False,
)
# 5D temporal latents (Qwen/Wan) sample with a HARD mask — same
# reasoning as Smart Inpaint (no clean soft-mask-during-denoise
# behaviour on non-zero-mean latent spaces).
op_sample_mask = op_mask
if latent_init.ndim == 5:
op_sample_mask = (op_mask >= 0.5).to(op_mask.dtype)
op_seed = int(seed)
callback = None if disable_live_preview else latent_preview.prepare_callback(model, steps)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
print(f"[Angelo outpaint] dir={outpaint_dir} amount={op_amt}px "
f"overlap={int(outpaint_overlap)}px "
f"new_latent=({latent_init.shape[-2]}x{latent_init.shape[-1]})")
noise = comfy.sample.prepare_noise(latent_init, op_seed, None)
op_refined = _do_sample(
guider=ov_guider, sampler=ov_sampler, sigmas=ov_sigmas,
model=model, noise=noise,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=op_positive, negative=op_negative,
source_latent=latent_init,
denoise=1.0,
noise_mask=op_sample_mask,
callback=callback,
disable_pbar=disable_pbar,
seed=op_seed,
)
op_img, op_refs = _decode_to_preview(vae, op_refined)
state["outpaint_pending"] = (op_refined, op_img)
state["outpaint_preview_refs"] = op_refs
state["click_seq"] = click_seq
state["refine_seed_at_run"] = op_seed
# ===== Quick Photo Refine: the one-click restoration recipe =====
# Whole canvas, the internal restoration prompt, and the Reference
# anchor — identity reconstructs from the reference while the
# texture re-renders. The toolbar DENOISE applies (1.0 = full
# re-render, the strongest restoration; lower = gentler clean-up);
# area prompt / toggles are ignored. The seed also applies, so
# seed_control = randomize lets you mash the button for variations.
new_quick = (
inpainting_mode == "Refine"
and quick_refine_seq > 0
and quick_refine_seq != state.get("quick_refine_seq", -1)
)
state["quick_refine_seq"] = quick_refine_seq
if new_quick:
# Per-model prompt: 5D latent = Qwen/Wan family (fuller, gentler
# instruction); 4D = FLUX family ("restore the photo").
qr_prompt = (_QUICK_REFINE_PROMPT_QWEN if current.dim() == 5
else _QUICK_REFINE_PROMPT)
if clip is not None:
tokens_q = clip.tokenize(qr_prompt)
qr_base = clip.encode_from_tokens_scheduled(tokens_q)
else:
# No CLIP → can't encode the restoration prompt; the main
# positive flows through. Still works, just less targeted.
qr_base = positive
qr_seed = int(seed)
callback = None if disable_live_preview else latent_preview.prepare_callback(model, steps)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
# Big canvases (e.g. after a 2× Restore Upscale) auto-route
# through the tiled engine — same recipe per ~1MP tile, each
# anchored to its own content, seams feathered — so the model
# never samples a latent beyond the size it renders well at.
canvas_mp = (image_w * image_h) / 1e6 if (image_w > 0 and image_h > 0) else 0.0
if canvas_mp > _QUICK_REFINE_TILE_THRESHOLD_MP:
print(f"[Angelo quick-refine] {canvas_mp:.1f}MP canvas — tiled restore pass")
qr_refined, qr_pixels, _n = _tiled_restore_pass(
model=model, vae=vae,
current=current, current_pixels=current_pixels,
scale=1.0,
positive_base=qr_base, negative=negative,
seed=qr_seed, steps=steps, cfg=cfg,
sampler_name=sampler_name, scheduler=scheduler,
callback=callback, disable_pbar=disable_pbar,
ov_guider=ov_guider, ov_sampler=ov_sampler, ov_sigmas=ov_sigmas,
)
else:
# Whole-image reference — IDENTICAL construction to the
# manual path (whole canvas painted), at the button's own
# fixed strength. Magic button: NO toolbar box affects it
# (only the seed, for variation mashing).
qr_positive = _apply_reference(qr_base, current.clone(), _QUICK_REFINE_REF)
print(f"[Angelo quick-refine] whole-canvas restoration pass, "
f"denoise={_QUICK_REFINE_DENOISE}, ref={_QUICK_REFINE_REF}, seed={qr_seed}")
noise = comfy.sample.prepare_noise(current, qr_seed, None)
qr_refined = _do_sample(
guider=ov_guider, sampler=ov_sampler, sigmas=ov_sigmas,
model=model, noise=noise,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=qr_positive, negative=negative,
source_latent=current,
denoise=_QUICK_REFINE_DENOISE,
noise_mask=None, # whole canvas — no mask
callback=callback,
disable_pbar=disable_pbar,
seed=qr_seed,
)
qr_pixels = None
state["history"].append((qr_refined, qr_pixels))
if len(state["history"]) > _HISTORY_CAP:
state["history"] = state["history"][-_HISTORY_CAP:]
state["redo_stack"] = []
state["vary_candidates"] = None
state["outpaint_pending"] = None
state["click_seq"] = click_seq
state["refine_seed_at_run"] = qr_seed
current = qr_refined
current_pixels = qr_pixels
# ===== ⬆ 2× Pixel: pure pixel-space upscale (no AI) =====
# Lanczos 2× of the decoded canvas, re-encoded, committed DIRECTLY
# as a fresh session base — dimension change = Load-Image semantics
# (history resets). Deterministic, so no review step. The AI
# enhancement is the user's next, separate move — e.g. ✨ Quick
# Photo Refine, which auto-tiles on the now-large canvas.
new_upscale = (
inpainting_mode == "Refine"
and upscale_seq > 0
and upscale_seq != state.get("upscale_seq", -1)
)
state["upscale_seq"] = upscale_seq
if new_upscale:
up_px_in = current_pixels if current_pixels is not None else _vae_decode(vae, current)
in_H, in_W = up_px_in.shape[1], up_px_in.shape[2]
out_W = max(16, int(round(in_W * 2.0 / 16.0)) * 16)
out_H = max(16, int(round(in_H * 2.0 / 16.0)) * 16)
up_chw = up_px_in.movedim(-1, 1)
up_chw = comfy.utils.common_upscale(up_chw, out_W, out_H, "lanczos", "disabled")
up_pixels = up_chw.movedim(1, -1).contiguous()
up_latent = _vae_encode(vae, up_pixels)
print(f"[Angelo 2x-pixel] {in_W}x{in_H} -> {out_W}x{out_H} (lanczos, no AI)")
_STATE[node_id] = {
"history": [(up_latent, up_pixels)],
"click_seq": click_seq,
"undo_seq": undo_seq,
"redo_seq": redo_seq,
"reroll_seq": reroll_seq,
"vary_seq": vary_seq,
"vary_pick_seq": vary_pick_seq,
"outpaint_seq": outpaint_seq,
"outpaint_accept_seq": outpaint_accept_seq,
"quick_refine_seq": quick_refine_seq,
"upscale_seq": upscale_seq,
"lir_seq": lir_seq,
"fingerprint": state.get("fingerprint"),
"loaded_seq": state.get("loaded_seq"),
"sampler_seed_at_run": state.get("sampler_seed_at_run", int(sampler_seed)),
"refine_seed_at_run": (state.get("refine_seed_at_run")
if state.get("refine_seed_at_run") is not None
else int(seed)),
"source_latent": up_latent,
"source_pixels": up_pixels,
}
state = _STATE[node_id]
current, current_pixels = up_latent, up_pixels
# ===== ▦ Large Image Refine: chess-pattern Xtra-Fine boxes =====
# The whole canvas refined box by box through the Xtra-Fine
# pipeline — bit-exact latent compositing per box, no pixel
# feathering. Chess order so the second colour's boxes see
# already-refined neighbours in their context pad. One history
# entry for the whole pass (one Undo reverts it all).
new_lir = (
inpainting_mode == "Refine"
and lir_seq > 0
and lir_seq != state.get("lir_seq", -1)
)
state["lir_seq"] = lir_seq
if new_lir:
if clip is not None:
tokens_l = clip.tokenize(_LIR_PROMPT)
lir_base = clip.encode_from_tokens_scheduled(tokens_l)
else:
lir_base = positive
lir_seed = int(seed)
callback = None if disable_live_preview else latent_preview.prepare_callback(model, steps)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
lir_px = current_pixels if current_pixels is not None else _vae_decode(vae, current)
H_pix, W_pix = lir_px.shape[1], lir_px.shape[2]
nlh, nlw = current.shape[-2], current.shape[-1]
sx = nlw / W_pix
sy = nlh / H_pix
box_edge = math.sqrt(_LIR_BOX_MP * 1e6)
nx = max(1, round(W_pix / box_edge))
ny = max(1, round(H_pix / box_edge))
x_edges = [round(i * nlw / nx) for i in range(nx + 1)]
y_edges = [round(j * nlh / ny) for j in range(ny + 1)]
boxes = [(i, j, x_edges[i], y_edges[j], x_edges[i + 1], y_edges[j + 1])
for j in range(ny) for i in range(nx)]
# Chess order: one colour, then the other.
order = ([b for b in boxes if (b[0] + b[1]) % 2 == 0]
+ [b for b in boxes if (b[0] + b[1]) % 2 == 1])
print(f"[Angelo large-refine] {W_pix}x{H_pix}: {nx}x{ny} boxes "
f"(~{_LIR_BOX_MP}MP target), chess order, ref={_LIR_REF}, "
f"denoise={_LIR_DENOISE}, pad={_LIR_CTX_PAD}px, feather 0")
work_lat, work_px = current, lir_px
for k, (bi, bj, x0, y0, x1, y1) in enumerate(order):
box_mask = torch.zeros((1, nlh, nlw), device=work_lat.device, dtype=torch.float32)
box_mask[0, y0:y1, x0:x1] = 1.0 # hard box — feather 0
bw_px = (x1 - x0) / sx
bh_px = (y1 - y0) / sy
crop_mp = max(0.05, ((bw_px + 2 * _LIR_CTX_PAD) * (bh_px + 2 * _LIR_CTX_PAD)) / 1e6)
work_lat, fresh_px = _refine_with_fine_upscaling(
model=model, vae=vae, current=work_lat,
current_pixels=work_px, mask=box_mask,
scale_x=sx, scale_y=sy,
target_mp=crop_mp * 1.02, # force the crop path — never
# the whole-canvas shortcut
max_linear=8.0, resize_method="lanczos",
context_pad_pixel=_LIR_CTX_PAD,
inpainting_mode="Refine",
reference_strength=_LIR_REF,
seed=(lir_seed + k * 9973) & 0xFFFFFFFFFFFFFFFF,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=lir_base, negative=negative,
denoise=_LIR_DENOISE,
callback=callback, disable_pbar=disable_pbar,
ov_guider=ov_guider, ov_sampler=ov_sampler, ov_sigmas=ov_sigmas,
)
if fresh_px is None:
work_px = _vae_decode(vae, work_lat)
else:
work_px = fresh_px
state["history"].append((work_lat, work_px))
if len(state["history"]) > _HISTORY_CAP:
state["history"] = state["history"][-_HISTORY_CAP:]
state["redo_stack"] = []
state["vary_candidates"] = None
state["outpaint_pending"] = None
state["click_seq"] = click_seq
state["refine_seed_at_run"] = lir_seed
current, current_pixels = work_lat, work_px
# Has the user clicked since our last execution for this node?
# Never treat a queue-hook click_seq bump as an edit in Outpaint
# mode — the canvas there is a direction picker, not a mask tool.
new_click = (
click_x >= 0
and click_y >= 0
and click_seq != state["click_seq"]
and inpainting_mode != "Outpaint"
)
# Source latent for every edit is the current cached latent, so all
# paths build ON TOP of the previous result:
# - normal clicks / paints iterate on the latest image
# - Persistent Mask Queue presses re-run the held mask on the
# latest result with a fresh seed, so the region gradually
# morphs further with each press (change something into
# something else)
# Re-roll is the ONE exception, and it gets there for free: it
# popped the last attempt above, so `current` is now that attempt's
# PRE-edit base — re-running from here gives a fresh variation on the
# ORIGINAL image, swapped in for the last attempt instead of stacked.
if new_click or reroll_now or vary_now:
# A re-roll re-runs the same mask widgets from the popped pre-
# edit base (current), so it flows through this same block.
#
# Vary ×4 also re-runs the same mask widgets, but from the
# UN-popped pre-edit base: history stays intact (the last
# attempt remains history[-1]) and only the sampling source is
# redirected to history[-2]. Nothing commits until the user
# picks a candidate, so cancelling the chooser costs nothing.
if vary_now:
hist_prev = state["history"][-2]
if isinstance(hist_prev, tuple):
current, current_pixels = hist_prev
else:
current, current_pixels = hist_prev, None
# Pixel → latent conversion. The JS sends us the actual image
# dimensions (image_w, image_h) so we can derive the per-axis
# scale dynamically rather than hardcoding a VAE downscale
# factor — that breaks for FLUX 2 (16×) vs FLUX 1 / SDXL (8×).
# Fall back to 8× only if image dims weren't provided (shouldn't
# happen in normal use).
latent_h = current.shape[-2]
latent_w = current.shape[-1]
if image_w > 0 and image_h > 0:
scale_x = latent_w / image_w
scale_y = latent_h / image_h
else:
# Fallback for headless tests / direct node use without the
# JS widget populating image_w/h. 8× is the most common VAE
# downscale (FLUX 1, SDXL, SD1.5) but breaks for FLUX 2 (16×).
scale_x = scale_y = 1.0 / 8.0
print("[Angelo] warning: image_w/h not set by JS; "
"falling back to 8x VAE assumption — may be wrong for FLUX 2")
# Geometric mean of the two axis scales (#28, from @KursatAs):
# using scale_x alone maps the on-screen circle to an ellipse in
# latent space whenever the axes scale differently (non-/16
# image dims), refining a taller/wider region than was painted.
# The geometric mean keeps the latent circle area-true on both
# portrait and landscape images.
scale_geom = math.sqrt(scale_x * scale_y)
r_latent = max(1.0, click_radius * scale_geom)
sigma_latent = (feather_radius * scale_geom) if feather_radius > 0 else 0.0
# Build the mask. Sources of mask shape, in priority:
# 1. Smart Guided Inpaint: full-image (no region — the whole
# image is edited, location comes from the prompt prefix).
# 2. Smart Inpaint: a single rectangle from rect_points.
# 3. Refine + paint_mode + stroke points: union of brush
# circles along the drag path.
# 4. Refine single-click: one circle at (click_x, click_y).
if inpainting_mode == "Smart Guided Inpaint":
mask = torch.ones((1, latent_h, latent_w),
device=current.device, dtype=torch.float32)
elif inpainting_mode == "Smart Inpaint":
rect = _parse_rect_points(rect_points)
if rect is not None:
mask = _rect_mask_latent(
latent_h, latent_w, rect,
scale_x, scale_y, current.device,
)
else:
# No rectangle drawn yet — nothing to do. Caller
# already gates on click_seq change so this is the
# "user switched into Smart Inpaint without dragging
# a rect" case; fall through with an empty mask and
# let downstream noise_mask handling preserve the
# cached latent.
mask = torch.zeros((1, latent_h, latent_w),
device=current.device, dtype=torch.float32)
else:
# Refine mask sources, in priority:
# 1. a brushed touch-up raster mask (Detect Shift/Alt brush)
# 2. a confirmed segmentation silhouette (SAM 3 / YOLO)
# 3. a paint stroke (union of brush circles)
# 4. a single click circle
raster_png = (seg_mask_png or "").strip()
seg_polys = _parse_seg_polygons(seg_polygon)
stroke_pts = _parse_stroke_points(stroke_points) if paint_mode else []
if raster_png:
mask = _raster_mask_latent(
latent_h, latent_w, raster_png, current.device,
)
elif seg_polys:
mask = _polygons_mask_latent(
latent_h, latent_w, seg_polys,
scale_x, scale_y, current.device,
)
elif stroke_pts:
mask = _stroke_mask_latent(
latent_h, latent_w,
stroke_pts, r_latent,
scale_x, scale_y, current.device,
)
else:
cx_latent = click_x * scale_x
cy_latent = click_y * scale_y
mask = _circle_mask_latent_direct(
latent_h, latent_w,
cx_latent, cy_latent, r_latent,
current.device,
)
if sigma_latent > 0:
mask = _gaussian_blur_2d(mask, max(0.5, sigma_latent))
mask = mask.clamp(0.0, 1.0)
# Restore brush (#12): Refine-only — the Smart modes regenerate
# content, so "restore to base" has no meaning there and the JS
# dims the toggle. Decided before the conditioning block so a
# restore never pays for a CLIP encode it won't use.
restore_now = bool(restore_mode) and inpainting_mode == "Refine"
# Area-prompt conditioning selection. When area_prompt is on AND a
# CLIP is connected, the refine uses the AREA text ONLY and NEVER
# the main prompt — even when the Area text is empty, in which case
# we encode the empty string (→ an empty conditioning) rather than
# falling back to `positive`. This is load-bearing for the edit
# modes: the main positive can carry whole-image reference_latents
# (a Klein edit workflow's ReferenceLatent), and letting it leak in
# made an empty-Area-Prompt Smart Inpaint reproduce the whole scene.
# Negative area text is optional — falls back to the main negative
# when empty (fine for CFG=1 / distilled models that ignore it).
#
# Smart Guided Inpaint prepends a location prefix to the positive
# text (e.g. "In the top left of the image, ") so the edit model
# places the content at the chosen spot.
#
# Without a CLIP we can't encode anything, so we must use the
# already-encoded main conditioning (an unavoidable degenerate case
# — area prompts need a CLIP connected).
area_pos_text = str(area_text_positive)
if inpainting_mode == "Smart Guided Inpaint":
prefix = _GUIDED_LOCATION_PREFIXES.get(str(guided_location), "")
area_pos_text = prefix + area_pos_text
if area_prompt and clip is not None and not restore_now:
tokens_p = clip.tokenize(area_pos_text)
refine_positive = clip.encode_from_tokens_scheduled(tokens_p)
if str(area_text_negative).strip():
tokens_n = clip.tokenize(str(area_text_negative))
refine_negative = clip.encode_from_tokens_scheduled(tokens_n)
else:
refine_negative = negative
else:
refine_positive = positive
refine_negative = negative
# Sample with the mask. Use the seed widget value as-is —
# NO click_seq offset. Per-Queue variation (when persistent_mask
# is on) and per-click variation (when user wants different
# attempts on the same spot) are now controlled by seed_control:
# fixed → same seed each run, repeatable result
# randomize → seed changes after each run (via JS after-gen),
# so each Queue / click produces a different result
# increment/decrement → +1/-1 each run
# An older version of this code did `(seed + click_seq) & mask`
# to fake per-click variation in the absence of after-gen
# control. That broke "fixed means fixed" — even with the seed
# widget locked, click_seq's increment still moved the effective
# sampling seed. Now the user has explicit control.
this_seed = int(seed)
# #21: skip the live-preview callback when the user has
# disabled it via Overrides — see Sampler Mode branch above.
callback = None if disable_live_preview else latent_preview.prepare_callback(model, steps)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
# Source latent = the current cached latent for every path (see
# the note above). Re-roll already exposed the pre-edit base as
# `current` via its pop, so it needs no special-casing here.
refine_source = current
# ===== Inpainting Mode pre-processing =====
# Refine: no pre-processing — partial denoise from the
# existing latent does the work.
# Smart Inpaint: forces fine_upscaling=True up top, so its
# pre-processing (latent zero + reference_latents) lives
# inside _refine_with_fine_upscaling.
# Smart Guided Inpaint: whole-image edit through the non-
# fine-upscale path below — inject the scene as
# reference_latents so Klein's edit branch keeps the rest
# of the image faithful while applying the location-guided
# change. POSITIVE ONLY (negative reference would steer
# CFG>1 samplers away from the scene).
if inpainting_mode == "Smart Guided Inpaint":
reference_latent = refine_source.clone()
refine_positive = node_helpers.conditioning_set_values(
refine_positive, {"reference_latents": [reference_latent]}, append=True,
)
# Reference strength (Refine only): anchor identity/content from
# the current image for the first reference_strength fraction of
# the schedule (timestep-range construction — _apply_reference),
# so high-denoise refines keep the subject. Plain path anchors
# on the WHOLE current image here; the Xtra-Fine path anchors on
# its upscaled crop instead (reference_strength passed below).
# Restore never samples, so it's excluded.
ref_strength = 0.0
if inpainting_mode == "Refine" and not restore_now and bool(refine_reference):
ref_strength = max(0.0, min(1.0, float(reference_strength)))
if ref_strength > 0.0 and not fine_upscaling:
refine_positive = _apply_reference(
refine_positive, refine_source.clone(), ref_strength)
# One pass normally; four for Vary ×4. The conditioning above is
# shared across passes — only the noise seed differs per pass.
n_passes = 4 if vary_now else 1
pass_results = []
for _k in range(n_passes):
pass_seed = this_seed if _k == 0 else (
(this_seed + _k * 1000003) & 0xFFFFFFFFFFFFFFFF
)
if restore_now:
# ===== Restore brush: heal back to the base, no sampling =====
# current = mask * base + (1-mask) * current, in latent space.
# Outside the (feathered) mask the current latent is kept
# bit-exact; inside it the session base is brought back. The
# result is pushed as a normal history entry so Undo / Redo /
# Re-roll keep working unchanged. (vary_now excludes
# restore_now, so n_passes is always 1 here.)
src = state.get("source_latent")
if src is None:
h0 = state["history"][0]
src = h0[0] if isinstance(h0, tuple) else h0
if tuple(src.shape) != tuple(current.shape):
# Shouldn't happen (source_latent resets with every base
# change), but never crash an edit session over it.
print("[Angelo restore] base/current latent shape mismatch "
f"({tuple(src.shape)} vs {tuple(current.shape)}) — restore skipped")
refined = current
else:
src = src.to(device=current.device, dtype=current.dtype)
alpha = mask
while alpha.dim() < current.dim():
alpha = alpha.unsqueeze(0) # [1,H,W] → [1,1,H,W] (→ [1,1,1,H,W] for 5D)
refined = src * alpha + current * (1.0 - alpha)
refined_pixels = None
elif fine_upscaling:
refined, refined_pixels = _refine_with_fine_upscaling(
model=model, vae=vae, current=refine_source, current_pixels=current_pixels, mask=mask,
scale_x=scale_x, scale_y=scale_y,
target_mp=float(min_megapixels),
max_linear=float(max_upscale),
resize_method=str(resize_method),
context_pad_pixel=int(fine_context_pad),
inpainting_mode=str(inpainting_mode),
reference_strength=ref_strength,
seed=pass_seed, steps=steps, cfg=cfg,
sampler_name=sampler_name, scheduler=scheduler,
positive=refine_positive, negative=refine_negative,
denoise=denoise, callback=callback, disable_pbar=disable_pbar,
# #8: pass the custom-sampler trio through so the
# internal sample call inside Fine Upscale uses the
# same dispatch logic.
ov_guider=ov_guider, ov_sampler=ov_sampler, ov_sigmas=ov_sigmas,
)
else:
noise = comfy.sample.prepare_noise(refine_source, pass_seed, None)
# #8: dispatched via _do_sample, see Sampler Mode note.
refined = _do_sample(
guider=ov_guider, sampler=ov_sampler, sigmas=ov_sigmas,
model=model, noise=noise,
steps=steps, cfg=cfg, sampler_name=sampler_name, scheduler=scheduler,
positive=refine_positive, negative=refine_negative,
source_latent=refine_source,
denoise=denoise,
noise_mask=mask,
callback=callback,
disable_pbar=disable_pbar,
seed=pass_seed,
)
refined_pixels = None
pass_results.append((refined, refined_pixels))
if vary_now:
# Stash candidates + one preview ref each for the JS chooser.
# History is untouched — the node keeps outputting the
# existing last attempt until the user picks (vary_pick).
state["vary_candidates"] = pass_results
vary_refs = []
for (lat_k, px_k) in pass_results:
if px_k is not None:
previewer = comfy_nodes.PreviewImage()
ui_k = previewer.save_images(px_k, filename_prefix="Angelo_vary")
vary_refs.append(ui_k["ui"]["images"][0])
else:
_img_k, refs_k = _decode_to_preview(vae, lat_k)
vary_refs.append(refs_k[0])
state["vary_preview_refs"] = vary_refs
state["click_seq"] = click_seq
state["refine_seed_at_run"] = int(seed)
# Reload the real last attempt for the node's outputs —
# `current` was redirected to history[-2] for sampling.
hist_last = state["history"][-1]
if isinstance(hist_last, tuple):
current, current_pixels = hist_last
else:
current, current_pixels = hist_last, None
else:
refined, refined_pixels = pass_results[0]
state["history"].append((refined, refined_pixels))
if len(state["history"]) > _HISTORY_CAP:
state["history"] = state["history"][-_HISTORY_CAP:]
# A genuine new edit (click or re-roll) invalidates the redo
# branch and any stale Vary / Outpaint stash.
state["redo_stack"] = []
state["vary_candidates"] = None
state["outpaint_pending"] = None
state["click_seq"] = click_seq
state["refine_seed_at_run"] = int(seed)
current = refined
current_pixels = refined_pixels
out_latent = {"samples": current}
ui_msg = {
"Angelo_preview": [],
"Angelo_mode": ["Edit Mode"],
"Angelo_refine_seed_at_run": [
int(state["refine_seed_at_run"])
if state.get("refine_seed_at_run") is not None else int(seed)
],
}
if current_pixels is not None:
image = current_pixels
previewer = comfy_nodes.PreviewImage()
ui = previewer.save_images(image, filename_prefix="Angelo_preview")
ui_msg["Angelo_preview"] = ui["ui"]["images"]
else:
image, image_refs = _decode_to_preview(vae, current)
ui_msg["Angelo_preview"] = image_refs
# Source image (#3/#9): the session base, decoded once and cached so
# repeated edits don't re-decode it (and so it survives history[0]
# eviction under _HISTORY_CAP). source_latent is set at every base
# (re)establishment; the history[0] fallback only covers pre-existing
# in-memory state from before this feature.
source_image = state.get("source_pixels")
if source_image is None:
src_latent = state.get("source_latent")
if src_latent is None:
h0 = state["history"][0]
src_latent = h0[0] if isinstance(h0, tuple) else h0
source_image = _vae_decode(vae, src_latent)
state["source_pixels"] = source_image
# Hold-to-compare: ship a viewable ref to the BASE image so the JS
# can flash it under the '\' key. Saved to temp once per session
# (per base) and the refs cached — repeat edits reuse them.
src_refs = state.get("source_preview_refs")
if src_refs is None:
previewer = comfy_nodes.PreviewImage()
ui_src = previewer.save_images(source_image, filename_prefix="Angelo_source")
src_refs = ui_src["ui"]["images"]
state["source_preview_refs"] = src_refs
ui_msg["Angelo_source_preview"] = src_refs
# Vary ×4: ship the candidate previews exactly once (popped, not
# read) — only the run that generated them opens the JS chooser;
# later runs must not re-open it.
vary_refs_out = state.pop("vary_preview_refs", None)
if vary_refs_out:
ui_msg["Angelo_vary_previews"] = vary_refs_out
# Outpaint: same pop-once contract for the review overlay.
op_refs_out = state.pop("outpaint_preview_refs", None)
if op_refs_out:
ui_msg["Angelo_outpaint_preview"] = op_refs_out
return {"ui": ui_msg, "result": (image, out_latent, source_image)}
class AngeloOverrides:
"""Companion node: bundle non-default Angelo settings into a single
ANGELO_OVERRIDES dict, wire it into Angelo's `overrides` slot, and
Angelo applies any non-sentinel field for that run. Originally added
for #25 (sampler-config overrides). Also carries display flags like
`disable_live_preview` (#21) so a single companion node covers every
"drive this from the workflow instead of the toolbar" use case.
Sampler-config fields use sentinel defaults ("(toolbar)" for the
combo dropdowns, -1 for INT / FLOAT) meaning "don't override this
one — use Angelo's toolbar value". Display flags default to the
same behaviour Angelo has without an Overrides node connected, so
leaving them alone is a no-op."""
_SENTINEL = "(toolbar)"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"steps": ("INT", {"default": -1, "min": -1, "max": 100, "step": 1,
"tooltip": "-1 = use Angelo's toolbar value."}),
"cfg": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 30.0, "step": 0.1,
"tooltip": "-1 = use Angelo's toolbar value."}),
"sampler_name": ([cls._SENTINEL] + list(comfy.samplers.KSampler.SAMPLERS),
{"default": cls._SENTINEL,
"tooltip": "(toolbar) = use Angelo's toolbar value."}),
"scheduler": ([cls._SENTINEL] + list(comfy.samplers.KSampler.SCHEDULERS),
{"default": cls._SENTINEL,
"tooltip": "(toolbar) = use Angelo's toolbar value."}),
"disable_live_preview": ("BOOLEAN", {"default": False,
"tooltip": "When ON, suppress ComfyUI's "
"live latent preview for this "
"Angelo node only (#21). "
"Useful if the global preview "
"is squashing Angelo's editor "
"area but you still want previews "
"on your other samplers."}),
},
"optional": {
# Custom-sampler trio (#8). Wire ALL THREE from a proper
# SamplerCustomAdvanced-style chain (a GUIDER node like
# CFGGuider/BasicGuider + a KSamplerSelect SAMPLER + a
# BasicScheduler/Karras/etc. SIGMAS at denoise=1.0) to
# replace Angelo's toolbar sampler/scheduler entirely.
# When all three are wired, Angelo uses the guider path
# and the toolbar's steps/cfg/sampler_name/scheduler are
# ignored (sigmas encode the schedule; guider encodes
# cfg + the sampling math). Partial wiring (e.g. sampler
# without sigmas) silently falls through to the default.
# Angelo's per-call denoise still applies — the sigmas
# are tail-sliced to honour it (same as
# SplitSigmasDenoise). Conds (Refine vs Area Prompt vs
# Smart Inpaint reference_latents) are re-set on the
# guider per call so the wired guider is generic.
"guider": ("GUIDER",),
"sampler": ("SAMPLER",),
"sigmas": ("SIGMAS",),
},
}
RETURN_TYPES = ("ANGELO_OVERRIDES",)
RETURN_NAMES = ("overrides",)
FUNCTION = "build"
CATEGORY = "sampling/Angelo"
DESCRIPTION = (
"Bundle Angelo settings into one wire that drives them from "
"your workflow instead of the toolbar: steps / cfg / sampler / "
"scheduler (per-field opt-in via sentinel defaults), plus the "
"disable_live_preview flag for #21."
)
def build(self, steps, cfg, sampler_name, scheduler, disable_live_preview,
guider=None, sampler=None, sigmas=None):
bundle = {
"steps": steps if isinstance(steps, int) and steps >= 1 else None,
"cfg": float(cfg) if isinstance(cfg, (int, float)) and cfg >= 0 else None,
"sampler_name": sampler_name if sampler_name != self._SENTINEL else None,
"scheduler": scheduler if scheduler != self._SENTINEL else None,
"disable_live_preview": bool(disable_live_preview),
# Custom-sampler trio (#8) — None if not wired, all-or-nothing
# at use time inside _do_sample.
"guider": guider,
"sampler": sampler,
"sigmas": sigmas,
}
return (bundle,)
NODE_CLASS_MAPPINGS = {
"AngeloRefine": AngeloRefine,
"AngeloOverrides": AngeloOverrides,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AngeloRefine": "Angelo — click to refine",
"AngeloOverrides": "Angelo — Overrides",
}
Reference in New Issue View Git Blame Copy Permalink