"""CADE 2.5: refined adaptive enhancer with reference clean and accumulation override. Builds on the CADE2 Beta: single clean iteration loop, optional latent-based parameter damping, CLIP-based reference clean, and per-run SageAttention accumulation override. """ from __future__ import annotations # moved/renamed module: mg_cade25 import torch import os import numpy as np import torch.nn.functional as F import traceback import nodes import comfy.model_management as model_management from .mg_adaptive import AdaptiveSamplerHelper from .mg_zesmart_sampler_v1_1 import _build_hybrid_sigmas import comfy.sample as _sample import comfy.samplers as _samplers import comfy.utils as _utils from .mg_upscale_module import MagicUpscaleModule, clear_gpu_and_ram_cache from .mg_controlfusion import _build_depth_map as _cf_build_depth_map from .mg_ids import IntelligentDetailStabilizer from .. import mg_sagpu_attention as sa_patch # FDG/NAG experimental paths removed for now; keeping code lean # Lazy CLIPSeg cache _CLIPSEG_MODEL = None _CLIPSEG_PROC = None _CLIPSEG_DEV = "cpu" _CLIPSEG_FORCE_CPU = True # pin CLIPSeg to CPU to avoid device drift # Cooperative cancel sentinel: set in callbacks when user interrupts _MG_CANCEL_REQUESTED = False # Per-iteration spatial guidance mask (B,1,H,W) in [0,1]; used by cfg_func when enabled # Kept for potential future use with non-ONNX masks (e.g., CLIPSeg/ControlFusion), # but not set by this node since ONNX paths are removed. CURRENT_ONNX_MASK_BCHW = None # ONNX runtime initialization removed def _try_init_clipseg(): """Lazy-load CLIPSeg processor + model and choose device. Returns True on success. """ global _CLIPSEG_MODEL, _CLIPSEG_PROC, _CLIPSEG_DEV if (_CLIPSEG_MODEL is not None) and (_CLIPSEG_PROC is not None): return True try: from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation # type: ignore except Exception: if not globals().get("_CLIPSEG_WARNED", False): print("[CADE2.5][CLIPSeg] transformers not available; CLIPSeg disabled.") globals()["_CLIPSEG_WARNED"] = True return False try: _CLIPSEG_PROC = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") _CLIPSEG_MODEL = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") if _CLIPSEG_FORCE_CPU: _CLIPSEG_DEV = "cpu" else: _CLIPSEG_DEV = "cuda" if torch.cuda.is_available() else "cpu" _CLIPSEG_MODEL = _CLIPSEG_MODEL.to(_CLIPSEG_DEV) _CLIPSEG_MODEL.eval() return True except Exception as e: print(f"[CADE2.5][CLIPSeg] failed to load model: {e}") return False def _clipseg_build_mask(image_bhwc: torch.Tensor, text: str, preview: int = 224, threshold: float = 0.4, blur: float = 7.0, dilate: int = 4, gain: float = 1.0, ref_embed: torch.Tensor | None = None, clip_vision=None, ref_threshold: float = 0.03) -> torch.Tensor | None: """Return BHWC single-channel mask [0,1] from CLIPSeg. - Uses cached CLIPSeg model; gracefully returns None on failure. - Applies optional threshold/blur/dilate and scaling gain. - If clip_vision + ref_embed provided, gates mask by CLIP-Vision distance. """ if not text or not isinstance(text, str): return None if not _try_init_clipseg(): return None try: # Prepare preview image (CPU PIL) target = int(max(16, min(1024, preview))) img = image_bhwc.detach().to('cpu') if img.ndim == 5: # squeeze depth if present if img.shape[1] == 1: img = img[:, 0] else: img = img[:, 0] B, H, W, C = img.shape x = img[0].movedim(-1, 0).unsqueeze(0) # 1,C,H,W x = F.interpolate(x, size=(target, target), mode='bilinear', align_corners=False) x = x.clamp(0, 1) arr = (x[0].movedim(0, -1).numpy() * 255.0).astype('uint8') from PIL import Image # lazy import pil_img = Image.fromarray(arr) # Run CLIPSeg import re prompts = [t.strip() for t in re.split(r"[\|,;\n]+", text) if t.strip()] if not prompts: prompts = [text.strip()] prompts = prompts[:8] inputs = _CLIPSEG_PROC(text=prompts, images=[pil_img] * len(prompts), return_tensors="pt") inputs = {k: v.to(_CLIPSEG_DEV) for k, v in inputs.items()} with torch.inference_mode(): outputs = _CLIPSEG_MODEL(**inputs) # type: ignore # logits: [N, H', W'] for N prompts logits = outputs.logits # [N,h,w] if logits.ndim == 2: logits = logits.unsqueeze(0) prob = torch.sigmoid(logits) # [N,h,w] # Soft-OR fuse across prompts prob = 1.0 - torch.prod(1.0 - prob.clamp(0, 1), dim=0, keepdim=True) # [1,h,w] prob = prob.unsqueeze(1) # [1,1,h,w] # Resize to original image size prob = F.interpolate(prob, size=(H, W), mode='bilinear', align_corners=False) m = prob[0, 0].to(dtype=image_bhwc.dtype, device=image_bhwc.device) # Threshold + blur (approx) if threshold > 0.0: m = torch.where(m > float(threshold), m, torch.zeros_like(m)) # Gaussian blur via our depthwise helper if blur > 0.0: rad = int(max(1, min(7, round(blur)))) m = _gaussian_blur_nchw(m.unsqueeze(0).unsqueeze(0), sigma=float(max(0.5, blur)), radius=rad)[0, 0] # Dilation via max-pool if int(dilate) > 0: k = int(dilate) * 2 + 1 p = int(dilate) m = F.max_pool2d(m.unsqueeze(0).unsqueeze(0), kernel_size=k, stride=1, padding=p)[0, 0] # Optional CLIP-Vision gating by reference distance if (clip_vision is not None) and (ref_embed is not None): try: cur = _encode_clip_image(image_bhwc, clip_vision, target_res=224) dist = _clip_cosine_distance(cur, ref_embed) if dist > float(ref_threshold): # up to +50% gain if distance exceeds the reference threshold gate = 1.0 + min(0.5, (dist - float(ref_threshold)) * 4.0) m = m * gate except Exception: pass m = (m * float(max(0.0, gain))).clamp(0, 1) out_mask = m.unsqueeze(0).unsqueeze(-1) # BHWC with B=1,C=1 # Best-effort release of temporaries to reduce RAM peak try: del inputs except Exception: pass try: del outputs except Exception: pass try: del logits except Exception: pass try: del prob except Exception: pass try: del pil_img except Exception: pass try: del arr except Exception: pass try: del x except Exception: pass try: del img except Exception: pass return out_mask except Exception as e: if not globals().get("_CLIPSEG_WARNED", False): print(f"[CADE2.5][CLIPSeg] mask failed: {e}") globals()["_CLIPSEG_WARNED"] = True return None def _np_to_mask_tensor(np_map: np.ndarray, out_h: int, out_w: int, device, dtype): """Convert numpy heatmap [H,W] or [1,H,W] or [H,W,1] to BHWC torch mask with B=1 and resize to out_h,out_w.""" if np_map.ndim == 3: np_map = np_map.reshape(np_map.shape[-2], np_map.shape[-1]) if (np_map.shape[0] == 1) else np_map.squeeze() if np_map.ndim != 2: return None t = torch.from_numpy(np_map.astype(np.float32)) t = t.clamp_min(0.0) t = t.unsqueeze(0).unsqueeze(0) # B=1,C=1,H,W t = F.interpolate(t, size=(out_h, out_w), mode="bilinear", align_corners=False) t = t.permute(0, 2, 3, 1).to(device=device, dtype=dtype) # B,H,W,C return t.clamp(0, 1) def _mask_to_like(mask_bhw1: torch.Tensor, like_bhwc: torch.Tensor) -> torch.Tensor: try: if mask_bhw1 is None or like_bhwc is None: return mask_bhw1 if mask_bhw1.ndim != 4 or like_bhwc.ndim != 4: return mask_bhw1 _, Ht, Wt, _ = like_bhwc.shape _, Hm, Wm, _ = mask_bhw1.shape if (Hm, Wm) == (Ht, Wt): return mask_bhw1 m = mask_bhw1.movedim(-1, 1) m = F.interpolate(m, size=(Ht, Wt), mode='bilinear', align_corners=False) return m.movedim(1, -1).clamp(0, 1) except Exception: return mask_bhw1 def _align_mask_pair(a_bhw1: torch.Tensor, b_bhw1: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: try: if a_bhw1 is None or b_bhw1 is None: return a_bhw1, b_bhw1 if a_bhw1.ndim != 4 or b_bhw1.ndim != 4: return a_bhw1, b_bhw1 _, Ha, Wa, _ = a_bhw1.shape _, Hb, Wb, _ = b_bhw1.shape if (Ha, Wa) == (Hb, Wb): return a_bhw1, b_bhw1 m = b_bhw1.movedim(-1, 1) m = F.interpolate(m, size=(Ha, Wa), mode='bilinear', align_corners=False) return a_bhw1, m.movedim(1, -1).clamp(0, 1) except Exception: return a_bhw1, b_bhw1 # --- Firefly/Hot-pixel remover (image space, BHWC in 0..1) --- def _median_pool3x3_bhwc(img_bhwc: torch.Tensor) -> torch.Tensor: B, H, W, C = img_bhwc.shape x = img_bhwc.permute(0, 3, 1, 2) # B,C,H,W unfold = F.unfold(x, kernel_size=3, padding=1) # B, 9*C, H*W unfold = unfold.view(B, x.shape[1], 9, H, W) # B,C,9,H,W med, _ = torch.median(unfold, dim=2) # B,C,H,W return med.permute(0, 2, 3, 1) # B,H,W,C def _despeckle_fireflies(img_bhwc: torch.Tensor, thr: float = 0.985, max_iso: float | None = None, grad_gate: float = 0.25) -> torch.Tensor: try: dev, dt = img_bhwc.device, img_bhwc.dtype B, H, W, C = img_bhwc.shape s = max(H, W) / 1024.0 k = 3 if s <= 1.1 else (5 if s <= 2.0 else 7) pad = k // 2 lum = (0.2126 * img_bhwc[..., 0] + 0.7152 * img_bhwc[..., 1] + 0.0722 * img_bhwc[..., 2]).to(device=dev, dtype=dt) try: q = float(torch.quantile(lum.reshape(-1), 0.9995).item()) thr_eff = max(float(thr), min(0.997, q)) except Exception: thr_eff = float(thr) # S/V based candidate: white, low saturation R, G, Bc = img_bhwc[..., 0], img_bhwc[..., 1], img_bhwc[..., 2] V = torch.maximum(R, torch.maximum(G, Bc)) mi = torch.minimum(R, torch.minimum(G, Bc)) S = 1.0 - (mi / (V + 1e-6)) v_thr = max(0.985, thr_eff) s_thr = 0.06 cand = (V > v_thr) & (S < s_thr) # gradient gate kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=dev, dtype=dt).view(1, 1, 3, 3) ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=dev, dtype=dt).view(1, 1, 3, 3) gx = F.conv2d(lum.unsqueeze(1), kx, padding=1) gy = F.conv2d(lum.unsqueeze(1), ky, padding=1) grad = torch.sqrt(gx * gx + gy * gy).squeeze(1) safe_gate = float(grad_gate) * (k / 3.0) ** 0.5 cand = cand & (grad < safe_gate) if cand.any(): try: import cv2, numpy as _np masks = [] for b in range(cand.shape[0]): msk = cand[b].detach().to('cpu').numpy().astype('uint8') * 255 num, labels, stats, _ = cv2.connectedComponentsWithStats(msk, connectivity=8) rem = _np.zeros_like(msk, dtype='uint8') area_max = int(max(3, round((k * k) * 0.6))) for lbl in range(1, num): area = stats[lbl, cv2.CC_STAT_AREA] if area <= area_max: rem[labels == lbl] = 255 masks.append(torch.from_numpy(rem > 0)) rm = torch.stack(masks, dim=0).to(device=dev) rm = rm.unsqueeze(-1) if rm.any(): med = _median_pool3x3_bhwc(img_bhwc) return torch.where(rm, med, img_bhwc) except Exception: pass # Fallback: density isolation bright = (img_bhwc.min(dim=-1).values > v_thr) dens = F.avg_pool2d(bright.float().unsqueeze(1), k, 1, pad).squeeze(1) max_iso_eff = (2.0 / (k * k)) if (max_iso is None) else float(max_iso) iso = bright & (dens < max_iso_eff) & (grad < safe_gate) if not iso.any(): return img_bhwc med = _median_pool3x3_bhwc(img_bhwc) return torch.where(iso.unsqueeze(-1), med, img_bhwc) except Exception: return img_bhwc def _try_heatmap_from_outputs(outputs: list, preview_hw: tuple[int, int]): """Return [H,W] heatmap from model outputs if possible. Supports: - Segmentation logits/probabilities (NCHW / NHWC) - Keypoints arrays -> gaussian disks on points - Bounding boxes -> soft rectangles """ if not outputs: return None Ht, Wt = int(preview_hw[0]), int(preview_hw[1]) def to_float(arr): if arr.dtype not in (np.float32, np.float64): try: arr = arr.astype(np.float32) except Exception: return None return arr def sigmoid(x): return 1.0 / (1.0 + np.exp(-x)) # 1) Prefer any spatial heatmap first for out in outputs: try: arr = np.asarray(out) except Exception: continue arr = to_float(arr) if arr is None: continue if arr.ndim == 4: n, a, b, c = arr.shape if c <= 4 and a >= 8 and b >= 8: if c == 1: hm = sigmoid(arr[0, :, :, 0]) if np.max(np.abs(arr)) > 1.5 else arr[0, :, :, 0] else: ex = np.exp(arr[0] - np.max(arr[0], axis=-1, keepdims=True)) prob = ex / np.clip(ex.sum(axis=-1, keepdims=True), 1e-6, None) hm = 1.0 - prob[..., 0] if prob.shape[-1] > 1 else prob[..., 0] return hm.astype(np.float32) else: if a == 1: ch = arr[0, 0] hm = sigmoid(ch) if np.max(np.abs(ch)) > 1.5 else ch return hm.astype(np.float32) else: x = arr[0] x = x - np.max(x, axis=0, keepdims=True) ex = np.exp(x) prob = ex / np.clip(np.sum(ex, axis=0, keepdims=True), 1e-6, None) bg = prob[0] if prob.shape[0] > 1 else prob[0] hm = 1.0 - bg return hm.astype(np.float32) if arr.ndim == 3: if arr.shape[0] == 1 and arr.shape[1] >= 8 and arr.shape[2] >= 8: return arr[0].astype(np.float32) if arr.ndim == 2 and arr.shape[0] >= 8 and arr.shape[1] >= 8: return arr.astype(np.float32) # 2) Try keypoints and boxes heat = np.zeros((Ht, Wt), dtype=np.float32) def draw_gaussian(hm, cx, cy, sigma=2.5, amp=1.0): r = max(1, int(3 * sigma)) xs = np.arange(-r, r + 1, dtype=np.float32) ys = np.arange(-r, r + 1, dtype=np.float32) gx = np.exp(-(xs**2) / (2 * sigma * sigma)) gy = np.exp(-(ys**2) / (2 * sigma * sigma)) g = np.outer(gy, gx) * float(amp) x0 = int(round(cx)) - r y0 = int(round(cy)) - r x1 = x0 + g.shape[1] y1 = y0 + g.shape[0] if x1 < 0 or y1 < 0 or x0 >= Wt or y0 >= Ht: return xs0 = max(0, x0) ys0 = max(0, y0) xs1 = min(Wt, x1) ys1 = min(Ht, y1) gx0 = xs0 - x0 gy0 = ys0 - y0 gx1 = gx0 + (xs1 - xs0) gy1 = gy0 + (ys1 - ys0) hm[ys0:ys1, xs0:xs1] = np.maximum(hm[ys0:ys1, xs0:xs1], g[gy0:gy1, gx0:gx1]) def draw_soft_rect(hm, x0, y0, x1, y1, edge=3.0): x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1) if x1 <= 0 or y1 <= 0 or x0 >= Wt or y0 >= Ht: return xs0 = max(0, min(x0, x1)) ys0 = max(0, min(y0, y1)) xs1 = min(Wt, max(x0, x1)) ys1 = min(Ht, max(y0, y1)) if xs1 - xs0 <= 0 or ys1 - ys0 <= 0: return hm[ys0:ys1, xs0:xs1] = np.maximum(hm[ys0:ys1, xs0:xs1], 1.0) # feather edges with simple blur-like falloff if edge > 0: rad = int(edge) if rad > 0: # quick separable triangle filter line = np.linspace(0, 1, rad + 1, dtype=np.float32)[1:] for d in range(1, rad + 1): w = line[d - 1] if ys0 - d >= 0: hm[ys0 - d:ys0, xs0:xs1] = np.maximum(hm[ys0 - d:ys0, xs0:xs1], w) if ys1 + d <= Ht: hm[ys1:ys1 + d, xs0:xs1] = np.maximum(hm[ys1:ys1 + d, xs0:xs1], w) if xs0 - d >= 0: hm[max(0, ys0 - d):min(Ht, ys1 + d), xs0 - d:xs0] = np.maximum( hm[max(0, ys0 - d):min(Ht, ys1 + d), xs0 - d:xs0], w) if xs1 + d <= Wt: hm[max(0, ys0 - d):min(Ht, ys1 + d), xs1:xs1 + d] = np.maximum( hm[max(0, ys0 - d):min(Ht, ys1 + d), xs1:xs1 + d], w) # Inspect outputs to find plausible keypoints/boxes for out in outputs: try: arr = np.asarray(out) except Exception: continue arr = to_float(arr) if arr is None: continue a = arr # Squeeze batch dims like [1,N,4] -> [N,4] while a.ndim > 2 and a.shape[0] == 1: a = np.squeeze(a, axis=0) # Keypoints: [N,2] or [N,3] or [K, N, 2/3] (relax N limit; subsample if huge) if a.ndim == 2 and a.shape[-1] in (2, 3): pts = a elif a.ndim == 3 and a.shape[-1] in (2, 3): pts = a.reshape(-1, a.shape[-1]) else: pts = None if pts is not None: # Coordinates range guess: if max>1.2 -> absolute; else normalized maxv = float(np.nanmax(np.abs(pts[:, :2]))) if pts.size else 0.0 for px, py, *rest in pts: if np.isnan(px) or np.isnan(py): continue if maxv <= 1.2: cx = float(px) * (Wt - 1) cy = float(py) * (Ht - 1) else: cx = float(px) cy = float(py) base_sig = max(1.5, min(Ht, Wt) / 128.0) if _ONNX_KPTS_ENABLE: draw_gaussian(heat, cx, cy, sigma=base_sig * float(_ONNX_KPTS_SIGMA), amp=float(_ONNX_KPTS_GAIN)) else: draw_gaussian(heat, cx, cy, sigma=base_sig) continue # Wholebody-style packed keypoints: [N, K*3] with triples (x,y,conf) if _ONNX_KPTS_ENABLE and a.ndim == 2 and a.shape[-1] >= 6 and (a.shape[-1] % 3) == 0: K = a.shape[-1] // 3 if K >= 5 and K <= 256: # Guess coordinate range once with np.errstate(invalid='ignore'): maxv = float(np.nanmax(np.abs(a[:, :2]))) if a.size else 0.0 for i in range(a.shape[0]): row = a[i] kp = row.reshape(K, 3) for (px, py, pc) in kp: if np.isnan(px) or np.isnan(py): continue if np.isfinite(pc) and pc < float(_ONNX_KPTS_CONF): continue if maxv <= 1.2: cx = float(px) * (Wt - 1) cy = float(py) * (Ht - 1) else: cx = float(px) cy = float(py) base_sig = max(1.0, min(Ht, Wt) / 128.0) draw_gaussian(heat, cx, cy, sigma=base_sig * float(_ONNX_KPTS_SIGMA), amp=float(_ONNX_KPTS_GAIN)) continue # Boxes: [N,4+] (x0,y0,x1,y1) or [N, (x,y,w,h, [conf, ...])]; relax N limit (handle YOLO-style outputs) if a.ndim == 2 and a.shape[-1] >= 4: boxes = a elif a.ndim == 3 and a.shape[-1] >= 4: # choose the smallest first two dims as N if a.shape[0] == 1: boxes = a.reshape(-1, a.shape[-1]) else: boxes = a.reshape(-1, a.shape[-1]) else: boxes = None if boxes is not None: # Optional score gating (try to find a confidence column) score = None if boxes.shape[-1] >= 6: score = boxes[:, 4] # if classes follow, mix in best class prob try: score = score * np.max(boxes[:, 5:], axis=-1) except Exception: pass elif boxes.shape[-1] == 5: score = boxes[:, 4] # Keep top-K by score if available if score is not None: try: order = np.argsort(-score) keep = order[: min(64, order.shape[0])] boxes = boxes[keep] score = score[keep] except Exception: score = None xy = boxes[:, :4] maxv = float(np.nanmax(np.abs(xy))) if xy.size else 0.0 if maxv <= 1.2: x0 = xy[:, 0] * (Wt - 1) y0 = xy[:, 1] * (Ht - 1) x1 = xy[:, 2] * (Wt - 1) y1 = xy[:, 3] * (Ht - 1) else: x0, y0, x1, y1 = xy[:, 0], xy[:, 1], xy[:, 2], xy[:, 3] # Heuristic: if many boxes are inverted, treat as [x,y,w,h] invalid = np.sum((x1 <= x0) | (y1 <= y0)) if invalid > 0.5 * x0.shape[0]: x, y, w, h = x0, y0, x1, y1 x0 = x - w * 0.5 y0 = y - h * 0.5 x1 = x + w * 0.5 y1 = y + h * 0.5 for i in range(x0.shape[0]): if score is not None and np.isfinite(score[i]) and score[i] < 0.2: continue draw_soft_rect(heat, x0[i], y0[i], x1[i], y1[i], edge=3.0) # Embedded keypoints in YOLO-style rows: try to parse trailing triples (x,y,conf) if _ONNX_KPTS_ENABLE and boxes.shape[-1] > 6: D = boxes.shape[-1] for i in range(boxes.shape[0]): row = boxes[i] parsed = False # try [xyxy, conf, cls, kpts] or [xyxy, conf, kpts] or [xyxy, kpts] for offset in (6, 5, 4): t = D - offset if t >= 6 and t % 3 == 0: k = t // 3 kp = row[offset:offset + 3 * k].reshape(k, 3) parsed = True break if not parsed: continue for (px, py, pc) in kp: if np.isnan(px) or np.isnan(py): continue if pc < float(_ONNX_KPTS_CONF): continue if maxv <= 1.2: cx = float(px) * (Wt - 1) cy = float(py) * (Ht - 1) else: cx = float(px) cy = float(py) base_sig = max(1.0, min(Ht, Wt) / 128.0) draw_gaussian(heat, cx, cy, sigma=base_sig * float(_ONNX_KPTS_SIGMA), amp=float(_ONNX_KPTS_GAIN)) if heat.max() > 0: heat = np.clip(heat, 0.0, 1.0) return heat return None def _onnx_build_mask(image_bhwc: torch.Tensor, preview: int, sensitivity: float, models_dir: str, anomaly_gain: float = 1.0) -> torch.Tensor: """Deprecated: ONNX path removed. Returns zero mask of input size.""" B, H, W, C = image_bhwc.shape return torch.zeros((B, H, W, 1), device=image_bhwc.device, dtype=image_bhwc.dtype) if not _try_init_onnx(models_dir): return torch.zeros((image_bhwc.shape[0], image_bhwc.shape[1], image_bhwc.shape[2], 1), device=image_bhwc.device, dtype=image_bhwc.dtype) if not _ONNX_SESS: return torch.zeros((image_bhwc.shape[0], image_bhwc.shape[1], image_bhwc.shape[2], 1), device=image_bhwc.device, dtype=image_bhwc.dtype) B, H, W, C = image_bhwc.shape device = image_bhwc.device dtype = image_bhwc.dtype # Process per-batch image masks = [] img_cpu = image_bhwc.detach().to('cpu') for b in range(B): masks_b = [] # Prepare input resized square preview target = int(max(16, min(1024, preview))) xb = img_cpu[b].movedim(-1, 0).unsqueeze(0) # 1,C,H,W x_stretch = F.interpolate(xb, size=(target, target), mode='bilinear', align_corners=False).clamp(0, 1) x_letter = _letterbox_nchw(xb, target).clamp(0, 1) # Try four variants: stretch RGB, letterbox RGB, stretch BGR, letterbox BGR variants = [ ("stretch-RGB", x_stretch), ("letterbox-RGB", x_letter), ("stretch-BGR", x_stretch[:, [2, 1, 0], :, :]), ("letterbox-BGR", x_letter[:, [2, 1, 0], :, :]), ] if _ONNX_DEBUG: try: print(f"[CADE2.5][ONNX] Build mask for image[{b}] -> preview {target}x{target}") except Exception: pass for name, sess in list(_ONNX_SESS.items()): try: inputs = sess.get_inputs() if not inputs: continue in_name = inputs[0].name in_shape = inputs[0].shape if hasattr(inputs[0], 'shape') else None # Choose layout automatically based on the presence of channel dim=3 if isinstance(in_shape, (list, tuple)) and len(in_shape) == 4: dim_vals = [] for d in in_shape: try: dim_vals.append(int(d)) except Exception: dim_vals.append(-1) if dim_vals[-1] == 3: layout = "NHWC" else: layout = "NCHW" else: layout = "NCHW?" if _ONNX_DEBUG: try: print(f"[CADE2.5][ONNX] Model '{name}' in_shape={in_shape} layout={layout}") except Exception: pass # Try multiple input variants and scales hm = None chosen = None for vname, vx in variants: if layout.startswith("NHWC"): xin = vx.permute(0, 2, 3, 1) else: xin = vx for scale in (1.0, 255.0): inp = (xin * float(scale)).numpy().astype(np.float32) feed = {in_name: inp} outs = sess.run(None, feed) if _ONNX_DEBUG: try: shapes = [] for o in outs: try: shapes.append(tuple(np.asarray(o).shape)) except Exception: shapes.append("?") print(f"[CADE2.5][ONNX] '{name}' {vname} scale={scale} -> outs shapes {shapes}") except Exception: pass hm = _try_heatmap_from_outputs(outs, (target, target)) if _ONNX_DEBUG: try: if hm is None: print(f"[CADE2.5][ONNX] '{name}' {vname} scale={scale}: no spatial heatmap detected") else: print(f"[CADE2.5][ONNX] '{name}' {vname} scale={scale}: heat stats min={np.min(hm):.4f} max={np.max(hm):.4f} mean={np.mean(hm):.4f}") except Exception: pass if hm is not None and np.max(hm) > 0: chosen = (vname, scale) break if hm is not None and np.max(hm) > 0: break if hm is None: continue # Scale by sensitivity and optional anomaly gain gain = float(max(0.0, sensitivity)) if 'anomaly' in name.lower(): gain *= float(max(0.0, anomaly_gain)) hm = np.clip(hm * gain, 0.0, 1.0) tmask = _np_to_mask_tensor(hm, H, W, device, dtype) if tmask is not None: masks_b.append(tmask) if _ONNX_DEBUG: try: area = float(tmask.movedim(-1,1).mean().item()) if chosen is not None: vname, scale = chosen print(f"[CADE2.5][ONNX] '{name}' via {vname} x{scale} area={area:.4f}") else: print(f"[CADE2.5][ONNX] '{name}' contribution area={area:.4f}") except Exception: pass except Exception: # Ignore failing models continue if not masks_b: masks.append(torch.zeros((1, H, W, 1), device=device, dtype=dtype)) else: # Soft-OR fusion: 1 - prod(1 - m) stack = torch.stack([masks_b[i] for i in range(len(masks_b))], dim=0) # M,1,H,W,1? actually B dims kept as 1 fused = 1.0 - torch.prod(1.0 - stack.clamp(0, 1), dim=0) # Light smoothing via bilinear down/up (anti alias) ch = fused.permute(0, 3, 1, 2) # B=1,C=1,H,W dd = F.interpolate(ch, scale_factor=0.5, mode='bilinear', align_corners=False, recompute_scale_factor=False) uu = F.interpolate(dd, size=(H, W), mode='bilinear', align_corners=False) fused = uu.permute(0, 2, 3, 1).clamp(0, 1) if _ONNX_DEBUG: try: area = float(fused.movedim(-1,1).mean().item()) print(f"[CADE2.5][ONNX] Fused area (image[{b}])={area:.4f}") except Exception: pass masks.append(fused) return torch.cat(masks, dim=0) def _sampler_names(): try: import comfy.samplers return comfy.samplers.KSampler.SAMPLERS except Exception: return ["euler"] def _scheduler_names(): try: import comfy.samplers scheds = list(comfy.samplers.KSampler.SCHEDULERS) if "MGHybrid" not in scheds: scheds.append("MGHybrid") return scheds except Exception: return ["normal", "MGHybrid"] def safe_decode(vae, lat, tile=512, ovlp=64): # Avoid building autograd graphs and release GPU memory early with torch.inference_mode(): h, w = lat["samples"].shape[-2:] if min(h, w) > 1024: # Increase overlap for ultra-hires to reduce seam artifacts ov = 128 if max(h, w) > 2048 else ovlp out = vae.decode_tiled(lat["samples"], tile_x=tile, tile_y=tile, overlap=ov) else: out = vae.decode(lat["samples"]) # Move to CPU and free VRAM ASAP try: try: out = out.detach() except Exception: pass out_cpu = out try: out_cpu = out_cpu.to('cpu') except Exception: pass try: del out except Exception: pass if torch.cuda.is_available(): try: torch.cuda.synchronize() except Exception: pass try: torch.cuda.empty_cache() except Exception: pass return out_cpu except Exception: return out def safe_encode(vae, img, tile=512, ovlp=64): import math, torch.nn.functional as F h, w = img.shape[1:3] try: stride = int(vae.spacial_compression_decode()) except Exception: stride = 8 if stride <= 0: stride = 8 def _align_up(x, s): return int(((x + s - 1) // s) * s) Ht = _align_up(h, stride) Wt = _align_up(w, stride) x = img if (Ht != h) or (Wt != w): # pad on bottom/right using replicate to avoid black borders pad_h = Ht - h pad_w = Wt - w x_nchw = img.movedim(-1, 1) x_nchw = F.pad(x_nchw, (0, pad_w, 0, pad_h), mode='replicate') x = x_nchw.movedim(1, -1) if min(Ht, Wt) > 1024: ov = 128 if max(Ht, Wt) > 2048 else ovlp return vae.encode_tiled(x[:, :, :, :3], tile_x=tile, tile_y=tile, overlap=ov) return vae.encode(x[:, :, :, :3]) def _gaussian_kernel(kernel_size: int, sigma: float, device=None): x, y = torch.meshgrid( torch.linspace(-1, 1, kernel_size, device=device), torch.linspace(-1, 1, kernel_size, device=device), indexing="ij", ) d = torch.sqrt(x * x + y * y) g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) return g / g.sum() def _sharpen_image(image: torch.Tensor, sharpen_radius: int, sigma: float, alpha: float): if sharpen_radius == 0: return (image,) image = image.to(model_management.get_torch_device()) batch_size, height, width, channels = image.shape kernel_size = sharpen_radius * 2 + 1 kernel = _gaussian_kernel(kernel_size, sigma, device=image.device) * -(alpha * 10) kernel = kernel.to(dtype=image.dtype) center = kernel_size // 2 kernel[center, center] = kernel[center, center] - kernel.sum() + 1.0 kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) tensor_image = image.permute(0, 3, 1, 2) tensor_image = F.pad(tensor_image, (sharpen_radius, sharpen_radius, sharpen_radius, sharpen_radius), 'reflect') sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels)[:, :, sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius] sharpened = sharpened.permute(0, 2, 3, 1) result = torch.clamp(sharpened, 0, 1) return (result.to(model_management.intermediate_device()),) def _encode_clip_image(image: torch.Tensor, clip_vision, target_res: int) -> torch.Tensor: # image: BHWC in [0,1] img = image.movedim(-1, 1) # BCHW img = F.interpolate(img, size=(target_res, target_res), mode="bilinear", align_corners=False) img = (img * 2.0) - 1.0 embeds = clip_vision.encode_image(img)["image_embeds"] embeds = F.normalize(embeds, dim=-1) return embeds def _clip_cosine_distance(a: torch.Tensor, b: torch.Tensor) -> float: if a.shape != b.shape: m = min(a.shape[0], b.shape[0]) a = a[:m] b = b[:m] sim = (a * b).sum(dim=-1).mean().clamp(-1.0, 1.0).item() return 1.0 - sim def _gaussian_blur_nchw(x: torch.Tensor, sigma: float = 1.0, radius: int = 1) -> torch.Tensor: """Lightweight depthwise Gaussian blur for NCHW or NCDHW tensors. Uses reflect padding and a normalized kernel built by _gaussian_kernel. """ if radius <= 0: return x ksz = radius * 2 + 1 kernel = _gaussian_kernel(ksz, sigma, device=x.device).to(dtype=x.dtype) # Support 5D by folding depth into batch if x.ndim == 5: b, c, d, h, w = x.shape x2 = x.permute(0, 2, 1, 3, 4).reshape(b * d, c, h, w) k = kernel.repeat(c, 1, 1).unsqueeze(1) # [C,1,K,K] x_pad = F.pad(x2, (radius, radius, radius, radius), mode='reflect') y2 = F.conv2d(x_pad, k, padding=0, groups=c) y = y2.reshape(b, d, c, h, w).permute(0, 2, 1, 3, 4) return y # 4D path if x.ndim == 4: b, c, h, w = x.shape k = kernel.repeat(c, 1, 1).unsqueeze(1) # [C,1,K,K] x_pad = F.pad(x, (radius, radius, radius, radius), mode='reflect') y = F.conv2d(x_pad, k, padding=0, groups=c) return y # Fallback: return input if unexpected dims return x def _letterbox_nchw(x: torch.Tensor, target: int, pad_val: float = 114.0 / 255.0) -> torch.Tensor: """Letterbox a BCHW tensor to target x target with constant padding (YOLO-style). Preserves aspect ratio, centers content, pads with pad_val. """ if x.ndim != 4: return F.interpolate(x, size=(target, target), mode='bilinear', align_corners=False) b, c, h, w = x.shape if h == 0 or w == 0: return F.interpolate(x, size=(target, target), mode='bilinear', align_corners=False) r = float(min(target / max(1, h), target / max(1, w))) nh = max(1, int(round(h * r))) nw = max(1, int(round(w * r))) y = F.interpolate(x, size=(nh, nw), mode='bilinear', align_corners=False) pt = (target - nh) // 2 pb = target - nh - pt pl = (target - nw) // 2 pr = target - nw - pl if pt < 0 or pb < 0 or pl < 0 or pr < 0: # Fallback stretch if rounding went weird return F.interpolate(x, size=(target, target), mode='bilinear', align_corners=False) return F.pad(y, (pl, pr, pt, pb), mode='constant', value=float(pad_val)) def _fdg_filter(delta: torch.Tensor, low_gain: float, high_gain: float, sigma: float = 1.0, radius: int = 1) -> torch.Tensor: """Frequency-Decoupled Guidance: split delta into low/high bands and reweight. delta: [B,C,H,W] """ low = _gaussian_blur_nchw(delta, sigma=sigma, radius=radius) high = delta - low return low * float(low_gain) + high * float(high_gain) def _fdg_split_three(delta: torch.Tensor, sigma_lo: float = 0.8, sigma_hi: float = 2.0, radius: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Tri-band split: returns (low, mid, high) for NCHW delta. low = G(sigma_hi) mid = G(sigma_lo) - G(sigma_hi) high = delta - G(sigma_lo) """ sig_lo = float(max(0.05, sigma_lo)) sig_hi = float(max(sig_lo + 1e-3, sigma_hi)) blur_lo = _gaussian_blur_nchw(delta, sigma=sig_lo, radius=radius) blur_hi = _gaussian_blur_nchw(delta, sigma=sig_hi, radius=radius) low = blur_hi mid = blur_lo - blur_hi high = delta - blur_lo return low, mid, high def _fdg_energy_fraction(delta: torch.Tensor, sigma: float = 1.0, radius: int = 1) -> torch.Tensor: """Return fraction of high-frequency energy: E_high / (E_low + E_high).""" low = _gaussian_blur_nchw(delta, sigma=sigma, radius=radius) high = delta - low e_low = (low * low).mean(dim=(1, 2, 3), keepdim=True) e_high = (high * high).mean(dim=(1, 2, 3), keepdim=True) frac = e_high / (e_low + e_high + 1e-8) return frac def _wrap_model_with_guidance(model, guidance_mode: str, rescale_multiplier: float, momentum_beta: float, cfg_curve: float, perp_damp: float, use_zero_init: bool=False, zero_init_steps: int=0, fdg_low: float = 0.6, fdg_high: float = 1.3, fdg_sigma: float = 1.0, ze_zero_steps: int = 0, ze_adaptive: bool = False, ze_r_switch_hi: float = 0.6, ze_r_switch_lo: float = 0.45, fdg_low_adaptive: bool = False, fdg_low_min: float = 0.45, fdg_low_max: float = 0.7, fdg_ema_beta: float = 0.8, use_local_mask: bool = False, mask_inside: float = 1.0, mask_outside: float = 1.0, midfreq_enable: bool = False, midfreq_gain: float = 0.0, midfreq_sigma_lo: float = 0.8, midfreq_sigma_hi: float = 2.0, mahiro_plus_enable: bool = False, mahiro_plus_strength: float = 0.5, eps_scale_enable: bool = False, eps_scale: float = 0.0, # NEW: CWN + AGC for Hard node too cwn_enable: bool = True, alpha_c: float = 1.0, alpha_u: float = 1.0, agc_enable: bool = True, agc_tau: float = 2.8, # NAG fallback nag_fb_enable: bool = False, nag_fb_scale: float = 4.0, nag_fb_tau: float = 2.5, nag_fb_alpha: float = 0.25): """Clone model and attach a cfg mixing function implementing RescaleCFG/FDG, CFGZero*/FD, or hybrid ZeResFDG. guidance_mode: 'default' | 'RescaleCFG' | 'RescaleFDG' | 'CFGZero*' | 'CFGZeroFD' | 'ZeResFDG' """ if guidance_mode == "default": return model m = model.clone() # State for momentum and sigma normalization across steps prev_delta = {"t": None} sigma_seen = {"max": None, "min": None} # Spectral switching/adaptive low state spec_state = {"ema": None, "mode": "CFGZeroFD"} # External reset hook to emulate fresh state per iteration without re-cloning the model def _mg_guidance_reset(): try: prev_delta["t"] = None sigma_seen["max"] = None sigma_seen["min"] = None spec_state["ema"] = None spec_state["mode"] = "CFGZeroFD" except Exception: pass try: setattr(m, "mg_guidance_reset", _mg_guidance_reset) except Exception: pass def cfg_func(args): cond = args["cond"] uncond = args["uncond"] cond_scale = args["cond_scale"] sigma = args.get("sigma", None) x_orig = args.get("input", None) # NAG fallback (noise-space) when CrossAttention patch inactive if bool(nag_fb_enable): try: from . import mg_sagpu_attention as _sa active = bool(getattr(_sa, "_nag_patch_active", False)) except Exception: active = False if not active: try: phi = float(nag_fb_scale); tau = float(nag_fb_tau); a = float(nag_fb_alpha) g = cond * phi - uncond * (phi - 1.0) def _l1(x): return torch.sum(torch.abs(x), dim=(1,2,3), keepdim=True).clamp_min(1e-6) s_pos = _l1(cond); s_g = _l1(g) scale = (s_pos * tau) / s_g g = torch.where((s_g > s_pos * tau), g * scale, g) cond = g * a + cond * (1.0 - a) except Exception: pass # Local spatial gain from CURRENT_ONNX_MASK_BCHW, resized to cond spatial size def _local_gain_for(hw): if not bool(use_local_mask): return None m = globals().get("CURRENT_ONNX_MASK_BCHW", None) if m is None: return None try: Ht, Wt = int(hw[0]), int(hw[1]) g = m.to(device=cond.device, dtype=cond.dtype) if g.shape[-2] != Ht or g.shape[-1] != Wt: g = F.interpolate(g, size=(Ht, Wt), mode='bilinear', align_corners=False) gi = float(mask_inside) go = float(mask_outside) gain = g * gi + (1.0 - g) * go # [B,1,H,W] return gain except Exception: return None # Allow hybrid switch per-step mode = guidance_mode if guidance_mode == "ZeResFDG": if bool(ze_adaptive): try: delta_raw = args["cond"] - args["uncond"] frac_b = _fdg_energy_fraction(delta_raw, sigma=float(fdg_sigma), radius=1) # [B,1,1,1] frac = float(frac_b.mean().clamp(0.0, 1.0).item()) except Exception: frac = 0.0 if spec_state["ema"] is None: spec_state["ema"] = frac else: beta = float(max(0.0, min(0.99, fdg_ema_beta))) spec_state["ema"] = beta * float(spec_state["ema"]) + (1.0 - beta) * frac r = float(spec_state["ema"]) # Hysteresis: switch up/down with two thresholds if spec_state["mode"] == "CFGZeroFD" and r >= float(ze_r_switch_hi): spec_state["mode"] = "RescaleFDG" elif spec_state["mode"] == "RescaleFDG" and r <= float(ze_r_switch_lo): spec_state["mode"] = "CFGZeroFD" mode = spec_state["mode"] else: try: sigmas = args["model_options"]["transformer_options"]["sample_sigmas"] matched_idx = (sigmas == args["timestep"][0]).nonzero() if len(matched_idx) > 0: current_idx = matched_idx.item() else: current_idx = 0 except Exception: current_idx = 0 mode = "CFGZeroFD" if current_idx <= int(ze_zero_steps) else "RescaleFDG" if mode in ("CFGZero*", "CFGZeroFD"): # Optional zero-init for the first N steps if use_zero_init and "model_options" in args and args.get("timestep") is not None: try: sigmas = args["model_options"]["transformer_options"]["sample_sigmas"] matched_idx = (sigmas == args["timestep"][0]).nonzero() if len(matched_idx) > 0: current_idx = matched_idx.item() else: # fallback lookup current_idx = 0 if current_idx <= int(zero_init_steps): return cond * 0.0 except Exception: pass # CWN for CFGZero branches: align energies before projection if bool(cwn_enable): try: _eps = 1e-6 sc = (cond.pow(2).mean(dim=(1, 2, 3), keepdim=True).sqrt() + _eps) su = (uncond.pow(2).mean(dim=(1, 2, 3), keepdim=True).sqrt() + _eps) g = 0.5 * (sc + su) cond = cond * (float(alpha_c) * g / sc) uncond = uncond * (float(alpha_u) * g / su) except Exception: pass # Project cond onto uncond subspace (batch-wise alpha) bsz = cond.shape[0] pos_flat = cond.view(bsz, -1) neg_flat = uncond.view(bsz, -1) dot = torch.sum(pos_flat * neg_flat, dim=1, keepdim=True) denom = torch.sum(neg_flat * neg_flat, dim=1, keepdim=True).clamp_min(1e-8) alpha = (dot / denom).view(bsz, *([1] * (cond.dim() - 1))) resid = cond - uncond * alpha # Adaptive low gain if enabled low_gain_eff = float(fdg_low) if bool(fdg_low_adaptive) and spec_state["ema"] is not None: s = float(spec_state["ema"]) # 0..1 fraction of high-frequency energy lmin = float(fdg_low_min) lmax = float(fdg_low_max) low_gain_eff = max(0.0, min(2.0, lmin + (lmax - lmin) * s)) if mode == "CFGZeroFD": resid = _fdg_filter(resid, low_gain=low_gain_eff, high_gain=fdg_high, sigma=float(fdg_sigma), radius=1) # Apply local spatial gain to residual guidance lg = _local_gain_for((cond.shape[-2], cond.shape[-1])) if lg is not None: resid = resid * lg.expand(-1, resid.shape[1], -1, -1) noise_pred = uncond * alpha + cond_scale * resid return noise_pred # RescaleCFG/FDG path (with optional momentum/perp damping and S-curve shaping) delta = cond - uncond pd = float(max(0.0, min(1.0, perp_damp))) if pd > 0.0 and (prev_delta["t"] is not None) and (prev_delta["t"].shape == delta.shape): prev = prev_delta["t"] denom = (prev * prev).sum(dim=(1,2,3), keepdim=True).clamp_min(1e-6) coeff = ((delta * prev).sum(dim=(1,2,3), keepdim=True) / denom) parallel = coeff * prev delta = delta - pd * parallel beta = float(max(0.0, min(0.95, momentum_beta))) if beta > 0.0: if prev_delta["t"] is None or prev_delta["t"].shape != delta.shape: prev_delta["t"] = delta.detach() delta = (1.0 - beta) * delta + beta * prev_delta["t"] prev_delta["t"] = delta.detach() cond = uncond + delta else: prev_delta["t"] = delta.detach() # Adaptive Guidance Clipping on delta (Rescale path) if bool(agc_enable): try: t = float(max(0.5, agc_tau)) delta = t * torch.tanh(delta / t) except Exception: pass # After momentum: optionally apply FDG and rebuild cond if mode == "RescaleFDG": # Adaptive low gain if enabled low_gain_eff = float(fdg_low) if bool(fdg_low_adaptive) and spec_state["ema"] is not None: s = float(spec_state["ema"]) # 0..1 lmin = float(fdg_low_min) lmax = float(fdg_low_max) low_gain_eff = max(0.0, min(2.0, lmin + (lmax - lmin) * s)) delta_fdg = _fdg_filter(delta, low_gain=low_gain_eff, high_gain=fdg_high, sigma=float(fdg_sigma), radius=1) # Optional mid-frequency emphasis (band-pass) blended on top if bool(midfreq_enable) and abs(float(midfreq_gain)) > 1e-6: lo, mid, hi = _fdg_split_three(delta, sigma_lo=float(midfreq_sigma_lo), sigma_hi=float(midfreq_sigma_hi), radius=1) # Respect local mask gain if present lg = _local_gain_for((cond.shape[-2], cond.shape[-1])) if lg is not None: mid = mid * lg.expand(-1, mid.shape[1], -1, -1) delta_fdg = delta_fdg + float(midfreq_gain) * mid lg = _local_gain_for((cond.shape[-2], cond.shape[-1])) if lg is not None: delta_fdg = delta_fdg * lg.expand(-1, delta_fdg.shape[1], -1, -1) cond = uncond + delta_fdg else: lg = _local_gain_for((cond.shape[-2], cond.shape[-1])) if lg is not None: delta = delta * lg.expand(-1, delta.shape[1], -1, -1) cond = uncond + delta cond_scale_eff = cond_scale if cfg_curve > 0.0 and (sigma is not None): s = sigma if s.ndim > 1: s = s.flatten() s_max = float(torch.max(s).item()) s_min = float(torch.min(s).item()) if sigma_seen["max"] is None: sigma_seen["max"] = s_max sigma_seen["min"] = s_min else: sigma_seen["max"] = max(sigma_seen["max"], s_max) sigma_seen["min"] = min(sigma_seen["min"], s_min) lo = max(1e-6, sigma_seen["min"]) hi = max(lo * (1.0 + 1e-6), sigma_seen["max"]) t = (torch.log(s + 1e-6) - torch.log(torch.tensor(lo, device=sigma.device))) / (torch.log(torch.tensor(hi, device=sigma.device)) - torch.log(torch.tensor(lo, device=sigma.device)) + 1e-6) t = t.clamp(0.0, 1.0) k = 6.0 * float(cfg_curve) s_curve = torch.tanh((t - 0.5) * k) gain = 1.0 + 0.15 * float(cfg_curve) * s_curve if gain.ndim > 0: gain = gain.mean().item() cond_scale_eff = cond_scale * float(gain) # Epsilon scaling (exposure bias correction): early steps get multiplier closer to (1 + eps_scale) eps_mult = 1.0 if bool(eps_scale_enable) and (sigma is not None): try: s = sigma if s.ndim > 1: s = s.flatten() s_max = float(torch.max(s).item()) s_min = float(torch.min(s).item()) if sigma_seen["max"] is None: sigma_seen["max"] = s_max sigma_seen["min"] = s_min else: sigma_seen["max"] = max(sigma_seen["max"], s_max) sigma_seen["min"] = min(sigma_seen["min"], s_min) lo = max(1e-6, sigma_seen["min"]) hi = max(lo * (1.0 + 1e-6), sigma_seen["max"]) t_lin = (torch.log(s + 1e-6) - torch.log(torch.tensor(lo, device=sigma.device))) / (torch.log(torch.tensor(hi, device=sigma.device)) - torch.log(torch.tensor(lo, device=sigma.device)) + 1e-6) t_lin = t_lin.clamp(0.0, 1.0) w_early = (1.0 - t_lin).mean().item() eps_mult = float(1.0 + eps_scale * w_early) except Exception: eps_mult = float(1.0 + eps_scale) if sigma is None or x_orig is None: return uncond + cond_scale * (cond - uncond) sigma_ = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1)) x = x_orig / (sigma_ * sigma_ + 1.0) v_cond = ((x - (x_orig - cond)) * (sigma_ ** 2 + 1.0) ** 0.5) / (sigma_) v_uncond = ((x - (x_orig - uncond)) * (sigma_ ** 2 + 1.0) ** 0.5) / (sigma_) # CWN in v-space for Rescale path (safer than eps-space) if bool(cwn_enable): try: _e = 1e-6 rc = (v_cond.pow(2).mean(dim=(1,2,3), keepdim=True).sqrt() + _e) ru = (v_uncond.pow(2).mean(dim=(1,2,3), keepdim=True).sqrt() + _e) v_cond_n = (v_cond / rc) * float(alpha_c) v_uncond_n = (v_uncond / ru) * float(alpha_u) except Exception: v_cond_n, v_uncond_n = v_cond, v_uncond else: v_cond_n, v_uncond_n = v_cond, v_uncond v_cfg = v_uncond_n + cond_scale_eff * (v_cond_n - v_uncond_n) ro_pos = torch.std(v_cond_n, dim=(1, 2, 3), keepdim=True) ro_cfg = torch.std(v_cfg, dim=(1, 2, 3), keepdim=True).clamp_min(1e-6) v_rescaled = v_cfg * (ro_pos / ro_cfg) v_final = float(rescale_multiplier) * v_rescaled + (1.0 - float(rescale_multiplier)) * v_cfg eps = x_orig - (x - (v_final * eps_mult) * sigma_ / (sigma_ * sigma_ + 1.0) ** 0.5) return eps m.set_model_sampler_cfg_function(cfg_func, disable_cfg1_optimization=True) # Note: ControlNet class-label injection wrapper removed to keep CADE neutral. # Optional directional post-mix inspired by Mahiro (global, no ONNX) if bool(mahiro_plus_enable): s_clamp = float(max(0.0, min(1.0, mahiro_plus_strength))) mb_state = {"ema": None} def _sqrt_sign(x: torch.Tensor) -> torch.Tensor: return x.sign() * torch.sqrt(x.abs().clamp_min(1e-12)) def _hp_split(x: torch.Tensor, radius: int = 1, sigma: float = 1.0): low = _gaussian_blur_nchw(x, sigma=sigma, radius=radius) high = x - low return low, high def _sched_gain(args) -> float: # Gentle mid-steps boost: triangle peak at the middle of schedule try: sigmas = args["model_options"]["transformer_options"]["sample_sigmas"] idx_t = args.get("timestep", None) if idx_t is None: return 1.0 matched = (sigmas == idx_t[0]).nonzero() if len(matched) == 0: return 1.0 i = float(matched.item()) n = float(sigmas.shape[0]) if n <= 1: return 1.0 phase = i / (n - 1.0) tri = 1.0 - abs(2.0 * phase - 1.0) return float(0.6 + 0.4 * tri) # 0.6 at edges -> 1.0 mid except Exception: return 1.0 def mahiro_plus_post(args): try: scale = args.get('cond_scale', 1.0) cond_p = args['cond_denoised'] uncond_p = args['uncond_denoised'] cfg = args['denoised'] # Orthogonalize positive to negative direction (batch-wise) bsz = cond_p.shape[0] pos_flat = cond_p.view(bsz, -1) neg_flat = uncond_p.view(bsz, -1) dot = torch.sum(pos_flat * neg_flat, dim=1, keepdim=True) denom = torch.sum(neg_flat * neg_flat, dim=1, keepdim=True).clamp_min(1e-8) alpha = (dot / denom).view(bsz, *([1] * (cond_p.dim() - 1))) c_orth = cond_p - uncond_p * alpha leap_raw = float(scale) * c_orth # Light high-pass emphasis for detail, protect low-frequency tone low, high = _hp_split(leap_raw, radius=1, sigma=1.0) leap = 0.35 * low + 1.00 * high # Directional agreement (global cosine over flattened dims) u_leap = float(scale) * uncond_p merge = 0.5 * (leap + cfg) nu = _sqrt_sign(u_leap).flatten(1) nm = _sqrt_sign(merge).flatten(1) sim = F.cosine_similarity(nu, nm, dim=1).mean() a = torch.clamp((sim + 1.0) * 0.5, 0.0, 1.0) # Small EMA for temporal smoothness if mb_state["ema"] is None: mb_state["ema"] = float(a) else: mb_state["ema"] = 0.8 * float(mb_state["ema"]) + 0.2 * float(a) a_eff = float(mb_state["ema"]) w = a_eff * cfg + (1.0 - a_eff) * leap # Gentle energy match to CFG dims = tuple(range(1, w.dim())) ro_w = torch.std(w, dim=dims, keepdim=True).clamp_min(1e-6) ro_cfg = torch.std(cfg, dim=dims, keepdim=True).clamp_min(1e-6) w_res = w * (ro_cfg / ro_w) # Schedule gain over steps (mid stronger) s_eff = s_clamp * _sched_gain(args) out = (1.0 - s_eff) * cfg + s_eff * w_res return out except Exception: return args['denoised'] try: m.set_model_sampler_post_cfg_function(mahiro_plus_post) except Exception: pass # Quantile clamp stabilizer (per-sample): soft range limit for denoised tensor # Always on, under the hood. Helps prevent rare exploding values. def _qclamp_post(args): try: x = args.get("denoised", None) if x is None: return args["denoised"] dt = x.dtype xf = x.to(dtype=torch.float32) B = xf.shape[0] lo_q, hi_q = 0.001, 0.999 out = [] for i in range(B): t = xf[i].reshape(-1) try: lo = torch.quantile(t, lo_q) hi = torch.quantile(t, hi_q) except Exception: n = t.numel() k_lo = max(1, int(n * lo_q)) k_hi = max(1, int(n * hi_q)) lo = torch.kthvalue(t, k_lo).values hi = torch.kthvalue(t, k_hi).values out.append(xf[i].clamp(min=lo, max=hi)) y = torch.stack(out, dim=0).to(dtype=dt) return y except Exception: return args["denoised"] try: m.set_model_sampler_post_cfg_function(_qclamp_post) except Exception: pass return m # --- AQClip-Lite: adaptive soft quantile clipping in latent space (tile overlap) --- @torch.no_grad() def _aqclip_lite(latent_bchw: torch.Tensor, tile: int = 32, stride: int = 16, alpha: float = 2.0, ema_state: dict | None = None, ema_beta: float = 0.8, H_override: torch.Tensor | None = None) -> tuple[torch.Tensor, dict]: try: z = latent_bchw B, C, H, W = z.shape dev, dt = z.device, z.dtype ksize = max(8, min(int(tile), min(H, W))) kstride = max(1, min(int(stride), ksize)) # Confidence map: attention entropy override or gradient proxy if (H_override is not None) and isinstance(H_override, torch.Tensor): hsrc = H_override.to(device=dev, dtype=dt) if hsrc.dim() == 3: hsrc = hsrc.unsqueeze(1) gpool = F.avg_pool2d(hsrc, kernel_size=ksize, stride=kstride) else: zm = z.mean(dim=1, keepdim=True) kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=dev, dtype=dt).view(1, 1, 3, 3) ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=dev, dtype=dt).view(1, 1, 3, 3) gx = F.conv2d(zm, kx, padding=1) gy = F.conv2d(zm, ky, padding=1) gmag = torch.sqrt(gx * gx + gy * gy) gpool = F.avg_pool2d(gmag, kernel_size=ksize, stride=kstride) gmax = gpool.amax(dim=(2, 3), keepdim=True).clamp_min(1e-6) Hn = (gpool / gmax).squeeze(1) # B,h',w' L = Hn.shape[1] * Hn.shape[2] Hn = Hn.reshape(B, L) # Map confidence -> quantiles ql = 0.5 * (Hn ** 2) qh = 1.0 - 0.5 * ((1.0 - Hn) ** 2) # Per-tile mean/std unf = F.unfold(z, kernel_size=ksize, stride=kstride) # B, C*ksize*ksize, L M = unf.shape[1] mu = unf.mean(dim=1).to(torch.float32) # B,L var = (unf.to(torch.float32) - mu.unsqueeze(1)).pow(2).mean(dim=1) sigma = (var + 1e-12).sqrt() # Normal inverse approximation: ndtri(q) = sqrt(2)*erfinv(2q-1) def _ndtri(q: torch.Tensor) -> torch.Tensor: return (2.0 ** 0.5) * torch.special.erfinv(q.mul(2.0).sub(1.0).clamp(-0.999999, 0.999999)) k_neg = _ndtri(ql).abs() k_pos = _ndtri(qh).abs() lo = mu - k_neg * sigma hi = mu + k_pos * sigma # EMA smooth if ema_state is None: ema_state = {} b = float(max(0.0, min(0.999, ema_beta))) if 'lo' in ema_state and 'hi' in ema_state and ema_state['lo'].shape == lo.shape: lo = b * ema_state['lo'] + (1.0 - b) * lo hi = b * ema_state['hi'] + (1.0 - b) * hi ema_state['lo'] = lo.detach() ema_state['hi'] = hi.detach() # Soft tanh clip (vectorized in unfold domain) mid = (lo + hi) * 0.5 half = (hi - lo) * 0.5 half = half.clamp_min(1e-6) y = (unf.to(torch.float32) - mid.unsqueeze(1)) / half.unsqueeze(1) y = torch.tanh(float(alpha) * y) unf_clipped = mid.unsqueeze(1) + half.unsqueeze(1) * y unf_clipped = unf_clipped.to(dt) out = F.fold(unf_clipped, output_size=(H, W), kernel_size=ksize, stride=kstride) ones = torch.ones((B, M, L), device=dev, dtype=dt) w = F.fold(ones, output_size=(H, W), kernel_size=ksize, stride=kstride).clamp_min(1e-6) out = out / w return out, ema_state except Exception: return latent_bchw, (ema_state or {}) class ComfyAdaptiveDetailEnhancer25: @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL", {}), "positive": ("CONDITIONING", {}), "negative": ("CONDITIONING", {}), "vae": ("VAE", {}), "latent": ("LATENT", {}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1}), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.0001}), "sampler_name": (_sampler_names(), {"default": _sampler_names()[0]}), "scheduler": (_scheduler_names(), {"default": _scheduler_names()[0]}), "iterations": ("INT", {"default": 1, "min": 1, "max": 1000}), "steps_delta": ("FLOAT", {"default": 0.0, "min": -1000.0, "max": 1000.0, "step": 0.01}), "cfg_delta": ("FLOAT", {"default": 0.0, "min": -100.0, "max": 100.0, "step": 0.01}), "denoise_delta": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.0001}), "apply_sharpen": ("BOOLEAN", {"default": False}), "apply_upscale": ("BOOLEAN", {"default": False}), "apply_ids": ("BOOLEAN", {"default": False}), "clip_clean": ("BOOLEAN", {"default": False}), "ids_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), "upscale_method": (MagicUpscaleModule.upscale_methods, {"default": "lanczos"}), "scale_by": ("FLOAT", {"default": 1.2, "min": 1.0, "max": 8.0, "step": 0.01}), "scale_delta": ("FLOAT", {"default": 0.0, "min": -8.0, "max": 8.0, "step": 0.01}), "noise_offset": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.5, "step": 0.01}), "threshold": ("FLOAT", {"default": 0.03, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "RMS latent drift threshold (smaller = more damping)."}), }, "optional": { "Sharpnes_strenght": ("FLOAT", {"default": 0.300, "min": 0.0, "max": 1.0, "step": 0.001}), "latent_compare": ("BOOLEAN", {"default": False, "tooltip": "Use latent drift to gently damp params (safer than overwriting latents)."}), "accumulation": (["default", "fp32+fp16", "fp32+fp32"], {"default": "default", "tooltip": "Override SageAttention PV accumulation mode for this node run."}), "reference_clean": ("BOOLEAN", {"default": False, "tooltip": "Use CLIP-Vision similarity to a reference image to stabilize output."}), "reference_image": ("IMAGE", {}), "clip_vision": ("CLIP_VISION", {}), "ref_preview": ("INT", {"default": 224, "min": 64, "max": 512, "step": 16}), "ref_threshold": ("FLOAT", {"default": 0.03, "min": 0.0, "max": 0.2, "step": 0.001}), "ref_cooldown": ("INT", {"default": 1, "min": 1, "max": 8}), # ONNX detectors removed # Guidance controls "guidance_mode": (["default", "RescaleCFG", "RescaleFDG", "CFGZero*", "CFGZeroFD", "ZeResFDG"], {"default": "RescaleCFG", "tooltip": "Rescale (stable), RescaleFDG (spectral), CFGZero*, CFGZeroFD, or hybrid ZeResFDG."}), "rescale_multiplier": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Blend between rescaled and plain CFG (like comfy RescaleCFG)."}), "momentum_beta": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "tooltip": "EMA momentum in eps-space for (cond-uncond), 0 to disable."}), "cfg_curve": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "S-curve shaping of cond_scale across steps (0=flat)."}), "perp_damp": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Remove a small portion of the component parallel to previous delta (0-1)."}), # Conditioning Weight Normalization (CWN) + Adaptive Guidance Clipping (AGC) "cwn_enable": ("BOOLEAN", {"default": True, "tooltip": "Normalize cond/uncond energy to steady CFG mixing."}), "alpha_c": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), "alpha_u": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), "agc_enable": ("BOOLEAN", {"default": True, "tooltip": "Soft-clip residual guidance to prevent rare spikes."}), "agc_tau": ("FLOAT", {"default": 2.8, "min": 0.5, "max": 6.0, "step": 0.1}), # NAG (Normalized Attention Guidance) toggles "use_nag": ("BOOLEAN", {"default": False, "tooltip": "Apply NAG inside CrossAttention (positive branch) during this node."}), "nag_scale": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 50.0, "step": 0.1}), "nag_tau": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 10.0, "step": 0.01}), "nag_alpha": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}), # AQClip-Lite (adaptive latent clipping) "aqclip_enable": ("BOOLEAN", {"default": False, "tooltip": "Adaptive soft tile clipping with overlap (reduces spikes on uncertain regions)."}), "aq_tile": ("INT", {"default": 32, "min": 8, "max": 128, "step": 1}), "aq_stride": ("INT", {"default": 16, "min": 4, "max": 128, "step": 1}), "aq_alpha": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 4.0, "step": 0.1}), "aq_ema_beta": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 0.99, "step": 0.01}), "aq_attn": ("BOOLEAN", {"default": False, "tooltip": "Use attention entropy as confidence (requires patched attention)."}), # CFGZero* extras "use_zero_init": ("BOOLEAN", {"default": False, "tooltip": "For CFGZero*, zero out first few steps."}), "zero_init_steps": ("INT", {"default": 0, "min": 0, "max": 20, "step": 1}), # FDG controls (placed last to avoid reordering existing fields) "fdg_low": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Low-frequency gain (<1 to restrain masses)."}), "fdg_high": ("FLOAT", {"default": 1.3, "min": 0.5, "max": 2.5, "step": 0.01, "tooltip": "High-frequency gain (>1 to boost details)."}), "fdg_sigma": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.5, "step": 0.05, "tooltip": "Gaussian sigma for FDG low-pass split."}), "ze_res_zero_steps": ("INT", {"default": 2, "min": 0, "max": 20, "step": 1, "tooltip": "Hybrid: number of initial steps to use CFGZeroFD before switching to RescaleFDG."}), # Adaptive spectral switch (ZeRes) and adaptive low gain "ze_adaptive": ("BOOLEAN", {"default": False, "tooltip": "Enable spectral switch: CFGZeroFD, RescaleFDG by HF/LF ratio (EMA)."}), "ze_r_switch_hi": ("FLOAT", {"default": 0.60, "min": 0.10, "max": 0.95, "step": 0.01, "tooltip": "Switch to RescaleFDG when EMA fraction of high-frequency."}), "ze_r_switch_lo": ("FLOAT", {"default": 0.45, "min": 0.05, "max": 0.90, "step": 0.01, "tooltip": "Switch back to CFGZeroFD when EMA fraction (hysteresis)."}), "fdg_low_adaptive": ("BOOLEAN", {"default": False, "tooltip": "Adapt fdg_low by HF fraction (EMA)."}), "fdg_low_min": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Lower bound for adaptive fdg_low."}), "fdg_low_max": ("FLOAT", {"default": 0.70, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Upper bound for adaptive fdg_low."}), "fdg_ema_beta": ("FLOAT", {"default": 0.80, "min": 0.0, "max": 0.99, "step": 0.01, "tooltip": "EMA smoothing for spectral ratio (higher = smoother)."}), # Mid-frequency stabilizer (hands/objects scale) "midfreq_enable": ("BOOLEAN", {"default": True, "tooltip": "Enable mid-frequency stabilizer (band-pass) to keep hands/objects stable at hi-res."}), "midfreq_gain": ("FLOAT", {"default": 0.65, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Blend amount of mid-frequency band added on top of FDG guidance (0..2)."}), "midfreq_sigma_lo": ("FLOAT", {"default": 0.55, "min": 0.05, "max": 2.0, "step": 0.01, "tooltip": "Lower Gaussian sigma for band split (controls smaller forms)."}), "midfreq_sigma_hi": ("FLOAT", {"default": 1.30, "min": 0.10, "max": 3.0, "step": 0.01, "tooltip": "Upper Gaussian sigma for band split (controls larger forms)."}), # ONNX local guidance and keypoints removed # Muse Blend global directional post-mix "muse_blend": ("BOOLEAN", {"default": False, "tooltip": "Enable Muse Blend (Mahiro+): gentle directional positive blend (global)."}), "muse_blend_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Overall influence of Muse Blend over baseline CFG (0..1)."}), # Exposure Bias Correction (epsilon scaling) "eps_scale_enable": ("BOOLEAN", {"default": False, "tooltip": "Exposure Bias Correction: scale predicted noise early in schedule."}), "eps_scale": ("FLOAT", {"default": 0.005, "min": -1.0, "max": 1.0, "step": 0.0005, "tooltip": "Signed scaling near early steps (recommended ~0.0045; use with care)."}), # KV pruning (self-attention speedup) "kv_prune_enable": ("BOOLEAN", {"default": False, "tooltip": "Speed: prune K/V tokens in self-attention by energy (safe on hi-res blocks)."}), "kv_keep": ("FLOAT", {"default": 0.85, "min": 0.5, "max": 1.0, "step": 0.01, "tooltip": "Fraction of tokens to keep when KV pruning is enabled."}), "kv_min_tokens": ("INT", {"default": 128, "min": 1, "max": 16384, "step": 1, "tooltip": "Minimum sequence length to apply KV pruning."}), "clipseg_enable": ("BOOLEAN", {"default": False, "tooltip": "Use CLIPSeg to build a text-driven mask (e.g., 'eyes | hands | face')."}), "clipseg_text": ("STRING", {"default": "", "multiline": False}), "clipseg_preview": ("INT", {"default": 224, "min": 64, "max": 512, "step": 16}), "clipseg_threshold": ("FLOAT", {"default": 0.40, "min": 0.0, "max": 1.0, "step": 0.05}), "clipseg_blur": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 15.0, "step": 0.1}), "clipseg_dilate": ("INT", {"default": 4, "min": 0, "max": 10, "step": 1}), "clipseg_gain": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), "clipseg_blend": (["fuse", "replace", "intersect"], {"default": "fuse", "tooltip": "How to combine CLIPSeg with any pre-mask (if present)."}), "clipseg_ref_gate": ("BOOLEAN", {"default": False, "tooltip": "If reference provided, boost mask when far from reference (CLIP-Vision)."}), "clipseg_ref_threshold": ("FLOAT", {"default": 0.03, "min": 0.0, "max": 0.2, "step": 0.001}), # Under-the-hood saving (disabled by default) "auto_save": ("BOOLEAN", {"default": False, "tooltip": "Save final IMAGE directly from CADE (uses low PNG compress to reduce RAM)."}), "save_prefix": ("STRING", {"default": "ComfyUI", "multiline": False}), "save_compress": ("INT", {"default": 1, "min": 0, "max": 9, "step": 1}), # Polish mode (final hi-res refinement) "polish_enable": ("BOOLEAN", {"default": False, "tooltip": "Polish: keep low-frequency shape from reference while allowing high-frequency details to refine."}), "polish_keep_low": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "How much low-frequency (global form, lighting) to take from reference image (0=use current, 1=use reference)."}), "polish_edge_lock": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Edge lock strength: protects edges from sideways drift (0=off, 1=strong)."}), "polish_sigma": ("FLOAT", {"default": 1.0, "min": 0.3, "max": 3.0, "step": 0.1, "tooltip": "Radius for low/high split: larger keeps bigger shapes as 'low' (global form)."}), "polish_start_after": ("INT", {"default": 1, "min": 0, "max": 3, "step": 1, "tooltip": "Enable polish after N iterations (0=immediately)."}), "polish_keep_low_ramp": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Starting share of low-frequency mix; ramps to polish_keep_low over remaining iterations."}), }, } RETURN_TYPES = ("LATENT", "IMAGE", "INT", "FLOAT", "FLOAT", "IMAGE") RETURN_NAMES = ("LATENT", "IMAGE", "steps", "cfg", "denoise", "mask_preview") FUNCTION = "apply_cade2" CATEGORY = "MagicNodes" def apply_cade2(self, model, vae, positive, negative, latent, seed, steps, cfg, denoise, sampler_name, scheduler, noise_offset, iterations=1, steps_delta=0.0, cfg_delta=0.0, denoise_delta=0.0, apply_sharpen=False, apply_upscale=False, apply_ids=False, clip_clean=False, ids_strength=0.5, upscale_method="lanczos", scale_by=1.2, scale_delta=0.0, Sharpnes_strenght=0.300, threshold=0.03, latent_compare=False, accumulation="default", reference_clean=False, reference_image=None, clip_vision=None, ref_preview=224, ref_threshold=0.03, ref_cooldown=1, guidance_mode="RescaleCFG", rescale_multiplier=0.7, momentum_beta=0.0, cfg_curve=0.0, perp_damp=0.0, cwn_enable=True, alpha_c=1.0, alpha_u=1.0, agc_enable=True, agc_tau=2.8, use_nag=False, nag_scale=4.0, nag_tau=2.5, nag_alpha=0.25, aqclip_enable=False, aq_tile=32, aq_stride=16, aq_alpha=2.0, aq_ema_beta=0.8, aq_attn=False, use_zero_init=False, zero_init_steps=0, fdg_low=0.6, fdg_high=1.3, fdg_sigma=1.0, ze_res_zero_steps=2, ze_adaptive=False, ze_r_switch_hi=0.60, ze_r_switch_lo=0.45, fdg_low_adaptive=False, fdg_low_min=0.45, fdg_low_max=0.70, fdg_ema_beta=0.80, midfreq_enable=True, midfreq_gain=0.65, midfreq_sigma_lo=0.55, midfreq_sigma_hi=1.30, muse_blend=False, muse_blend_strength=0.5, eps_scale_enable=False, eps_scale=0.005, clipseg_enable=False, clipseg_text="", clipseg_preview=224, clipseg_threshold=0.40, clipseg_blur=7.0, clipseg_dilate=4, clipseg_gain=1.0, clipseg_blend="fuse", clipseg_ref_gate=False, clipseg_ref_threshold=0.03, polish_enable=False, polish_keep_low=0.4, polish_edge_lock=0.2, polish_sigma=1.0, polish_start_after=1, polish_keep_low_ramp=0.2, auto_save=False, save_prefix="ComfyUI", save_compress=1, kv_prune_enable=False, kv_keep=0.85, kv_min_tokens=128): # Hard reset of any sticky globals from prior runs try: global CURRENT_ONNX_MASK_BCHW CURRENT_ONNX_MASK_BCHW = None except Exception: pass image = safe_decode(vae, latent) tuned_steps, tuned_cfg, tuned_denoise = AdaptiveSamplerHelper().tune( image, steps, cfg, denoise) current_steps = tuned_steps current_cfg = tuned_cfg current_denoise = tuned_denoise # Work on a detached copy to avoid mutating input latent across runs try: current_latent = {"samples": latent["samples"].clone()} except Exception: current_latent = {"samples": latent["samples"]} current_scale = scale_by ref_embed = None if reference_clean and (clip_vision is not None) and (reference_image is not None): try: ref_embed = _encode_clip_image(reference_image, clip_vision, ref_preview) except Exception: ref_embed = None # Pre-disable any lingering NAG patch from previous runs and set PV accumulation for this node try: sa_patch.enable_crossattention_nag_patch(False) except Exception: pass prev_accum = getattr(sa_patch, "CURRENT_PV_ACCUM", None) sa_patch.CURRENT_PV_ACCUM = None if accumulation == "default" else accumulation # Enable NAG patch if requested try: sa_patch.enable_crossattention_nag_patch(bool(use_nag), float(nag_scale), float(nag_tau), float(nag_alpha)) except Exception: pass # Enable attention-entropy probe for AQClip Attn-mode try: if hasattr(sa_patch, "enable_attention_entropy_capture"): sa_patch.enable_attention_entropy_capture(bool(aq_attn), max_tokens=1024, max_heads=4) except Exception: pass # Visual separation and start marker try: print("") except Exception: pass try: print("\x1b[32m==== Starting main job ====\x1b[0m") except Exception: pass # Enable KV pruning (self-attention) if requested try: if hasattr(sa_patch, "set_kv_prune"): sa_patch.set_kv_prune(bool(kv_prune_enable), float(kv_keep), int(kv_min_tokens)) except Exception: pass mask_last = None try: with torch.inference_mode(): __cade_noop = 0 # ensure non-empty with-block # Preflight: reset sticky state and build external masks once (CPU-pinned) try: CURRENT_ONNX_MASK_BCHW = None except Exception: pass pre_mask = None pre_area = 0.0 # ONNX mask removed # Build CLIPSeg mask once if bool(clipseg_enable) and isinstance(clipseg_text, str) and clipseg_text.strip() != "": try: cmask = _clipseg_build_mask(image, clipseg_text, int(clipseg_preview), float(clipseg_threshold), float(clipseg_blur), int(clipseg_dilate), float(clipseg_gain), None, None, float(clipseg_ref_threshold)) if cmask is not None: if pre_mask is None: pre_mask = cmask else: pre_mask, cmask = _align_mask_pair(pre_mask, cmask) if clipseg_blend == "replace": pre_mask = cmask elif clipseg_blend == "intersect": pre_mask = (pre_mask * cmask).clamp(0, 1) else: pre_mask = (1.0 - (1.0 - pre_mask) * (1.0 - cmask)).clamp(0, 1) except Exception: pass if pre_mask is not None: mask_last = pre_mask om = pre_mask.movedim(-1, 1) pre_area = float(om.mean().item()) # One-time gentle damping from area (disabled to preserve outline precision) # try: # if pre_area > 0.005: # damp = 1.0 - min(0.10, 0.02 + pre_area * 0.08) # current_denoise = max(0.10, current_denoise * damp) # current_cfg = max(1.0, current_cfg * (1.0 - 0.005)) # except Exception: # pass # Compact status try: clipseg_status = "on" if bool(clipseg_enable) and isinstance(clipseg_text, str) and clipseg_text.strip() != "" else "off" # print preflight info only in debug sessions (muted by default) if False: print(f"[CADE2.5][preflight] clipseg={clipseg_status} device={'cpu' if _CLIPSEG_FORCE_CPU else _CLIPSEG_DEV} mask_area={pre_area:.4f}") except Exception: pass # Freeze per-iteration external mask rebuild clipseg_enable = False # Depth gate cache for micro-detail injection (reuse per resolution) depth_gate_cache = {"size": None, "mask": None} # Release preflight temporaries to avoid keeping big tensors alive try: del cmask except Exception: pass try: del om except Exception: pass try: del pre_mask except Exception: pass try: del image except Exception: pass # Prepare guided sampler once per node run to avoid cloning model each iteration sampler_model = _wrap_model_with_guidance( model, guidance_mode, rescale_multiplier, momentum_beta, cfg_curve, perp_damp, use_zero_init=bool(use_zero_init), zero_init_steps=int(zero_init_steps), fdg_low=float(fdg_low), fdg_high=float(fdg_high), fdg_sigma=float(fdg_sigma), midfreq_enable=bool(midfreq_enable), midfreq_gain=float(midfreq_gain), midfreq_sigma_lo=float(midfreq_sigma_lo), midfreq_sigma_hi=float(midfreq_sigma_hi), ze_zero_steps=int(ze_res_zero_steps), ze_adaptive=bool(ze_adaptive), ze_r_switch_hi=float(ze_r_switch_hi), ze_r_switch_lo=float(ze_r_switch_lo), fdg_low_adaptive=bool(fdg_low_adaptive), fdg_low_min=float(fdg_low_min), fdg_low_max=float(fdg_low_max), fdg_ema_beta=float(fdg_ema_beta), use_local_mask=False, mask_inside=1.0, mask_outside=1.0, mahiro_plus_enable=bool(muse_blend), mahiro_plus_strength=float(muse_blend_strength), eps_scale_enable=bool(eps_scale_enable), eps_scale=float(eps_scale), cwn_enable=bool(cwn_enable), alpha_c=float(alpha_c), alpha_u=float(alpha_u), agc_enable=bool(agc_enable), agc_tau=float(agc_tau), nag_fb_enable=bool(use_nag), nag_fb_scale=float(nag_scale), nag_fb_tau=float(nag_tau), nag_fb_alpha=float(nag_alpha) ) # early interruption check before starting the loop try: model_management.throw_exception_if_processing_interrupted() except Exception: # ensure finally-block cleanup runs and exception propagates raise for i in range(iterations): # cooperative cancel at the start of each iteration model_management.throw_exception_if_processing_interrupted() if i % 2 == 0: clear_gpu_and_ram_cache() # Reset guidance internal state so each iteration starts clean try: if hasattr(sampler_model, "mg_guidance_reset"): sampler_model.mg_guidance_reset() except Exception: pass prev_samples = current_latent["samples"].clone().detach() iter_seed = seed + i * 7777 if noise_offset > 0.0: # Deterministic noise offset tied to iter_seed fade = 1.0 - (i / max(1, iterations)) try: gen = torch.Generator(device='cpu') except Exception: gen = torch.Generator() gen.manual_seed(int(iter_seed) & 0xFFFFFFFF) eps = torch.randn( size=current_latent["samples"].shape, dtype=current_latent["samples"].dtype, device='cpu', generator=gen, ).to(current_latent["samples"].device) current_latent["samples"] = current_latent["samples"] + (noise_offset * fade) * eps try: del eps except Exception: pass # ONNX pre-sampling detectors removed # CLIPSeg mask (optional) try: if bool(clipseg_enable) and isinstance(clipseg_text, str) and clipseg_text.strip() != "": img_prev2 = safe_decode(vae, current_latent) cmask = _clipseg_build_mask(img_prev2, clipseg_text, int(clipseg_preview), float(clipseg_threshold), float(clipseg_blur), int(clipseg_dilate), float(clipseg_gain), ref_embed if bool(clipseg_ref_gate) else None, clip_vision if bool(clipseg_ref_gate) else None, float(clipseg_ref_threshold)) if cmask is not None: if mask_last is None: fused = cmask else: mask_last, cmask = _align_mask_pair(mask_last, cmask) if clipseg_blend == "replace": fused = cmask elif clipseg_blend == "intersect": fused = (mask_last * cmask).clamp(0, 1) else: fused = (1.0 - (1.0 - mask_last) * (1.0 - cmask)).clamp(0, 1) mask_last = fused om = fused.movedim(-1, 1) area = float(om.mean().item()) if area > 0.005: damp = 1.0 - min(0.10, 0.02 + area * 0.08) current_denoise = max(0.10, current_denoise * damp) current_cfg = max(1.0, current_cfg * (1.0 - 0.005)) # No local guidance toggles here; keep optional mask hook clear except Exception: pass # release heavy temporaries from CLIPSeg path try: del img_prev2 except Exception: pass try: del cmask except Exception: pass try: del fused except Exception: pass try: del om except Exception: pass # Sampler model prepared once above; reuse it here (no-op assignment) sampler_model = sampler_model if str(scheduler) == "MGHybrid": try: # Build ZeSmart hybrid sigmas with safe defaults sigmas = _build_hybrid_sigmas( sampler_model, int(current_steps), str(sampler_name), "hybrid", mix=0.5, denoise=float(current_denoise), jitter=0.01, seed=int(iter_seed), _debug=False, tail_smooth=0.15, auto_hybrid_tail=True, auto_tail_strength=0.4, ) # Prepare latent + noise like in MG_ZeSmartSampler lat_img = current_latent["samples"] lat_img = _sample.fix_empty_latent_channels(sampler_model, lat_img) batch_inds = current_latent.get("batch_index", None) noise = _sample.prepare_noise(lat_img, int(iter_seed), batch_inds) noise_mask = current_latent.get("noise_mask", None) callback = _wrap_interruptible_callback(sampler_model, int(current_steps)) # cooperative cancel just before entering sampler model_management.throw_exception_if_processing_interrupted() disable_pbar = not _utils.PROGRESS_BAR_ENABLED sampler_obj = _samplers.sampler_object(str(sampler_name)) samples = _sample.sample_custom( sampler_model, noise, float(current_cfg), sampler_obj, sigmas, positive, negative, lat_img, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=int(iter_seed) ) current_latent = {**current_latent} current_latent["samples"] = samples except Exception as e: # Before any fallback, propagate user cancel if set try: model_management.throw_exception_if_processing_interrupted() except Exception: globals()["_MG_CANCEL_REQUESTED"] = False raise # Do not swallow user interruption; also check sentinel just in case if isinstance(e, model_management.InterruptProcessingException) or globals().get("_MG_CANCEL_REQUESTED", False): globals()["_MG_CANCEL_REQUESTED"] = False raise # Fallback to original path if anything goes wrong print(f"[CADE2.5][MGHybrid] fallback to common_ksampler due to: {e}") current_latent, = _interruptible_ksampler( sampler_model, iter_seed, int(current_steps), current_cfg, sampler_name, _scheduler_names()[0], positive, negative, current_latent, denoise=current_denoise) else: current_latent, = _interruptible_ksampler( sampler_model, iter_seed, int(current_steps), current_cfg, sampler_name, scheduler, positive, negative, current_latent, denoise=current_denoise) # cooperative cancel right after sampling, before further heavy work model_management.throw_exception_if_processing_interrupted() # release sampler temporaries (best-effort) try: del lat_img except Exception: pass try: del noise except Exception: pass try: del noise_mask except Exception: pass try: del callback except Exception: pass try: del sampler_obj except Exception: pass try: del sigmas except Exception: pass if bool(latent_compare): _cur = current_latent["samples"] _prev = prev_samples try: if _prev.device != _cur.device: _prev = _prev.to(_cur.device) if _prev.dtype != _cur.dtype: _prev = _prev.to(dtype=_cur.dtype) except Exception: pass latent_diff = _cur - _prev rms = torch.sqrt(torch.mean(latent_diff * latent_diff)) drift = float(rms.item()) if drift > float(threshold): overshoot = max(0.0, drift - float(threshold)) damp = 1.0 - min(0.15, overshoot * 2.0) current_denoise = max(0.20, current_denoise * damp) cfg_damp = 0.997 if damp > 0.9 else 0.99 current_cfg = max(1.0, current_cfg * cfg_damp) try: del prev_samples except Exception: pass # AQClip-Lite: adaptive soft clipping in latent space (before decode) try: if bool(aqclip_enable): if 'aq_state' not in locals(): aq_state = None H_override = None if bool(aq_attn) and hasattr(sa_patch, "get_attention_entropy_map"): try: Hm = sa_patch.get_attention_entropy_map(clear=False) if Hm is not None: H_override = F.interpolate(Hm, size=(current_latent["samples"].shape[-2], current_latent["samples"].shape[-1]), mode="bilinear", align_corners=False) except Exception: H_override = None z_new, aq_state = _aqclip_lite( current_latent["samples"], tile=int(aq_tile), stride=int(aq_stride), alpha=float(aq_alpha), ema_state=aq_state, ema_beta=float(aq_ema_beta), H_override=H_override, ) current_latent["samples"] = z_new try: del H_override except Exception: pass try: del Hm except Exception: pass except Exception: pass image = safe_decode(vae, current_latent) # allow cancel between sampling and post-decode logic model_management.throw_exception_if_processing_interrupted() # Polish mode: keep global form (low frequencies) from reference while letting details refine if bool(polish_enable) and (i >= int(polish_start_after)): try: # Prepare tensors img = image ref = reference_image if (reference_image is not None) else img if ref.shape[1] != img.shape[1] or ref.shape[2] != img.shape[2]: # resize reference to match current image ref_n = ref.movedim(-1, 1) ref_n = F.interpolate(ref_n, size=(img.shape[1], img.shape[2]), mode='bilinear', align_corners=False) ref = ref_n.movedim(1, -1) x = img.movedim(-1, 1) r = ref.movedim(-1, 1) # Low/high split via Gaussian blur rad = max(1, int(round(float(polish_sigma) * 2))) low_x = _gaussian_blur_nchw(x, sigma=float(polish_sigma), radius=rad) low_r = _gaussian_blur_nchw(r, sigma=float(polish_sigma), radius=rad) high_x = x - low_x # Mix low from reference and current with ramp # a starts from polish_keep_low_ramp and linearly ramps to polish_keep_low over remaining iterations try: denom = max(1, int(iterations) - int(polish_start_after)) t = max(0.0, min(1.0, (i - int(polish_start_after)) / denom)) except Exception: t = 1.0 a0 = float(polish_keep_low_ramp) at = float(polish_keep_low) a = a0 + (at - a0) * t low_mix = low_r * a + low_x * (1.0 - a) new = low_mix + high_x # Micro-detail injection on tail: very light HF boost gated by edges+depth try: phase = (i + 1) / max(1, int(iterations)) # ramp starts late (>=0.70 of iterations), slightly earlier and wider ramp = max(0.0, min(1.0, (phase - 0.70) / 0.30)) if ramp > 0.0: # fine-scale high-pass micro = x - _gaussian_blur_nchw(x, sigma=0.6, radius=1) # edge gate: suppress near strong edges to avoid halos gray = x.mean(dim=1, keepdim=True) sobel_x = torch.tensor([[[-1,0,1],[-2,0,2],[-1,0,1]]], dtype=gray.dtype, device=gray.device).unsqueeze(1) sobel_y = torch.tensor([[[-1,-2,-1],[0,0,0],[1,2,1]]], dtype=gray.dtype, device=gray.device).unsqueeze(1) gx = F.conv2d(gray, sobel_x, padding=1) gy = F.conv2d(gray, sobel_y, padding=1) mag = torch.sqrt(gx*gx + gy*gy) m_edge = (mag - mag.amin()) / (mag.amax() - mag.amin() + 1e-8) g_edge = (1.0 - m_edge).clamp(0.0, 1.0).pow(0.65) # prefer flats/meso-areas # depth gate: prefer nearer surfaces when depth is available try: sz = (int(img.shape[1]), int(img.shape[2])) if depth_gate_cache.get("size") != sz or depth_gate_cache.get("mask") is None: model_path = os.path.join(os.path.dirname(__file__), '..', 'depth-anything', 'depth_anything_v2_vitl.pth') dm = _cf_build_depth_map(img, res=512, model_path=model_path, hires_mode=True) depth_gate_cache = {"size": sz, "mask": dm} dm = depth_gate_cache.get("mask") if dm is not None: g_depth = (dm.movedim(-1, 1).clamp(0,1)) ** 1.35 else: g_depth = torch.ones_like(g_edge) except Exception: g_depth = torch.ones_like(g_edge) g = (g_edge * g_depth).clamp(0.0, 1.0) micro_boost = 0.018 * ramp # very gentle, slightly higher new = new + micro_boost * (micro * g) except Exception: pass # Edge-lock: protect edges from drift by biasing toward low_mix along edges el = float(polish_edge_lock) if el > 1e-6: # Sobel edge magnitude on grayscale gray = x.mean(dim=1, keepdim=True) sobel_x = torch.tensor([[[-1,0,1],[-2,0,2],[-1,0,1]]], dtype=gray.dtype, device=gray.device).unsqueeze(1) sobel_y = torch.tensor([[[-1,-2,-1],[0,0,0],[1,2,1]]], dtype=gray.dtype, device=gray.device).unsqueeze(1) gx = F.conv2d(gray, sobel_x, padding=1) gy = F.conv2d(gray, sobel_y, padding=1) mag = torch.sqrt(gx*gx + gy*gy) m = (mag - mag.amin()) / (mag.amax() - mag.amin() + 1e-8) # Blend toward low_mix near edges new = new * (1.0 - el*m) + (low_mix) * (el*m) img2 = new.movedim(1, -1).clamp(0,1) # Feed back to latent for next steps current_latent = {"samples": safe_encode(vae, img2)} image = img2 # best-effort release of large temporaries try: del x del r del low_x del low_r del high_x del low_mix del new del micro del gray del sobel_x del sobel_y del gx del gy del mag del m_edge del g_edge del g_depth del g del ref_n del ref del img except Exception: pass try: clear_gpu_and_ram_cache() except Exception: pass except Exception: pass # ONNX detectors removed if reference_clean and (ref_embed is not None) and (i % max(1, ref_cooldown) == 0): try: cur_embed = _encode_clip_image(image, clip_vision, ref_preview) dist = _clip_cosine_distance(cur_embed, ref_embed) if dist > ref_threshold: current_denoise = max(0.10, current_denoise * 0.9) current_cfg = max(1.0, current_cfg * 0.99) except Exception: pass if apply_upscale and current_scale != 1.0: current_latent, image = MagicUpscaleModule().process_upscale( current_latent, vae, upscale_method, current_scale) # After upscale at large sizes, add a tiny HF sprinkle gated by edges+depth try: H, W = int(image.shape[1]), int(image.shape[2]) if max(H, W) > 1536: blur = _gaussian_blur(image, radius=1.0, sigma=0.8) hf = (image - blur).clamp(-1, 1) # Edge gate in image space (luma Sobel) lum = (0.2126 * image[..., 0] + 0.7152 * image[..., 1] + 0.0722 * image[..., 2]) kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=lum.device, dtype=lum.dtype).view(1, 1, 3, 3) ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=lum.device, dtype=lum.dtype).view(1, 1, 3, 3) g = torch.sqrt(F.conv2d(lum.unsqueeze(1), kx, padding=1)**2 + F.conv2d(lum.unsqueeze(1), ky, padding=1)**2).squeeze(1) m = (g - g.amin()) / (g.amax() - g.amin() + 1e-8) g_edge = (1.0 - m).clamp(0,1).pow(0.5).unsqueeze(-1) # Depth gate (once per resolution) try: sz = (H, W) if depth_gate_cache.get("size") != sz or depth_gate_cache.get("mask") is None: model_path = os.path.join(os.path.dirname(__file__), '..', 'depth-anything', 'depth_anything_v2_vitl.pth') dm = _cf_build_depth_map(image, res=512, model_path=model_path, hires_mode=True) depth_gate_cache = {"size": sz, "mask": dm} dm = depth_gate_cache.get("mask") if dm is not None: g_depth = dm.clamp(0,1) ** 1.2 else: g_depth = torch.ones_like(g_edge) except Exception: g_depth = torch.ones_like(g_edge) g_tot = (g_edge * g_depth).clamp(0,1) image = (image + 0.045 * hf * g_tot).clamp(0,1) except Exception: pass current_cfg = max(4.0, current_cfg * (1.0 / current_scale)) current_denoise = max(0.15, current_denoise * (1.0 / current_scale)) current_steps = max(1, current_steps - steps_delta) current_cfg = max(0.0, current_cfg - cfg_delta) current_denoise = max(0.0, current_denoise - denoise_delta) current_scale = max(1.0, current_scale - scale_delta) if apply_upscale and current_scale != 1.0 and max(image.shape[1:3]) > 1024: current_latent = {"samples": safe_encode(vae, image)} finally: # Always disable NAG patch and clear local mask, even on errors try: sa_patch.enable_crossattention_nag_patch(False) except Exception: pass # Turn off attention-entropy probe to avoid holding last maps try: if hasattr(sa_patch, "enable_attention_entropy_capture"): sa_patch.enable_attention_entropy_capture(False) except Exception: pass try: sa_patch.CURRENT_PV_ACCUM = prev_accum except Exception: pass try: CURRENT_ONNX_MASK_BCHW = None except Exception: pass # reset cancel sentinel and cleanup cache try: globals()["_MG_CANCEL_REQUESTED"] = False clear_gpu_and_ram_cache() except Exception: pass # best-effort cleanup of GPU/CPU caches on cancel or error try: clear_gpu_and_ram_cache() except Exception: pass if apply_ids: image, = IntelligentDetailStabilizer().stabilize(image, ids_strength) if apply_sharpen: image, = _sharpen_image(image, 2, 1.0, Sharpnes_strenght) # Mask preview as IMAGE (RGB) if mask_last is None: mask_last = torch.zeros((image.shape[0], image.shape[1], image.shape[2], 1), device=image.device, dtype=image.dtype) onnx_mask_img = mask_last.repeat(1, 1, 1, 3).clamp(0, 1) # Final pass: remove isolated hot whites ("fireflies") without touching real edges/highlights try: image = _despeckle_fireflies(image, thr=0.998, max_iso=4.0/9.0, grad_gate=0.15) except Exception: pass # Under-the-hood preview downscale for UI/output IMAGE to cap RAM during save/preview try: B, H, W, C = image.shape max_side = max(int(H), int(W)) cap = 4096 if max_side > cap: scale = float(cap) / float(max_side) nh = max(1, int(round(H * scale))) nw = max(1, int(round(W * scale))) x = image.movedim(-1, 1) x = F.interpolate(x, size=(nh, nw), mode='bilinear', align_corners=False) image = x.movedim(1, -1).clamp(0, 1).to(dtype=image.dtype) except Exception: pass # Optional: save from node with low PNG compress to reduce RAM spike; ignore UI wiring try: if bool(auto_save): from comfy_api.latest._ui import ImageSaveHelper, FolderType _ = ImageSaveHelper.save_images( [image], filename_prefix=str(save_prefix), folder_type=FolderType.output, cls=ComfyAdaptiveDetailEnhancer25, compress_level=int(save_compress)) except Exception: pass # Cleanup KV pruning state to avoid leaking into other nodes try: if hasattr(sa_patch, "set_kv_prune"): sa_patch.set_kv_prune(False, 1.0, int(kv_min_tokens)) except Exception: pass return current_latent, image, int(current_steps), float(current_cfg), float(current_denoise), onnx_mask_img def _wrap_interruptible_callback(model, steps): base_cb = nodes.latent_preview.prepare_callback(model, int(steps)) def _cb(step, x0, x, total_steps): # mark sentinel so outer layers avoid fallbacks on cancel if model_management.processing_interrupted(): globals()["_MG_CANCEL_REQUESTED"] = True raise model_management.InterruptProcessingException() return base_cb(step, x0, x, total_steps) return _cb def _interruptible_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0): lat_img = _sample.fix_empty_latent_channels(model, latent["samples"]) batch_inds = latent.get("batch_index", None) noise = _sample.prepare_noise(lat_img, int(seed), batch_inds) noise_mask = latent.get("noise_mask", None) callback = _wrap_interruptible_callback(model, int(steps)) # cooperative cancel just before sampler entry model_management.throw_exception_if_processing_interrupted() disable_pbar = not _utils.PROGRESS_BAR_ENABLED samples = _sample.sample( model, noise, int(steps), float(cfg), str(sampler_name), str(scheduler), positive, negative, lat_img, denoise=float(denoise), disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=int(seed) ) out = {**latent} out["samples"] = samples return (out,)