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

feat: 2x Restore Upscale — tiled, reference-anchored creative upscaling

Browse Source 
New '⬆ 2x Restore' button (Refine mode). Pipeline: decode -> 2x
lanczos PIXEL upscale (latent upscaling smears — the Xtra-Fine
lesson) -> re-encode -> the Quick Photo Refine recipe run over
overlapping ~1MP tiles, each at denoise 1.0 anchored to ITS OWN
pre-refine latent via reference_latents -> feathered pixel composite
(separable ramp weights, normalised) -> final encode.

The per-tile anchor is the load-bearing idea: it prevents the two
classic tiled-upscaler failures (per-tile hallucination + seam
drift) without a tile ControlNet — adjacent tiles agree because
they're anchored to the same underlying image.

Dimension change = the Outpaint review/accept flow, reused: shared
overlay (title + Try-again dispatch via _AngeloPendingOp), Accept
installs the result as a fresh session base.

Big-canvas safety: the model never SAMPLES a latent beyond ~1MP —
during the upscale (tiles) and after it: ✨ Quick Photo Refine now
auto-routes through the tiled engine at scale 1.0 above 1.6MP, so
the magic button stays usable on upscaled canvases. Post-accept
toast steers spot edits to Xtra-Fine (which crops, so it's
size-independent). Full-size VAE encode/decode remain (comfy
auto-tiles VAE on OOM).

New hidden widget upscale_seq declared LAST + anchored in all three
state-replacement sites.
This commit is contained in:
shootthesound
2026-06-10 15:45:19 +01:00
parent 4c47cd6e6d
commit ac4d60448e
2 changed files with 292 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
angelo_nodes.py
+235 -25
View File
@@ -842,6 +842,16 @@ _QUICK_REFINE_PROMPT_QWEN = ("lightly restore this old photo, remove dust and sc
_QUICK_REFINE_DENOISE = 1.0
_QUICK_REFINE_REF = 1.0
# 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 = 1024
_TILE_OVERLAP_PX = 128
_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
@@ -877,6 +887,137 @@ def _apply_reference(positive, ref_latent: torch.Tensor, strength: float):
return head + tail
def _tile_positions(total: int, tile: int, stride: int) -> list[int]:
"""Start offsets covering [0, total) with `tile`-sized windows at
`stride` spacing, the last window pulled flush to the end so the far
edge is always covered exactly once (no sliver tiles)."""
if total <= tile:
return [0]
pos = list(range(0, total - tile + 1, stride))
if pos[-1] + tile < total:
pos.append(total - tile)
return pos
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,
) -> 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)
ov_lx = max(1, _TILE_OVERLAP_PX // ppl_x)
ov_ly = max(1, _TILE_OVERLAP_PX // ppl_y)
xs = _tile_positions(nlw, tile_lx, tile_lx - ov_lx)
ys = _tile_positions(nlh, tile_ly, 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"overlap {_TILE_OVERLAP_PX}px)")
out = torch.zeros_like(big_pixels)
wsum = torch.zeros((1, new_H, new_W, 1), dtype=big_pixels.dtype, device=big_pixels.device)
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)
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(), _QUICK_REFINE_REF)
tile_seed = (seed + idx * 1000003) & 0xFFFFFFFFFFFFFFFF
noise = comfy.sample.prepare_noise(tile_lat, tile_seed, None)
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=_QUICK_REFINE_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, _TILE_OVERLAP_PX, big_pixels.device)
wx = _ramp(ptw, _TILE_OVERLAP_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
@@ -1553,6 +1694,14 @@ class AngeloRefine:
"anchored and free predictions. 0 = "
"off, 1 = fully anchored. Set via the "
"Ref box on the toolbar."}),
# 2× Restore Upscale: the ⬆ button bumps this. Pixel-space
# 2× lanczos + re-encode, then the tiled restore engine
# (_tiled_restore_pass) over overlapping ~1MP tiles. Result
# is STASHED for review (reuses the outpaint pending-base
# flow — Accept installs it as a fresh session base, since
# the dimensions changed). Declared LAST.
"upscale_seq": ("INT", {"default": 0, "min": 0, "max": 0x7FFFFFFF}),
},
"optional": {
# CLIP / text encoder for the Area Prompt. Optional —
@@ -1660,6 +1809,7 @@ class AngeloRefine:
refine_reference=False,
quick_refine_seq=0,
reference_strength=0.0,
upscale_seq=0,
latent=None,
clip=None,
overrides=None,
@@ -1850,6 +2000,7 @@ class AngeloRefine:
"outpaint_seq": outpaint_seq,
"outpaint_accept_seq": outpaint_accept_seq,
"quick_refine_seq": quick_refine_seq,
"upscale_seq": upscale_seq,
"fingerprint": incoming_fp,
"sampler_seed_at_run": int(sampler_seed),
"loaded_seq": loaded_seq,
@@ -1927,6 +2078,7 @@ class AngeloRefine:
"outpaint_seq": outpaint_seq,
"outpaint_accept_seq": outpaint_accept_seq,
"quick_refine_seq": quick_refine_seq,
"upscale_seq": upscale_seq,
"fingerprint": incoming_fp,
"loaded_seq": loaded_seq,
# Source image (#3/#9): capture the base once, independent of
@@ -2022,6 +2174,7 @@ class AngeloRefine:
"outpaint_seq": outpaint_seq,
"outpaint_accept_seq": outpaint_accept_seq,
"quick_refine_seq": quick_refine_seq,
"upscale_seq": upscale_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.
@@ -2185,36 +2338,54 @@ class AngeloRefine:
else _QUICK_REFINE_PROMPT)
if clip is not None:
tokens_q = clip.tokenize(qr_prompt)
qr_positive = clip.encode_from_tokens_scheduled(tokens_q)
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_positive = positive
# 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_strength = _QUICK_REFINE_REF
qr_positive = _apply_reference(qr_positive, current.clone(), qr_strength)
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
print(f"[Angelo quick-refine] whole-canvas restoration pass, "
f"denoise={_QUICK_REFINE_DENOISE}, ref={qr_strength}, 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,
)
state["history"].append((qr_refined, None))
# 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"] = []
@@ -2223,7 +2394,46 @@ class AngeloRefine:
state["click_seq"] = click_seq
state["refine_seed_at_run"] = qr_seed
current = qr_refined
current_pixels = None
current_pixels = qr_pixels
# ===== 2× Restore Upscale: tiled, reference-anchored upscaling =====
# Pixel 2× + re-encode, then the tiled restore engine. The result
# changes the canvas DIMENSIONS, so it goes through the same
# review/accept flow as Outpaint: stash + overlay, Accept installs
# it as a fresh session base (history resets), Cancel costs nothing.
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_prompt = (_QUICK_REFINE_PROMPT_QWEN if current.dim() == 5
else _QUICK_REFINE_PROMPT)
if clip is not None:
tokens_u = clip.tokenize(up_prompt)
up_base = clip.encode_from_tokens_scheduled(tokens_u)
else:
up_base = positive
up_seed = int(seed)
callback = None if disable_live_preview else latent_preview.prepare_callback(model, steps)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
up_latent, up_pixels, _n = _tiled_restore_pass(
model=model, vae=vae,
current=current, current_pixels=current_pixels,
scale=2.0,
positive_base=up_base, negative=negative,
seed=up_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,
)
previewer = comfy_nodes.PreviewImage()
ui_up = previewer.save_images(up_pixels, filename_prefix="Angelo_upscale")
state["outpaint_pending"] = (up_latent, up_pixels)
state["outpaint_preview_refs"] = ui_up["ui"]["images"]
state["click_seq"] = click_seq
state["refine_seed_at_run"] = up_seed
# Has the user clicked since our last execution for this node?
# Never treat a queue-hook click_seq bump as an edit in Outpaint
web/angelo.js
+57 -6
View File
@@ -604,6 +604,21 @@ function attachPreviewCanvas(node) {
row1.appendChild(quickFixBtn);
node._AngeloQuickFixBtn = quickFixBtn;
const upscaleBtn = makeActionButton("⬆ 2× Restore", () => triggerRestoreUpscale(node), "quickfix");
upscaleBtn.title = "2× Restore Upscale — creative upscaling built from the restoration recipe. "
+ "The image is upscaled 2× in pixel space (lanczos), re-encoded, then restored tile by "
+ "tile: overlapping ~1MP squares, each at denoise 1.0 anchored to its own content as a "
+ "reference, composited with feathered seams. The anchor is what stops tiles "
+ "hallucinating new content and keeps neighbours agreeing — the job tile-ControlNets do "
+ "in other pipelines.\n\n"
+ "The result is shown in a review overlay first — Accept commits it as a NEW session "
+ "base (history resets, like Outpaint), Try again re-rolls it, Cancel costs nothing. "
+ "After accepting: use Xtra-Fine for spot edits (it crops, so it stays fast at any "
+ "canvas size), and ✨ Quick Photo Refine auto-tiles on big canvases.\n\n"
+ "Edit models + CLIP; Refine mode only. A 1MP source ≈ 9 tiles ≈ 15–20s on Klein.";
row1.appendChild(upscaleBtn);
node._AngeloUpscaleBtn = upscaleBtn;
row1.appendChild(makeSeparator());
const persistentMaskToggle = makeToggleButton("Persistent Mask", () => {
@@ -1458,6 +1473,7 @@ function attachPreviewCanvas(node) {
opEsc.style.cssText = "font-weight:normal; font-size:10px; color:#8aa;";
opHeader.appendChild(opTitle);
opHeader.appendChild(opEsc);
node._AngeloOutpaintTitle = opTitle;
const opImgWrap = document.createElement("div");
opImgWrap.style.cssText = "flex:1 1 auto; min-height:0; display:flex; "
+ "align-items:center; justify-content:center;";
@@ -3987,6 +4003,25 @@ function triggerQuickPhotoRefine(node) {
queuePrompt();
}
// 2× Restore Upscale: tiled, reference-anchored creative upscaling. The
// result is a dimension change, so it rides the same review/accept
// (pending-base) flow as Outpaint — _AngeloPendingOp tells the shared
// overlay which op "Try again" should re-fire.
function triggerRestoreUpscale(node) {
if (isAnySmartMode(node) || isOutpaintMode(node)) return; // dimmed there anyway
if (!node._AngeloImg) {
_angeloToast("Generate or load an image first");
return;
}
const ws = findWidget(node, "upscale_seq");
if (!ws) return;
node._AngeloPendingOp = "upscale";
setWidget(ws, ((ws.value || 0) + 1) & 0x7FFFFFFF);
_angeloToast("⬆ 2× Restore Upscale — restoring tile by tile…");
dbg("queue restore upscale", { upscale_seq: ws.value });
queuePrompt();
}
// ===== Outpaint — directional canvas extension with review-before-commit =====
function isOutpaintMode(node) {
@@ -4005,6 +4040,7 @@ function triggerOutpaint(node, dir) {
return;
}
setWidget(wd, dir);
node._AngeloPendingOp = "outpaint";
setWidget(ws, ((ws.value || 0) + 1) & 0x7FFFFFFF);
// The direction just changed — refresh the combined-prompt preview.
syncOutpaintPromptPreview(node);
@@ -4028,6 +4064,13 @@ function showOutpaintReview(node, ref) {
const ov = node._AngeloOutpaintOverlay;
const img = node._AngeloOutpaintImg;
if (!ov || !img) return;
// Title follows the pending op — the overlay is shared between
// Outpaint and 2× Restore Upscale.
if (node._AngeloOutpaintTitle) {
node._AngeloOutpaintTitle.textContent = node._AngeloPendingOp === "upscale"
? "2× Restore Upscale — keep it?"
: "Outpaint preview — keep it?";
}
img.src = makeViewUrl(ref);
ov.style.display = "flex";
}
@@ -4049,15 +4092,19 @@ function triggerOutpaintAccept(node) {
const wpr = findWidget(node, "outpaint_protect");
if (wpr) setWidget(wpr, "");
syncOutpaintControls(node);
_angeloToast("Committing the new canvas — this is your new base");
_angeloToast(node._AngeloPendingOp === "upscale"
? "Committing the 2× canvas — use Xtra-Fine for spot edits at this size"
: "Committing the new canvas — this is your new base");
dbg("queue outpaint accept", { outpaint_accept_seq: ws.value });
queuePrompt();
}
// Same direction + amount, fresh seed. The stale stash is simply
// overwritten by the new run.
// Same operation, fresh seed. The stale stash is simply overwritten by
// the new run. Dispatches on the pending op — the review overlay is
// shared between Outpaint and 2× Restore Upscale.
function triggerOutpaintRetry(node) {
const ws = findWidget(node, "outpaint_seq");
const seqName = node._AngeloPendingOp === "upscale" ? "upscale_seq" : "outpaint_seq";
const ws = findWidget(node, seqName);
if (!ws) return;
const wseed = findWidget(node, "seed");
if (wseed) {
@@ -4065,7 +4112,9 @@ function triggerOutpaintRetry(node) {
syncSeedInput(node);
}
setWidget(ws, ((ws.value || 0) + 1) & 0x7FFFFFFF);
_angeloToast("Trying the extension again with a fresh seed…");
_angeloToast(node._AngeloPendingOp === "upscale"
? "Re-rolling the 2× restore with a fresh seed…"
: "Trying the extension again with a fresh seed…");
queuePrompt();
}
@@ -4440,7 +4489,7 @@ function syncSmartInpaintLockedWidgets(node) {
_dimControls(node, ["_AngeloRerollBtn", "_AngeloVaryBtn"], outp);
// Quick Photo Refine is a Refine-mode action.
_dimControls(node, ["_AngeloQuickFixBtn"], anySmart || outp);
_dimControls(node, ["_AngeloQuickFixBtn", "_AngeloUpscaleBtn"], anySmart || outp);
// Outpaint row visibility + input mirrors.
syncOutpaintControls(node);
@@ -4860,6 +4909,8 @@ function hideMechanicalWidgets(node) {
"refine_reference", "reference_strength",
// Quick Photo Refine — driven by the ✨ button
"quick_refine_seq",
// 2× Restore Upscale — driven by the ⬆ button
"upscale_seq",
// Toolbar-driven (visible via the bar above the canvas)
"persistent_mask", "area_prompt", "paint_mode", "fine_upscaling",
"click_radius", "feather_radius", "denoise",