"""CADE 2.5: refined adaptive enhancer with reference clean and 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 ..hard.mg_adaptive import AdaptiveSamplerHelper from ..hard.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 ..hard.mg_upscale_module import MagicUpscaleModule, clear_gpu_and_ram_cache from ..hard.mg_controlfusion import _build_depth_map as _cf_build_depth_map from ..hard.mg_ids import IntelligentDetailStabilizer from .. import mg_sagpu_attention as sa_patch from .preset_loader import get as load_preset from ..hard.mg_controlfusion import _pyracanny as _cf_pyracanny, _build_depth_map as _cf_build_depth # FDG/NAG experimental paths removed for now; keeping code lean _ONNX_RT = None _ONNX_SESS = {} # name -> onnxruntime.InferenceSession _ONNX_WARNED = False _ONNX_DEBUG = False _ONNX_FORCE_CPU = True # pin ONNX to CPU for deterministic behavior _ONNX_COUNT_DEBUG = True # print detected counts (faces/hands/persons) when True (temporarily forced ON) # Lazy CLIPSeg cache _CLIPSEG_MODEL = None _CLIPSEG_PROC = None _CLIPSEG_DEV = "cpu" _CLIPSEG_FORCE_CPU = True # pin CLIPSeg to CPU to avoid device drift # ONNX keypoints (wholebody/pose) parsing toggles (set by UI at runtime) _ONNX_KPTS_ENABLE = False _ONNX_KPTS_SIGMA = 2.5 _ONNX_KPTS_GAIN = 1.5 _ONNX_KPTS_CONF = 0.20 # Per-iteration spatial guidance mask (B,1,H,W) in [0,1]; used by cfg_func when enabled CURRENT_ONNX_MASK_BCHW = None # --- 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) -> 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 proxy: gradient magnitude on channel-mean latent 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 {}) def _try_init_onnx(models_dir: str): """Initialize onnxruntime and load all .onnx models in models_dir. We prefer GPU providers when available, but gracefully fall back to CPU. """ global _ONNX_RT, _ONNX_SESS, _ONNX_WARNED import os if _ONNX_RT is None: try: import onnxruntime as ort _ONNX_RT = ort except Exception: if not _ONNX_WARNED: print("[CADE2.5][ONNX] onnxruntime not available, skipping ONNX detectors.") _ONNX_WARNED = True return False # Build provider preference list try: avail = set(_ONNX_RT.get_available_providers()) except Exception: avail = set() pref = [] if _ONNX_FORCE_CPU: pref = ["CPUExecutionProvider"] else: for p in ("CUDAExecutionProvider", "DmlExecutionProvider", "CPUExecutionProvider"): if p in avail or p == "CPUExecutionProvider": pref.append(p) if _ONNX_DEBUG: try: print(f"[CADE2.5][ONNX] Available providers: {sorted(list(avail))}") print(f"[CADE2.5][ONNX] Provider preference: {pref}") except Exception: pass # Load any .onnx in models_dir try: for fname in os.listdir(models_dir): if not fname.lower().endswith('.onnx'): continue if fname in _ONNX_SESS: continue full = os.path.join(models_dir, fname) try: _ONNX_SESS[fname] = _ONNX_RT.InferenceSession(full, providers=pref) if _ONNX_DEBUG: try: print(f"[CADE2.5][ONNX] Loaded model: {fname}") except Exception: pass except Exception as e: if not _ONNX_WARNED: print(f"[CADE2.5][ONNX] failed to load {fname}: {e}") except Exception as e: if not _ONNX_WARNED: print(f"[CADE2.5][ONNX] cannot list models in {models_dir}: {e}") if not _ONNX_SESS and not _ONNX_WARNED: print("[CADE2.5][ONNX] No ONNX models found in", models_dir) _ONNX_WARNED = True return len(_ONNX_SESS) > 0 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) # --- 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 # Scale-aware window s = max(H, W) / 1024.0 k = 3 if s <= 1.1 else (5 if s <= 2.0 else 7) pad = k // 2 # Value/Saturation from RGB (fast, no colorspace conv required) R, G, Bc = img_bhwc[..., 0], img_bhwc[..., 1], img_bhwc[..., 2] V = torch.maximum(R, torch.maximum(G, Bc)) m = torch.minimum(R, torch.minimum(G, Bc)) S = 1.0 - (m / (V + 1e-6)) # Dynamic bright threshold from top tail; allow manual override for very high thr try: q = float(torch.quantile(V.reshape(-1), 0.9995).item()) thr_eff = float(thr) if float(thr) >= 0.99 else max(float(thr), min(0.997, q)) except Exception: thr_eff = float(thr) v_thr = max(0.985, thr_eff) s_thr = 0.06 cand = (V > v_thr) & (S < s_thr) # gradient gate to protect real edges/highlights lum = (0.2126 * R + 0.7152 * G + 0.0722 * Bc) 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 not cand.any(): return img_bhwc # Prefer connected components (OpenCV) to drop small bright specks try: import cv2 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') # Size threshold grows with k area_max = int(max(3, round((k * k) * 0.8))) 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) # B,H,W (bool) rm = rm.unsqueeze(-1) # B,H,W,1 if not rm.any(): return img_bhwc med = _median_pool3x3_bhwc(img_bhwc) return torch.where(rm, med, img_bhwc) except Exception: # Fallback: isolation via local density dens = F.avg_pool2d(cand.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 = cand & (dens < max_iso_eff) 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 # 1D edge-case: single detection row (post-processed model output) if _ONNX_KPTS_ENABLE and a.ndim == 1 and a.shape[0] >= 6: D = a.shape[0] parsed = False # try [xyxy, conf, cls, kpts] or [xyxy, conf, kpts] or [xyxy, kpts] or [kpts] for offset in (6, 5, 4, 0): t = D - offset if t >= 6 and (t % 3) == 0: k = t // 3 kp = a[offset:offset + 3 * k].reshape(k, 3) parsed = True break if parsed: with np.errstate(invalid='ignore'): maxv = float(np.nanmax(np.abs(kp[:, :2]))) if kp.size else 0.0 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] >= 5: score = boxes[:, 4] # If trailing columns look like probabilities in [0,1], mix the best one; if they look like class ids, ignore if boxes.shape[-1] > 5: try: tail = boxes[:, 5:] tmax = np.max(tail, axis=-1) # Heuristic: treat as probs if within [0,1] and not integer-like; else assume class ids maybe_prob = np.all((tmax >= 0.0) & (tmax <= 1.0)) frac = np.abs(tmax - np.round(tmax)) maybe_classid = (np.mean(frac < 1e-6) > 0.9) and (np.max(tmax) >= 1.0) if maybe_prob and not maybe_classid: score = score * tmax except Exception: pass # Keep top-K by score if available if score is not None: try: order = np.argsort(-score) keep = order[: min(12, 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.05: 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: """Return BHWC single-channel mask [0,1] fused from all auto-detected ONNX models. - Auto-loads any .onnx in models_dir. - Heuristically extracts spatial heatmaps; non-spatial outputs are ignored. - Uses soft-OR fusion across models. Models whose filename contains 'anomaly' are scaled by anomaly_gain. """ if not _try_init_onnx(models_dir): # Explicit hint when debugging counts try: if globals().get("_ONNX_COUNT_DEBUG", False): print("[CADE2.5][ONNX] inactive: onnxruntime not available or init failed") except Exception: pass 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: # Explicit hint when debugging counts try: if globals().get("_ONNX_COUNT_DEBUG", False): print(f"[CADE2.5][ONNX] inactive: no .onnx models loaded from {models_dir}") except Exception: pass return torch.zeros((image_bhwc.shape[0], image_bhwc.shape[1], image_bhwc.shape[2], 1), device=image_bhwc.device, dtype=image_bhwc.dtype) # One-time session summary when counting is enabled if globals().get("_ONNX_COUNT_DEBUG", False): try: names = list(_ONNX_SESS.keys()) short = names[:3] more = "..." if len(names) > 3 else "" print(f"[CADE2.5][ONNX] sessions={len(names)} models={short}{more}") except Exception: pass 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 = [] counts_b: dict[str, int] = {} if globals().get("_ONNX_COUNT_DEBUG", False): try: print(f"[CADE2.5][ONNX] build mask image[{b}] preview={int(max(16, min(1024, preview)))}") except Exception: pass # Prepare base BCHW tensor and default preview size; per-model resize comes later target = int(max(16, min(1024, preview))) xb = img_cpu[b].movedim(-1, 0).unsqueeze(0) # 1,C,H,W 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 # Build per-model sized variants (respect fixed input shapes when provided) th, tw = target, target try: if isinstance(in_shape, (list, tuple)) and len(in_shape) == 4: dd = [] for d in in_shape: try: dd.append(int(d)) except Exception: dd.append(-1) if layout == "NCHW" and dd[2] > 8 and dd[3] > 8: th, tw = int(dd[2]), int(dd[3]) if layout.startswith("NHWC") and dd[1] > 8 and dd[2] > 8: th, tw = int(dd[1]), int(dd[2]) except Exception: th, tw = target, target x_stretch_m = F.interpolate(xb, size=(th, tw), mode='bilinear', align_corners=False).clamp(0, 1) if th == tw: x_letter_m = _letterbox_nchw(xb, th).clamp(0, 1) else: sq = max(th, tw) x_letter_sq = _letterbox_nchw(xb, sq).clamp(0, 1) x_letter_m = F.interpolate(x_letter_sq, size=(th, tw), mode='bilinear', align_corners=False).clamp(0, 1) variants = [ ("stretch-RGB", x_stretch_m), ("letterbox-RGB", x_letter_m), ("stretch-BGR", x_stretch_m[:, [2, 1, 0], :, :]), ("letterbox-BGR", x_letter_m[:, [2, 1, 0], :, :]), ] # 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) # Heuristic rejection of stripe artifacts (horizontal/vertical banding) try: rm = np.mean(hm, axis=1) # row means (H) cm = np.mean(hm, axis=0) # col means (W) rstd = float(np.std(rm)) cstd = float(np.std(cm)) zig_r = float(np.mean(np.abs(np.diff(rm)))) zig_c = float(np.mean(np.abs(np.diff(cm)))) horiz_bands = (rstd > 10.0 * max(cstd, 1e-6)) and (zig_r > 0.02) vert_bands = (cstd > 10.0 * max(rstd, 1e-6)) and (zig_c > 0.02) if horiz_bands or vert_bands: if _ONNX_DEBUG: print(f"[CADE2.5][ONNX] '{name}' rejected as stripe artifact (rstd={rstd:.4f} cstd={cstd:.4f} zig_r={zig_r:.4f} zig_c={zig_c:.4f})") hm = None except Exception: pass tmask = _np_to_mask_tensor(hm, H, W, device, dtype) if tmask is not None: masks_b.append(tmask) # Optional counting for debugging: derive counts from raw outputs if possible try: if globals().get("_ONNX_COUNT_DEBUG", False): lower = name.lower() is_face = ("face" in lower) is_hand = ("hand" in lower) is_pose = any(w in lower for w in ["wholebody", "pose", "person", "body"]) and not (is_face or is_hand) boxes_cnt = 0 persons_cnt = 0 wrists_cnt = 0 try: for outx in outs: arr0 = np.asarray(outx) if arr0 is None: continue a = arr0.astype(np.float32, copy=False) while a.ndim > 2 and a.shape[0] == 1: try: a = np.squeeze(a, axis=0) except Exception: break # Face/Hand: count only plausible (N,D) detection tables, not spatial maps/grids if (is_face or is_hand) and a.ndim == 2: N, D = a.shape if D >= 4 and D <= 64 and N <= 512 and not (N >= 32 and D >= 32): n = N # Try to locate a confidence column conf = None if D >= 5: c = a[:, 4] if np.all(np.isfinite(c)) and 0.0 <= float(np.nanmin(c)) and float(np.nanmax(c)) <= 1.0: conf = c if conf is None: c = a[:, -1] if np.all(np.isfinite(c)) and 0.0 <= float(np.nanmin(c)) and float(np.nanmax(c)) <= 1.0: conf = c if conf is not None: n = int(np.sum(conf >= 0.25)) boxes_cnt = max(boxes_cnt, int(n)) # Pose: prefer keypoints formats only if is_pose: # Packed per row [N, K*3] if a.ndim == 2 and a.shape[-1] >= 6 and (a.shape[-1] % 3) == 0: persons_cnt = max(persons_cnt, int(a.shape[0])) try: K = a.shape[-1] // 3 if K >= 11: lw = a[:, 9*3 + 2] rw = a[:, 10*3 + 2] wrists_cnt = max(wrists_cnt, int(np.sum(lw >= 0.2) + np.sum(rw >= 0.2))) except Exception: pass # [N,K,2/3] if a.ndim == 3 and a.shape[-1] in (2, 3): persons_cnt = max(persons_cnt, int(a.shape[0])) try: if a.shape[1] >= 11 and a.shape[-1] == 3: lw = a[:, 9, 2] rw = a[:, 10, 2] wrists_cnt = max(wrists_cnt, int(np.sum(lw >= 0.2) + np.sum(rw >= 0.2))) except Exception: pass except Exception: pass # Map to categories by model name using derived counts if is_face: if boxes_cnt > 0: counts_b["faces"] = counts_b.get("faces", 0) + boxes_cnt elif is_hand: if boxes_cnt > 0: counts_b["hands"] = counts_b.get("hands", 0) + boxes_cnt elif is_pose: if persons_cnt > 0: counts_b["persons"] = counts_b.get("persons", 0) + persons_cnt # Fallback hands from wrists or 2 per person if wrists_cnt > 0: counts_b["hands"] = counts_b.get("hands", 0) + wrists_cnt elif persons_cnt > 0: counts_b["hands"] = counts_b.get("hands", 0) + (2 * persons_cnt) except Exception: pass 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)) if _ONNX_DEBUG or globals().get("_ONNX_COUNT_DEBUG", False): try: print(f"[CADE2.5][ONNX] Detected (image[{b}]): none (no contributing models)") except Exception: pass 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 or globals().get("_ONNX_COUNT_DEBUG", False): try: area = float(fused.movedim(-1,1).mean().item()) if _ONNX_DEBUG: print(f"[CADE2.5][ONNX] Fused area (image[{b}])={area:.4f}") # Print per-image counts if requested if globals().get("_ONNX_COUNT_DEBUG", False): if counts_b: faces = counts_b.get("faces", 0) hands = counts_b.get("hands", 0) persons = counts_b.get("persons", 0) print(f"[CADE2.5][ONNX] Detected (image[{b}]): faces={faces} hands={hands} persons={persons}") else: print(f"[CADE2.5][ONNX] Detected (image[{b}]): counts unavailable (no categories or cv2 missing), area={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=128, to_fp32: bool = False): # Ensure we don't build autograd graphs during final decode steps 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 detach to release VRAM/graphs early try: try: out = out.detach() except Exception: pass out_cpu = out try: out_cpu = out_cpu.to('cpu') except Exception: pass # Optional: force fp32 decode output (after moving to CPU to save VRAM) try: if bool(to_fp32) and out_cpu.dtype != torch.float32: out_cpu = out_cpu.float() 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 _match_latent_channels(vae, latent: dict, model=None): """Align latent channel count to model/VAE expectations (e.g., FLUX/Z_image 16ch) with variance preservation.""" if not isinstance(latent, dict) or ("samples" not in latent): return latent z = latent.get("samples", None) if z is None: return latent try: target_c = None # Prefer model latent_format if available (more reliable than VAE decoder) if model is not None: try: lf = model.get_model_object("latent_format") target_c = int(getattr(lf, "latent_channels", None) or 0) or None except Exception: target_c = None fs = getattr(vae, "first_stage_model", None) dec = getattr(fs, "decoder", None) if dec is not None and hasattr(dec, "conv_in"): target_c = target_c or int(dec.conv_in.in_channels) if target_c is None and hasattr(fs, "latent_channels"): target_c = int(getattr(fs, "latent_channels")) if target_c is None and hasattr(vae, "latent_channels"): target_c = int(getattr(vae, "latent_channels")) if target_c is None: return latent cur_c = int(z.shape[1]) if cur_c == target_c: return latent # Repeat channels when divisible (common case: 4 -> 16) if target_c % cur_c == 0 and cur_c > 0: rep = target_c // cur_c reps = [1, rep] + [1] * (z.ndim - 2) z_fixed = z.repeat(*reps) # Preserve variance after channel replication z_fixed = z_fixed / (rep ** 0.5) else: # Fallback: pad zeros or slice to match if target_c > cur_c: pad = target_c - cur_c pad_tensor = torch.zeros_like(z[:, :1, ...]).repeat(1, pad, *([1] * (z.ndim - 2))) z_fixed = torch.cat([z, pad_tensor], dim=1) else: z_fixed = z[:, :target_c, ...] latent = {**latent, "samples": z_fixed} except Exception: pass return latent def _harmonize_cond_tokens(cond_list): """Pad/truncate cond tokens + masks to a common length to avoid mismatches (e.g., 499 vs 528 or 981 vs 1286).""" if not isinstance(cond_list, list): return cond_list # pass 1: find max token length across cross_attn max_len = 0 for c in cond_list: if isinstance(c, dict): ca = c.get("cross_attn", None) if ca is not None: try: max_len = max(max_len, int(ca.shape[1])) except Exception: pass if max_len <= 0: return cond_list fixed = [] for c in cond_list: if not isinstance(c, dict): fixed.append(c) continue d = c.copy() ca = d.get("cross_attn", None) am = d.get("attention_mask", None) # Harmonize cross_attn length if ca is not None: try: ca_len = int(ca.shape[1]) if ca_len < max_len: pad_shape = list(ca.shape) pad_shape[1] = max_len - ca_len ca_pad = torch.zeros(pad_shape, device=ca.device, dtype=ca.dtype) ca = torch.cat([ca, ca_pad], dim=1) elif ca_len > max_len: ca = ca[:, :max_len, ...] d["cross_attn"] = ca except Exception: pass # Harmonize mask length to cross_attn length if ca is not None: ca_len = int(ca.shape[1]) if am is None: am = torch.ones((ca.shape[0], ca_len), device=ca.device, dtype=ca.dtype) try: am_len = int(am.shape[-1] if am.dim() == 2 else am.shape[1]) if am_len < ca_len: pad = ca_len - am_len pad_shape = list(am.shape) pad_shape[-1] = pad pad_tensor = torch.zeros(pad_shape, device=am.device, dtype=am.dtype) am = torch.cat([am, pad_tensor], dim=-1) elif am_len > ca_len: am = am[..., :ca_len] d["attention_mask"] = am try: d["num_tokens"] = int(torch.count_nonzero(am, dim=-1).max().item()) except Exception: d["num_tokens"] = ca_len except Exception: pass fixed.append(d) return fixed def _summarize_conds(label, conds): out = [] if isinstance(conds, list): for idx, c in enumerate(conds): try: ca = c.get("cross_attn", None) if isinstance(c, dict) else None am = c.get("attention_mask", None) if isinstance(c, dict) else None out.append(f"{label}[{idx}]: ca={None if ca is None else list(ca.shape)}, am={None if am is None else list(am.shape)}") except Exception: pass return "; ".join(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) with torch.inference_mode(): if min(Ht, Wt) > 1024: ov = 128 if max(Ht, Wt) > 2048 else ovlp out = vae.encode_tiled(x[:, :, :, :3], tile_x=tile, tile_y=tile, overlap=ov) else: out = vae.encode(x[:, :, :, :3]) try: torch.cuda.synchronize() if torch.cuda.is_available() else None except Exception: pass return out 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 _soft_symmetry_blend(image_bhwc: torch.Tensor, mask_bhwc: torch.Tensor, alpha: float = 0.03, lp_sigma: float = 1.5) -> torch.Tensor: """Gently mix a mirrored low-frequency component inside mask. - image_bhwc: [B,H,W,C] in [0,1] - mask_bhwc: [B,H,W,1] in [0,1] """ try: if image_bhwc is None or mask_bhwc is None: return image_bhwc if image_bhwc.ndim != 4 or mask_bhwc.ndim != 4: return image_bhwc B, H, W, C = image_bhwc.shape if C < 3: return image_bhwc # Mirror along width mirror = torch.flip(image_bhwc, dims=[2]) # Low-pass both x = image_bhwc.movedim(-1, 1) y = mirror.movedim(-1, 1) rad = max(1, int(round(lp_sigma))) x_lp = _gaussian_blur_nchw(x, sigma=float(lp_sigma), radius=rad) y_lp = _gaussian_blur_nchw(y, sigma=float(lp_sigma), radius=rad) # High-pass from original hp = x - x_lp # Blend LPs inside mask m = mask_bhwc.movedim(-1, 1).clamp(0, 1) a = float(max(0.0, min(0.2, alpha))) base = (1.0 - a * m) * x_lp + (a * m) * y_lp res = (base + hp).movedim(1, -1).clamp(0, 1) return res.to(image_bhwc.dtype) except Exception: return image_bhwc 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_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 _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]: 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 _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, cfg_sched_type: str = "off", cfg_sched_min: float = 0.0, cfg_sched_max: float = 0.0, cfg_sched_gamma: float = 1.5, cfg_sched_u_pow: float = 1.0, cwn_enable: bool = True, alpha_c: float = 1.0, alpha_u: float = 1.0, agc_enable: bool = True, agc_tau: float = 2.8, 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 _reset_state(): 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", _reset_state) 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 in noise-space (when CrossAttention patch is inactive) --- if bool(nag_fb_enable): try: active = bool(getattr(sa_patch, "_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 # Compute effective cond scale before any branch, so schedules apply in all modes cond_scale_eff = cond_scale curve_gain = 1.0 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) g = 1.0 + 0.15 * float(cfg_curve) * s_curve if g.ndim > 0: g = g.mean().item() curve_gain = float(g) cond_scale_eff = cond_scale * curve_gain if isinstance(cfg_sched_type, str) and cfg_sched_type.lower() != "off" 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 = (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) if t.ndim > 0: t_val = float(t.mean().item()) else: t_val = float(t.item()) cmin = float(max(0.0, cfg_sched_min)) cmax = float(max(cmin, cfg_sched_max)) tp = cfg_sched_type.lower() if tp == "cosine": import math cfg_val = cmax - (cmax - cmin) * 0.5 * (1.0 + math.cos(math.pi * t_val)) elif tp in ("warmup", "warm-up", "linear"): g = float(max(0.0, min(1.0, t_val))) ** float(max(0.1, cfg_sched_gamma)) cfg_val = cmin + (cmax - cmin) * g elif tp in ("u", "u-shape", "ushape"): e = 4.0 * (t_val - 0.5) * (t_val - 0.5) e = float(min(1.0, max(0.0, e))) e = e ** float(max(0.1, cfg_sched_u_pow)) cfg_val = cmin + (cmax - cmin) * e else: cfg_val = cond_scale_eff cond_scale_eff = float(cfg_val) * float(curve_gain) except Exception: pass # 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: energy align cond/uncond 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) # --- AGC for CFGZero branches --- if bool(agc_enable): try: t = float(max(0.5, agc_tau)) resid = t * torch.tanh(resid / t) except Exception: pass noise_pred = uncond * alpha + cond_scale_eff * 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 (AGC) --- 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 blended on top if bool(midfreq_enable) and abs(float(midfreq_gain)) > 1e-6: lo_b, mid_b, hi_b = _fdg_split_three(delta, sigma_lo=float(midfreq_sigma_lo), sigma_hi=float(midfreq_sigma_hi), radius=1) lg = _local_gain_for((cond.shape[-2], cond.shape[-1])) if lg is not None: mid_b = mid_b * lg.expand(-1, mid_b.shape[1], -1, -1) delta_fdg = delta_fdg + float(midfreq_gain) * mid_b 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 # 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 (more stable than eps-space) if bool(cwn_enable): try: _eps = 1e-6 rc = (v_cond.pow(2).mean(dim=(1,2,3), keepdim=True).sqrt() + _eps) ru = (v_uncond.pow(2).mean(dim=(1,2,3), keepdim=True).sqrt() + _eps) 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 (Muse Blend), 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 def _edge_mask(image_bhwc: torch.Tensor, threshold: float = 0.20, blur: float = 1.0) -> torch.Tensor: """Return a simple edge mask BHWC [0,1] using Sobel magnitude on luminance. It is resolution-agnostic and intentionally lightweight. """ try: img = image_bhwc lum = (0.2126 * img[..., 0] + 0.7152 * img[..., 1] + 0.0722 * img[..., 2]) kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=img.device, dtype=img.dtype).view(1, 1, 3, 3) ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=img.device, dtype=img.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)) g = g.squeeze(1) # Robust normalization via 98th percentile try: q = torch.quantile(g.flatten(), 0.98).clamp_min(1e-6) except Exception: q = torch.topk(g.flatten(), max(1, int(g.numel() * 0.02))).values.min().clamp_min(1e-6) m = (g / q).clamp(0, 1) if threshold > 0.0: m = (m > float(threshold)).to(img.dtype) if blur > 0.0: rad = int(max(1, min(5, round(float(blur))))) m = _gaussian_blur_nchw(m.unsqueeze(1), sigma=float(max(0.5, blur)), radius=rad).squeeze(1) return m.unsqueeze(-1) except Exception: return torch.zeros((image_bhwc.shape[0], image_bhwc.shape[1], image_bhwc.shape[2], 1), device=image_bhwc.device, dtype=image_bhwc.dtype) def _cf_edges_post(acc_t: torch.Tensor, edge_width: float, edge_smooth: float, edge_single_line: bool, edge_single_strength: float) -> torch.Tensor: try: import cv2, numpy as _np img = (acc_t.clamp(0,1).detach().to('cpu').numpy()*255.0).astype(_np.uint8) # Thickness adjust if float(edge_width) != 0.0: s = float(edge_width) k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3)) op = cv2.dilate if s > 0 else cv2.erode it = int(max(0, min(6, round(abs(s) * 4.0)))) frac = abs(s) * 4.0 - it for _ in range(max(0, it)): img = op(img, k, iterations=1) if frac > 1e-6: y2 = op(img, k, iterations=1) img = ((1.0-frac)*img.astype(_np.float32) + frac*y2.astype(_np.float32)).astype(_np.uint8) # Collapse double lines to single centerline if bool(edge_single_line) and float(edge_single_strength) > 1e-6: try: s = float(edge_single_strength) close = cv2.morphologyEx(img, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)), iterations=1) if hasattr(cv2, 'ximgproc') and hasattr(cv2.ximgproc, 'thinning'): sk = cv2.ximgproc.thinning(close) else: iters = max(1, int(round(2 + 6*s))) sk = _np.zeros_like(close) src = close.copy() elem = cv2.getStructuringElement(cv2.MORPH_CROSS, (3,3)) for _ in range(iters): er = cv2.erode(src, elem, iterations=1) opn = cv2.morphologyEx(er, cv2.MORPH_OPEN, elem) tmp = cv2.subtract(er, opn) sk = cv2.bitwise_or(sk, tmp) src = er if not _np.any(src): break img = ((1.0 - s) * img.astype(_np.float32) + s * sk.astype(_np.float32)).astype(_np.uint8) except Exception: pass # Smooth if float(edge_smooth) > 1e-6: sigma = max(0.1, min(2.0, float(edge_smooth) * 1.2)) img = cv2.GaussianBlur(img, (0,0), sigmaX=sigma) out = torch.from_numpy((img.astype(_np.float32)/255.0)).to(device=acc_t.device, dtype=acc_t.dtype) return out.clamp(0,1) except Exception: # Torch fallback: light blur-only and basic thicken/thin y = acc_t if float(edge_width) > 1e-6: k = max(1, int(round(float(edge_width) * 2))) p = k y = F.max_pool2d(y.unsqueeze(0).unsqueeze(0), kernel_size=2*k+1, stride=1, padding=p)[0,0] if float(edge_width) < -1e-6: k = max(1, int(round(abs(float(edge_width)) * 2))) p = k maxed = F.max_pool2d((1.0 - y).unsqueeze(0).unsqueeze(0), kernel_size=2*k+1, stride=1, padding=p)[0,0] y = 1.0 - maxed if float(edge_smooth) > 1e-6: s = max(1, int(round(float(edge_smooth)*2))) y = F.avg_pool2d(y.unsqueeze(0).unsqueeze(0), kernel_size=2*s+1, stride=1, padding=s)[0,0] return y.clamp(0,1) def _build_cf_edge_mask_from_step(image_bhwc: torch.Tensor, preset_step: str) -> torch.Tensor | None: try: p = load_preset("mg_controlfusion", preset_step) # Safe converters (preset values may be blank strings) def _safe_int(val, default): try: iv = int(val) return iv if iv > 0 else default except Exception: return default def _safe_float(val, default): try: return float(val) except Exception: return default # Read CF params with safe defaults enable_depth = bool(p.get('enable_depth', True)) depth_model_path = str(p.get('depth_model_path', '')) depth_resolution = _safe_int(p.get('depth_resolution', 768), 768) hires_mask_auto = bool(p.get('hires_mask_auto', True)) pyra_low = _safe_int(p.get('pyra_low', 109), 109) pyra_high = _safe_int(p.get('pyra_high', 147), 147) pyra_resolution = _safe_int(p.get('pyra_resolution', 1024), 1024) edge_thin_iter = int(p.get('edge_thin_iter', 0)) edge_boost = _safe_float(p.get('edge_boost', 0.0), 0.0) smart_tune = bool(p.get('smart_tune', False)) smart_boost = _safe_float(p.get('smart_boost', 0.2), 0.2) edge_width = _safe_float(p.get('edge_width', 0.0), 0.0) edge_smooth = _safe_float(p.get('edge_smooth', 0.0), 0.0) edge_single_line = bool(p.get('edge_single_line', False)) edge_single_strength = _safe_float(p.get('edge_single_strength', 0.0), 0.0) edge_depth_gate = bool(p.get('edge_depth_gate', False)) edge_depth_gamma = _safe_float(p.get('edge_depth_gamma', 1.5), 1.5) edge_alpha = _safe_float(p.get('edge_alpha', 1.0), 1.0) # Treat blend_factor as extra gain for edges (depth is not mixed here) blend_factor = _safe_float(p.get('blend_factor', 0.02), 0.02) # ControlNet multipliers: use edge_strength_mul as an additional gain for the edge mask edge_strength_mul = _safe_float(p.get('edge_strength_mul', 1.0), 1.0) # Build edges with CF PyraCanny ed = _cf_pyracanny(image_bhwc, pyra_low, pyra_high, pyra_resolution, edge_thin_iter, edge_boost, smart_tune, smart_boost, preserve_aspect=bool(hires_mask_auto)) ed = _cf_edges_post(ed, edge_width, edge_smooth, edge_single_line, edge_single_strength) # Depth-gate edges if enabled if edge_depth_gate and enable_depth: try: depth = _cf_build_depth(image_bhwc, int(depth_resolution), str(depth_model_path), bool(hires_mask_auto)) g = depth.clamp(0,1) ** float(edge_depth_gamma) ed = (ed * g).clamp(0,1) except Exception: pass # Apply opacity + edge strength (keep blend_factor only for ControlNet stage, not for mask amplitude) total_gain = max(0.0, float(edge_alpha)) * max(0.0, float(edge_strength_mul)) ed = (ed * total_gain).clamp(0,1) # Return BHWC single-channel return ed.unsqueeze(0).unsqueeze(-1) except Exception: return None 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, C = 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, Ca = a_bhw1.shape _, Hb, Wb, Cb = 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 def _mask_dilate(mask_bhw1: torch.Tensor, k: int = 3) -> torch.Tensor: if k <= 1: return mask_bhw1 m = mask_bhw1.movedim(-1, 1) m = F.max_pool2d(m, kernel_size=k, stride=1, padding=k // 2) return m.movedim(1, -1) def _mask_erode(mask_bhw1: torch.Tensor, k: int = 3) -> torch.Tensor: if k <= 1: return mask_bhw1 m = mask_bhw1.movedim(-1, 1) e = 1.0 - F.max_pool2d(1.0 - m, kernel_size=k, stride=1, padding=k // 2) return e.movedim(1, -1) class ComfyAdaptiveDetailEnhancer25: @classmethod def INPUT_TYPES(cls): return { "required": { "preset_step": (["Step 1", "Step 2", "Step 3", "Step 4"], {"default": "Step 1", "tooltip": "Choose the Step preset. Toggle Custom below to apply UI values; otherwise Step preset values are used."}), "custom": ("BOOLEAN", {"default": False, "tooltip": "Custom override: when enabled, your UI values override the selected Step for visible controls; hidden parameters still come from the Step preset."}), "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": 0.7, "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 (beta) unified toggle for Hands/Face/Pose "onnx_detectors": ("BOOLEAN", {"default": False, "tooltip": "Use auto ONNX detectors (any .onnx in models) to refine artifact mask."}), "onnx_preview": ("INT", {"default": 224, "min": 64, "max": 512, "step": 16, "tooltip": "Square preview size fed to ONNX models."}), "onnx_sensitivity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Global gain for fused ONNX mask."}), "onnx_anomaly_gain": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "tooltip": "Extra gain for 'anomaly' models (e.g., anomaly_det.onnx)."}), # 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)."}), # 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}), # 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)."}), # ONNX local guidance (placed last to avoid reordering) "onnx_local_guidance": ("BOOLEAN", {"default": False, "tooltip": "Modulate guidance spatially by ONNX mask."}), "onnx_mask_inside": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Multiplier for guidance inside mask (protects)."}), "onnx_mask_outside": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0, "step": 0.01, "tooltip": "Multiplier for guidance outside mask."}), "onnx_debug": ("BOOLEAN", {"default": False, "tooltip": "Print ONNX mask area per iteration."}), # ONNX wholebody keypoints local heatmap (placed last) "onnx_kpts_enable": ("BOOLEAN", {"default": False, "tooltip": "Parse YOLO wholebody keypoints and add local heatmap."}), "onnx_kpts_sigma": ("FLOAT", {"default": 2.5, "min": 0.5, "max": 8.0, "step": 0.1, "tooltip": "Keypoint Gaussian sigma multiplier."}), "onnx_kpts_gain": ("FLOAT", {"default": 1.5, "min": 0.1, "max": 5.0, "step": 0.1, "tooltip": "Keypoint heat amplitude multiplier."}), "onnx_kpts_conf": ("FLOAT", {"default": 0.20, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Keypoint confidence threshold."}), # Muse Blend global directional post-mix "muse_blend": ("BOOLEAN", {"default": False, "tooltip": "Enable Muse Blend: 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)."}), "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 ONNX mask."}), "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}), # Preview/output image cap (helps RAM during save/preview) "preview_downscale": ("BOOLEAN", {"default": True, "tooltip": "Cap final IMAGE to max 1920 on the longer side to reduce RAM spike during save/preview. Disable for full-res output."}), # Under-the-hood saving (disabled by default to avoid duplicate saves) "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, onnx_detectors=False, onnx_preview=224, onnx_sensitivity=0.5, onnx_anomaly_gain=1.0, guidance_mode="RescaleCFG", rescale_multiplier=0.7, momentum_beta=0.0, cfg_curve=0.0, perp_damp=0.0, use_nag=False, nag_scale=4.0, nag_tau=2.5, nag_alpha=0.25, 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, onnx_local_guidance=False, onnx_mask_inside=1.0, onnx_mask_outside=1.0, onnx_debug=False, onnx_kpts_enable=False, onnx_kpts_sigma=2.5, onnx_kpts_gain=1.5, onnx_kpts_conf=0.20, 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, preview_downscale=False, auto_save=False, save_prefix="ComfyUI", save_compress=1, preset_step="Step 1", custom_override=False): # Cooperative cancel before any heavy work model_management.throw_exception_if_processing_interrupted() # Load base preset for the selected Step. When custom_override is True, # visible UI controls (top-level) are kept from UI; hidden ones still come from preset. try: p = load_preset("mg_cade25", preset_step) if isinstance(preset_step, str) else {} except Exception: p = {} def pv(name, cur, top=False): return cur if (top and bool(custom_override)) else p.get(name, cur) seed = int(pv("seed", seed, top=True)) steps = int(pv("steps", steps, top=True)) cfg = float(pv("cfg", cfg, top=True)) denoise = float(pv("denoise", denoise, top=True)) sampler_name = str(pv("sampler_name", sampler_name, top=True)) scheduler = str(pv("scheduler", scheduler, top=True)) iterations = int(pv("iterations", iterations)) # Smart-seed per-step toggle (defaults to True if not present in preset) smart_seed_enable = bool(pv("smart_seed_enable", True)) smart_seed_k = int(pv("smart_seed_k", 3)) smart_seed_steps = int(pv("smart_seed_steps", 3)) smart_seed_diversity = float(pv("smart_seed_diversity", 0.0)) steps_delta = float(pv("steps_delta", steps_delta)) cfg_delta = float(pv("cfg_delta", cfg_delta)) denoise_delta = float(pv("denoise_delta", denoise_delta)) apply_sharpen = bool(pv("apply_sharpen", apply_sharpen)) apply_upscale = bool(pv("apply_upscale", apply_upscale)) apply_ids = bool(pv("apply_ids", apply_ids)) clip_clean = bool(pv("clip_clean", clip_clean)) ids_strength = float(pv("ids_strength", ids_strength)) upscale_method = str(pv("upscale_method", upscale_method)) scale_by = float(pv("scale_by", scale_by)) scale_delta = float(pv("scale_delta", scale_delta)) noise_offset = float(pv("noise_offset", noise_offset)) threshold = float(pv("threshold", threshold)) Sharpnes_strenght = float(pv("Sharpnes_strenght", Sharpnes_strenght)) latent_compare = bool(pv("latent_compare", latent_compare)) accumulation = str(pv("accumulation", accumulation)) reference_clean = bool(pv("reference_clean", reference_clean)) ref_preview = int(pv("ref_preview", ref_preview)) ref_threshold = float(pv("ref_threshold", ref_threshold)) ref_cooldown = int(pv("ref_cooldown", ref_cooldown)) onnx_detectors = bool(pv("onnx_detectors", onnx_detectors)) onnx_preview = int(pv("onnx_preview", onnx_preview)) onnx_sensitivity = float(pv("onnx_sensitivity", onnx_sensitivity)) onnx_anomaly_gain = float(pv("onnx_anomaly_gain", onnx_anomaly_gain)) guidance_mode = str(pv("guidance_mode", guidance_mode)) rescale_multiplier = float(pv("rescale_multiplier", rescale_multiplier)) momentum_beta = float(pv("momentum_beta", momentum_beta)) cfg_curve = float(pv("cfg_curve", cfg_curve)) perp_damp = float(pv("perp_damp", perp_damp)) use_nag = bool(pv("use_nag", use_nag)) nag_scale = float(pv("nag_scale", nag_scale)) nag_tau = float(pv("nag_tau", nag_tau)) nag_alpha = float(pv("nag_alpha", nag_alpha)) use_zero_init = bool(pv("use_zero_init", use_zero_init)) zero_init_steps = int(pv("zero_init_steps", zero_init_steps)) fdg_low = float(pv("fdg_low", fdg_low)) fdg_high = float(pv("fdg_high", fdg_high)) fdg_sigma = float(pv("fdg_sigma", fdg_sigma)) ze_res_zero_steps = int(pv("ze_res_zero_steps", ze_res_zero_steps)) ze_adaptive = bool(pv("ze_adaptive", ze_adaptive)) ze_r_switch_hi = float(pv("ze_r_switch_hi", ze_r_switch_hi)) ze_r_switch_lo = float(pv("ze_r_switch_lo", ze_r_switch_lo)) fdg_low_adaptive = bool(pv("fdg_low_adaptive", fdg_low_adaptive)) fdg_low_min = float(pv("fdg_low_min", fdg_low_min)) fdg_low_max = float(pv("fdg_low_max", fdg_low_max)) fdg_ema_beta = float(pv("fdg_ema_beta", fdg_ema_beta)) # AQClip-Lite (hidden in Easy UI, controllable via presets) aqclip_enable = bool(pv("aqclip_enable", False)) aq_tile = int(pv("aq_tile", 32)) aq_stride = int(pv("aq_stride", 16)) aq_alpha = float(pv("aq_alpha", 2.0)) aq_ema_beta = float(pv("aq_ema_beta", 0.85)) midfreq_enable = bool(pv("midfreq_enable", False)) midfreq_gain = float(pv("midfreq_gain", 0.0)) midfreq_sigma_lo = float(pv("midfreq_sigma_lo", 0.8)) midfreq_sigma_hi = float(pv("midfreq_sigma_hi", 2.0)) onnx_local_guidance = bool(pv("onnx_local_guidance", onnx_local_guidance)) onnx_mask_inside = float(pv("onnx_mask_inside", onnx_mask_inside)) onnx_mask_outside = float(pv("onnx_mask_outside", onnx_mask_outside)) onnx_debug = bool(pv("onnx_debug", onnx_debug)) onnx_kpts_enable = bool(pv("onnx_kpts_enable", onnx_kpts_enable)) onnx_kpts_sigma = float(pv("onnx_kpts_sigma", onnx_kpts_sigma)) onnx_kpts_gain = float(pv("onnx_kpts_gain", onnx_kpts_gain)) onnx_kpts_conf = float(pv("onnx_kpts_conf", onnx_kpts_conf)) muse_blend = bool(pv("muse_blend", muse_blend)) muse_blend_strength = float(pv("muse_blend_strength", muse_blend_strength)) eps_scale_enable = bool(pv("eps_scale_enable", eps_scale_enable)) eps_scale = float(pv("eps_scale", eps_scale)) clipseg_enable = bool(pv("clipseg_enable", clipseg_enable)) clipseg_text = str(pv("clipseg_text", clipseg_text, top=True)) clipseg_preview = int(pv("clipseg_preview", clipseg_preview)) clipseg_threshold = float(pv("clipseg_threshold", clipseg_threshold)) clipseg_blur = float(pv("clipseg_blur", clipseg_blur)) clipseg_dilate = int(pv("clipseg_dilate", clipseg_dilate)) clipseg_gain = float(pv("clipseg_gain", clipseg_gain)) clipseg_blend = str(pv("clipseg_blend", clipseg_blend)) clipseg_ref_gate = bool(pv("clipseg_ref_gate", clipseg_ref_gate)) clipseg_ref_threshold = float(pv("clipseg_ref_threshold", clipseg_ref_threshold)) # CFG scheduling (internal-only; configured via presets) cfg_sched = str(pv("cfg_sched", "off")) cfg_sched_min = float(pv("cfg_sched_min", max(0.0, cfg * 0.5))) cfg_sched_max = float(pv("cfg_sched_max", cfg)) cfg_sched_gamma = float(pv("cfg_sched_gamma", 1.5)) cfg_sched_u_pow = float(pv("cfg_sched_u_pow", 1.0)) # VAE decode: allow forcing fp32 output (default false) vae_decode_fp32 = bool(pv("vae_decode_fp32", False)) # CWN + AGC defaults (hidden in Easy; controlled via presets) cwn_enable = bool(pv("cwn_enable", True)) alpha_c = float(pv("alpha_c", 1.0)) alpha_u = float(pv("alpha_u", 1.0)) agc_enable = bool(pv("agc_enable", True)) agc_tau = float(pv("agc_tau", 2.8)) # Latent buffer (internal-only; configured via presets) latent_buffer = bool(pv("latent_buffer", True)) lb_inject = float(pv("lb_inject", 0.25)) lb_ema = float(pv("lb_ema", 0.75)) lb_every = int(pv("lb_every", 1)) lb_anchor_every = int(pv("lb_anchor_every", 6)) lb_masked = bool(pv("lb_masked", True)) lb_rebase_thresh = float(pv("lb_rebase_thresh", 0.10)) lb_rebase_rate = float(pv("lb_rebase_rate", 0.25)) polish_enable = bool(pv("polish_enable", polish_enable)) polish_keep_low = float(pv("polish_keep_low", polish_keep_low)) polish_edge_lock = float(pv("polish_edge_lock", polish_edge_lock)) polish_sigma = float(pv("polish_sigma", polish_sigma)) polish_start_after = int(pv("polish_start_after", polish_start_after)) polish_keep_low_ramp = float(pv("polish_keep_low_ramp", polish_keep_low_ramp)) # CADE Seg: per-step toggle to include CF edges into Seg mask seg_use_cf_edges = bool(pv("seg_use_cf_edges", True)) # Hard reset of any sticky globals from prior runs try: global CURRENT_ONNX_MASK_BCHW, _ONNX_KPTS_ENABLE, _ONNX_KPTS_SIGMA, _ONNX_KPTS_GAIN, _ONNX_KPTS_CONF CURRENT_ONNX_MASK_BCHW = None # Reset KPTS toggles to sane defaults; they will be set again if enabled below _ONNX_KPTS_ENABLE = False _ONNX_KPTS_SIGMA = 2.5 _ONNX_KPTS_GAIN = 1.5 _ONNX_KPTS_CONF = 0.20 except Exception: pass # Align latent channels to VAE/model (e.g., Z_image/FLUX use 16ch latents) latent = _match_latent_channels(vae, latent, model) # Harmonize cond token lengths to prevent rare MGHybrid size mismatches positive = _harmonize_cond_tokens(positive) negative = _harmonize_cond_tokens(negative) image = safe_decode(vae, latent, to_fp32=bool(vae_decode_fp32)) # allow user cancel right after initial decode model_management.throw_exception_if_processing_interrupted() 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 # Derive a user-friendly step tag for logs try: _ps = str(preset_step) _num = ''.join(ch for ch in _ps if ch.isdigit()) step_tag = f"Step:{_num}" if _num else _ps except Exception: step_tag = str(preset_step) # Smart seed selection (Sobol + light probing) when effective seed==0 and not in custom override mode try: if int(seed) == 0 and not bool(custom_override) and bool(smart_seed_enable): seed = _smart_seed_select( model, vae, positive, negative, current_latent, str(sampler_name), str(scheduler), float(current_cfg), float(current_denoise), base_seed=0, step_tag=step_tag, k=int(max(1, smart_seed_k)), probe_steps=int(max(1, smart_seed_steps)), clip_vision=clip_vision, reference_image=reference_image, clipseg_text=str(clipseg_text), diversity=float(max(0.0, smart_seed_diversity))) except Exception as e: # propagate user cancel; swallow only non-interrupt errors if isinstance(e, model_management.InterruptProcessingException): raise pass # Visual separation and start marker after seed is finalized try: print("") except Exception: pass try: print(f"\x1b[32m==== {step_tag}, Starting main job ====\x1b[0m") except Exception: pass 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 (read from preset) try: aq_attn = bool(p.get("aq_attn", False)) if isinstance(p, dict) else False if hasattr(sa_patch, "enable_attention_entropy_capture"): sa_patch.enable_attention_entropy_capture(aq_attn, max_tokens=1024, max_heads=4) except Exception: pass # Enable KV pruning for self-attention (read from preset) try: kv_enable = bool(p.get("kv_prune_enable", False)) if isinstance(p, dict) else False kv_keep = float(p.get("kv_keep", 0.85)) if isinstance(p, dict) else 0.85 kv_min_tokens = int(p.get("kv_min_tokens", 128)) if isinstance(p, dict) else 128 if hasattr(sa_patch, "set_kv_prune"): sa_patch.set_kv_prune(kv_enable, kv_keep, kv_min_tokens) except Exception: pass onnx_mask_last = None try: with torch.inference_mode(): __cade_noop = 0 # ensure non-empty with-block # Latent buffer runtime state lb_state = {"z_ema": None, "anchor": None, "drift_last": None, "ref_dist_last": None} # Pre-initialize EMA from the incoming latent so that a 2-iteration node already benefits on iter=1 try: if bool(latent_buffer) and (iterations > 1): z0 = current_latent.get("samples", None) if isinstance(z0, torch.Tensor): lb_state["z_ema"] = z0.clone().detach() lb_state["anchor"] = z0.clone().detach() except Exception: pass # 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 detectors disabled in Easy: prefer CLIPSeg + edge fusion onnx_detectors = False # 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 not None: pre_mask = _mask_to_like(pre_mask, image) cmask = _mask_to_like(cmask, image) 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 # Edge mask from ControlFusion Step (with depth gating) when enabled; fallback to Sobel if bool(seg_use_cf_edges): try: emask = _build_cf_edge_mask_from_step(image, str(preset_step)) except Exception: emask = None if emask is None: try: emask = _edge_mask(image, threshold=0.20, blur=1.0) except Exception: emask = None if emask is not None: if pre_mask is not None: pre_mask, emask = _align_mask_pair(pre_mask, emask) pre_mask = emask if pre_mask is None else (1.0 - (1.0 - pre_mask) * (1.0 - emask)).clamp(0, 1) if pre_mask is not None: onnx_mask_last = pre_mask om = pre_mask.movedim(-1, 1) pre_area = float(om.mean().item()) if bool(onnx_local_guidance): try: if 0.02 <= pre_area <= 0.35: CURRENT_ONNX_MASK_BCHW = om.clamp(0, 1).to(model_management.get_torch_device()) else: CURRENT_ONNX_MASK_BCHW = None except Exception: CURRENT_ONNX_MASK_BCHW = None try: del onnx_mask except Exception: pass try: del om except Exception: pass try: del img_preview except Exception: pass # One-time damping from area (disabled by default) if False: try: if pre_area > 0.005: damp = 1.0 - min(0.04, 0.008 + pre_area * 0.02) current_denoise = max(0.10, current_denoise * damp) current_cfg = max(1.0, current_cfg * (1.0 - 0.003)) except Exception: pass # Preflight symmetry disabled (kept for experiments only) if False: try: img0 = image sym_mask = _clipseg_build_mask(img0, "face | head | torso | shoulders", preview=int(clipseg_preview), threshold=0.45, blur=5.0, dilate=2, gain=1.0) if sym_mask is not None: img_sym = _soft_symmetry_blend(img0, sym_mask, alpha=0.012, lp_sigma=1.75) current_latent = {"samples": safe_encode(vae, img_sym)} image = img_sym except Exception: pass # Compact status try: provs = [] if _ONNX_RT is not None: provs = list(_ONNX_RT.get_available_providers()) clipseg_status = "on" if bool(clipseg_enable) and isinstance(clipseg_text, str) and clipseg_text.strip() != "" else "off" kpts = f"kpts={'on' if bool(onnx_kpts_enable) else 'off'} sigma={float(onnx_kpts_sigma):.2f} gain={float(onnx_kpts_gain):.2f} conf={float(onnx_kpts_conf):.2f}" # print preflight info only in debug sessions (muted by default) if False: print(f"[CADE2.5][preflight] onnx_sessions={len(_ONNX_SESS)} providers={provs if provs else ['CPU']} clipseg={clipseg_status} device={'cpu' if _CLIPSEG_FORCE_CPU else _CLIPSEG_DEV} mask_area={pre_area:.4f} {kpts}") except Exception: pass # Freeze per-iteration external mask rebuild onnx_detectors = False clipseg_enable = False # Depth gate cache for micro-detail injection (reuse per resolution) depth_gate_cache = {"size": None, "mask": None} # 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=bool(onnx_local_guidance), mask_inside=float(onnx_mask_inside), mask_outside=float(onnx_mask_outside), 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), cfg_sched_type=str(cfg_sched), cfg_sched_min=float(cfg_sched_min), cfg_sched_max=float(cfg_sched_max), cfg_sched_gamma=float(cfg_sched_gamma), cfg_sched_u_pow=float(cfg_sched_u_pow), 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) ) # check once more right before the loop starts model_management.throw_exception_if_processing_interrupted() 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 # Pre-sampling ONNX detectors: handled once below (kept compact) # Pre-sampling ONNX detectors (build mask and optionally adjust params for this iteration) if onnx_detectors and (i % max(1, 1) == 0): try: import os models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "models") img_preview = safe_decode(vae, current_latent, to_fp32=bool(vae_decode_fp32)) # Set toggles for this iteration globals()["_ONNX_DEBUG"] = bool(onnx_debug) globals()["_ONNX_COUNT_DEBUG"] = True # force counts ON for debugging session globals()["_ONNX_KPTS_ENABLE"] = bool(onnx_kpts_enable) globals()["_ONNX_KPTS_SIGMA"] = float(onnx_kpts_sigma) globals()["_ONNX_KPTS_GAIN"] = float(onnx_kpts_gain) globals()["_ONNX_KPTS_CONF"] = float(onnx_kpts_conf) onnx_mask = _onnx_build_mask(img_preview, int(onnx_preview), float(onnx_sensitivity), models_dir, float(onnx_anomaly_gain)) if onnx_mask is not None: onnx_mask_last = onnx_mask om = onnx_mask.movedim(-1, 1) area = float(om.mean().item()) if bool(onnx_debug): print(f"[CADE2.5][ONNX] iter={i} mask_area={area:.4f}") if area > 0.005: damp = 1.0 - min(0.25, 0.06 + onnx_sensitivity * 0.05 + area * 0.25) current_denoise = max(0.10, current_denoise * damp) current_cfg = max(1.0, current_cfg * (1.0 - 0.015 * onnx_sensitivity)) # Prepare spatial mask for cfg_func if requested if bool(onnx_local_guidance): # store BCHW mask in global for this iteration (only for reasonable areas) if 0.02 <= area <= 0.35: CURRENT_ONNX_MASK_BCHW = om.clamp(0, 1).to(model_management.get_torch_device()) else: CURRENT_ONNX_MASK_BCHW = None else: CURRENT_ONNX_MASK_BCHW = None except Exception: CURRENT_ONNX_MASK_BCHW = None # CF edge mask (from current image) and fusion (only when enabled) if bool(seg_use_cf_edges): try: img_prev2 = safe_decode(vae, current_latent, to_fp32=bool(vae_decode_fp32)) em2 = _build_cf_edge_mask_from_step(img_prev2, str(preset_step)) if em2 is not None: if onnx_mask_last is None: onnx_mask_last = em2 else: onnx_mask_last, em2 = _align_mask_pair(onnx_mask_last, em2) onnx_mask_last = (1.0 - (1.0 - onnx_mask_last) * (1.0 - em2)).clamp(0, 1) om = onnx_mask_last.movedim(-1, 1) area = float(om.mean().item()) if bool(onnx_local_guidance): if 0.02 <= area <= 0.35: CURRENT_ONNX_MASK_BCHW = om.clamp(0, 1).to(model_management.get_torch_device()) else: CURRENT_ONNX_MASK_BCHW = None except Exception: pass # CLIPSeg mask (optional) and fusion with ONNX try: if bool(clipseg_enable) and isinstance(clipseg_text, str) and clipseg_text.strip() != "": 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 onnx_mask_last is None: fused = cmask else: if clipseg_blend == "replace": fused = cmask elif clipseg_blend == "intersect": onnx_mask_last, cmask = _align_mask_pair(onnx_mask_last, cmask) fused = (onnx_mask_last * cmask).clamp(0, 1) else: onnx_mask_last, cmask = _align_mask_pair(onnx_mask_last, cmask) fused = (1.0 - (1.0 - onnx_mask_last) * (1.0 - cmask)).clamp(0, 1) onnx_mask_last = fused om = fused.movedim(-1, 1) area = float(om.mean().item()) if bool(onnx_debug): print(f"[CADE2.5][MASK] iter={i} mask_area={area:.4f}") if area > 0.005: damp = 1.0 - min(0.10, 0.02 + float(onnx_sensitivity) * 0.02 + area * 0.08) current_denoise = max(0.10, current_denoise * damp) current_cfg = max(1.0, current_cfg * (1.0 - 0.008 * float(onnx_sensitivity))) if bool(onnx_local_guidance): if 0.02 <= area <= 0.35: CURRENT_ONNX_MASK_BCHW = om.clamp(0, 1).to(model_management.get_torch_device()) else: CURRENT_ONNX_MASK_BCHW = None except Exception: pass try: del img_prev2 except Exception: pass try: del em2 except Exception: pass try: del cmask del fused del om except Exception: pass # Latent buffer: pre-sampling injection if bool(latent_buffer) and (iterations > 1) and (i > 0) and (i % max(1, lb_every) == 0): try: z = current_latent["samples"] if lb_state["z_ema"] is None or tuple(lb_state["z_ema"].shape) != tuple(z.shape): lb_state["z_ema"] = z.clone().detach() inj = float(max(0.0, min(0.95, lb_inject))) # optional boost when far from reference (uses last known distance) try: if (ref_embed is not None) and (lb_state.get("ref_dist_last") is not None) and (lb_state["ref_dist_last"] > float(ref_threshold)): inj = min(0.95, inj * (1.0 + min(0.5, (lb_state["ref_dist_last"] - float(ref_threshold)) * 0.75))) except Exception: pass z_ema = lb_state["z_ema"] if bool(lb_masked) and (onnx_mask_last is not None): m = onnx_mask_last.movedim(-1, 1) m = F.interpolate(m, size=(z.shape[-2], z.shape[-1]), mode='bilinear', align_corners=False).clamp(0, 1) m = m.expand(-1, z.shape[1], -1, -1) z = z * (1.0 - inj * m) + z_ema * (inj * m) else: z = z * (1.0 - inj) + z_ema * inj current_latent["samples"] = z except Exception: pass # Sampler model prepared once above; reused across iterations (no-op here) sampler_model = sampler_model # Local best-of-2 in ROI (hands/face), only on early iteration to limit overhead try: do_local_refine = False # disable local best-of-2 by default if do_local_refine: img_roi = safe_decode(vae, current_latent, to_fp32=bool(vae_decode_fp32)) roi = _clipseg_build_mask(img_roi, "hand | hands | face", preview=max(192, int(clipseg_preview//2)), threshold=0.40, blur=5.0, dilate=2, gain=1.0) if roi is None and onnx_mask_last is not None: roi = torch.clamp(onnx_mask_last, 0.0, 1.0) if roi is not None: # Area gating try: ra = float(roi.mean().item()) except Exception: ra = 0.0 if not (0.02 <= ra <= 0.15): raise Exception("ROI area out of range; skip local_refine") # Light erosion to avoid halo influence try: m = roi[..., 0].unsqueeze(1) # disable erosion effect (kernel=1) ero = 1.0 - F.max_pool2d(1.0 - m, kernel_size=1, stride=1, padding=0) roi = ero.clamp(0, 1).movedim(1, -1) except Exception: pass # micro sampling params micro_steps = int(max(2, min(4, round(max(1, current_steps) * 0.05)))) micro_denoise = float(min(0.22, max(0.10, current_denoise * 0.30))) s1 = int((iter_seed ^ 0x9E3779B1) & 0xFFFFFFFFFFFFFFFF) s2 = int((iter_seed ^ 0x85EBCA77) & 0xFFFFFFFFFFFFFFFF) # Candidate A lat_in_a = {"samples": current_latent["samples"].clone()} lat_a, = nodes.common_ksampler( sampler_model, s1, micro_steps, current_cfg, sampler_name, scheduler, positive, negative, lat_in_a, denoise=micro_denoise) img_a = safe_decode(vae, lat_a, to_fp32=bool(vae_decode_fp32)) # Candidate B lat_in_b = {"samples": current_latent["samples"].clone()} lat_b, = nodes.common_ksampler( sampler_model, s2, micro_steps, current_cfg, sampler_name, scheduler, positive, negative, lat_in_b, denoise=micro_denoise) img_b = safe_decode(vae, lat_b, to_fp32=bool(vae_decode_fp32)) # Score inside ROI def _roi_stats(img, roi_mask): try: m = roi_mask[..., 0].clamp(0, 1) R, G, Bc = img[..., 0], img[..., 1], img[..., 2] lum = (0.2126 * R + 0.7152 * G + 0.0722 * Bc) # edges kx = torch.tensor([[-1,0,1],[-2,0,2],[-1,0,1]], device=img.device, dtype=img.dtype).view(1,1,3,3) ky = torch.tensor([[-1,-2,-1],[0,0,0],[1,2,1]], device=img.device, dtype=img.dtype).view(1,1,3,3) gx = F.conv2d(lum.unsqueeze(1), kx, padding=1) gy = F.conv2d(lum.unsqueeze(1), ky, padding=1) g = torch.sqrt(gx*gx + gy*gy).squeeze(1) wmean = (g*m).mean() / (m.mean()+1e-6) # speckles V = torch.maximum(R, torch.maximum(G, Bc)) mi = torch.minimum(R, torch.minimum(G, Bc)) S = 1.0 - (mi / (V + 1e-6)) cand = (V > 0.98) & (S < 0.12) speck = (cand.float()*m).mean() / (m.mean()+1e-6) lmean = (lum*m).mean() / (m.mean()+1e-6) return float(wmean.item()), float(speck.item()), float(lmean.item()) except Exception: return 0.0, 0.5, 0.5 ed_a, sp_a, lm_a = _roi_stats(img_a, roi) ed_b, sp_b, lm_b = _roi_stats(img_b, roi) edge_target = 0.08 score_a = -abs(ed_a - edge_target) - 0.8*sp_a - 0.10*abs(lm_a - 0.5) score_b = -abs(ed_b - edge_target) - 0.8*sp_b - 0.10*abs(lm_b - 0.5) # Optional CLIP-Vision ref boost if ref_embed is not None and clip_vision is not None: try: emb_a = _encode_clip_image(img_a, clip_vision, target_res=224) emb_b = _encode_clip_image(img_b, clip_vision, target_res=224) sim_a = float((emb_a * ref_embed).sum(dim=-1).mean().clamp(-1.0, 1.0).item()) sim_b = float((emb_b * ref_embed).sum(dim=-1).mean().clamp(-1.0, 1.0).item()) score_a += 0.25 * (0.5*(sim_a+1.0)) score_b += 0.25 * (0.5*(sim_b+1.0)) except Exception: pass if score_b > score_a: current_latent = lat_b else: current_latent = lat_a try: del img_roi except Exception: pass try: del roi except Exception: pass try: del lat_in_a del lat_a del img_a except Exception: pass try: del lat_in_b del lat_b del img_b except Exception: pass except Exception: pass 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 = _match_latent_channels(vae, {"samples": lat_img}, sampler_model)["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: try: print(f"[CADE2.5][MGHybrid][debug] sigmas={list(sigmas.shape)} lat={list(current_latent['samples'].shape)}") print(_summarize_conds("pos", positive)) print(_summarize_conds("neg", negative)) except Exception: pass try: traceback.print_exc() except Exception: pass # 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 immediately after sampling model_management.throw_exception_if_processing_interrupted() # Release heavy temporaries from sampler path 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 # Latent buffer: post-sampling EMA update and drift measure try: z_now = current_latent["samples"].detach() if lb_state["z_ema"] is None or tuple(lb_state["z_ema"].shape) != tuple(z_now.shape): lb_state["z_ema"] = z_now.clone() lb_state["anchor"] = z_now.clone() else: lb = float(max(0.0, min(0.99, lb_ema))) lb_state["z_ema"] = lb * lb_state["z_ema"] + (1.0 - lb) * z_now if int(lb_anchor_every) > 0 and ((i + 1) % int(lb_anchor_every) == 0): lb_state["anchor"] = lb_state["z_ema"].clone() except Exception: pass # local RMS drift (independent of UI) try: _cur = current_latent["samples"] _prev = prev_samples if _prev.device != _cur.device: _prev = _prev.to(_cur.device) if _prev.dtype != _cur.dtype: _prev = _prev.to(dtype=_cur.dtype) _diff = _cur - _prev lb_state["drift_last"] = float(torch.sqrt(torch.mean(_diff * _diff)).item()) 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) # Latent buffer: optional rebase toward anchor on overshoot if bool(latent_buffer) and (iterations > 1) and (lb_state.get("anchor") is not None): try: dval = lb_state.get("drift_last", None) if (dval is not None) and (dval > float(lb_rebase_thresh)): rb = float(max(0.0, min(1.0, lb_rebase_rate))) z = current_latent["samples"] a = lb_state["anchor"] if bool(lb_masked) and (onnx_mask_last is not None): m = onnx_mask_last.movedim(-1, 1) m = F.interpolate(m, size=(z.shape[-2], z.shape[-1]), mode='bilinear', align_corners=False).clamp(0, 1) m = m.expand(-1, z.shape[1], -1, -1) z = z * (1.0 - rb * m) + a * (rb * m) else: z = z * (1.0 - rb) + a * rb current_latent["samples"] = z except Exception: pass 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 try: if bool(aq_attn) and hasattr(sa_patch, "get_attention_entropy_map"): 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, to_fp32=bool(vae_decode_fp32)) # and again after decode before post-processing 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 = max(0.0, min(1.0, (phase - 0.70) / 0.30)) if ramp > 0.0: micro = x - _gaussian_blur_nchw(x, sigma=0.6, radius=1) 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) 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 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 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_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 (beta): fuse hands/face/pose mask if available (post-sampling; skip if already set) if onnx_detectors and (i % max(1, 1) == 0) and (onnx_mask_last is None): try: import os models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "models") globals()["_ONNX_DEBUG"] = False globals()["_ONNX_KPTS_ENABLE"] = bool(onnx_kpts_enable) globals()["_ONNX_KPTS_SIGMA"] = float(onnx_kpts_sigma) globals()["_ONNX_KPTS_GAIN"] = float(onnx_kpts_gain) globals()["_ONNX_KPTS_CONF"] = float(onnx_kpts_conf) onnx_mask = _onnx_build_mask(image, int(onnx_preview), float(onnx_sensitivity), models_dir, float(onnx_anomaly_gain)) if onnx_mask is not None: onnx_mask_last = onnx_mask om = onnx_mask.movedim(-1,1) area = float(om.mean().item()) # verbose post-mask log removed; keep single compact log above if area > 0.005: damp = 1.0 - min(0.25, 0.06 + onnx_sensitivity*0.05 + area*0.25) current_denoise = max(0.10, current_denoise * damp) current_cfg = max(1.0, current_cfg * (1.0 - 0.015*onnx_sensitivity)) except Exception: pass 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) # store for next-iter latent buffer injection boost try: lb_state["ref_dist_last"] = float(dist) except Exception: pass 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: # Simple BHWC blur def _gb_bhwc(im: torch.Tensor, radius: float, sigma: float) -> torch.Tensor: if radius <= 0.0: return im pad = int(max(1, round(radius))) ksz = pad * 2 + 1 k = _gaussian_kernel(ksz, sigma, device=im.device).to(dtype=im.dtype) k = k.unsqueeze(0).unsqueeze(0) b, h, w, c = im.shape xch = im.permute(0, 3, 1, 2) y = F.conv2d(F.pad(xch, (pad, pad, pad, pad), mode='reflect'), k.repeat(c, 1, 1, 1), groups=c) return y.permute(0, 2, 3, 1) blur = _gb_bhwc(image, radius=1.0, sigma=0.8) hf = (image - blur).clamp(-1, 1) 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) 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.35 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 (AQClip Attn-mode) to avoid holding last maps try: if hasattr(sa_patch, "enable_attention_entropy_capture"): sa_patch.enable_attention_entropy_capture(False) except Exception: pass # Disable KV pruning as well (avoid leaking state) try: if hasattr(sa_patch, "set_kv_prune"): sa_patch.set_kv_prune(False, 1.0, 128) except Exception: pass try: sa_patch.CURRENT_PV_ACCUM = prev_accum except Exception: pass try: CURRENT_ONNX_MASK_BCHW = None except Exception: pass try: globals()["_MG_CANCEL_REQUESTED"] = False clear_gpu_and_ram_cache() except Exception: pass # best-effort cache cleanup 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) # ONNX mask preview as IMAGE (RGB) if onnx_mask_last is None: onnx_mask_last = torch.zeros((image.shape[0], image.shape[1], image.shape[2], 1), device=image.device, dtype=image.dtype) onnx_mask_img = onnx_mask_last.repeat(1, 1, 1, 3).clamp(0, 1) # Final pass: remove isolated hot whites ("fireflies") without touching real edges/highlights 6.0/9.0, 0.05 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 preview_downscale = False # hard-coded default (can be toggled here if needed) try: if bool(preview_downscale): B, H, W, C = image.shape max_side = max(int(H), int(W)) cap = 1920 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=CADEEasyUI, compress_level=int(save_compress)) except Exception: pass return current_latent, image, int(current_steps), float(current_cfg), float(current_denoise), onnx_mask_img # === Easy UI wrapper: show only top-level controls === class CADEEasyUI(ComfyAdaptiveDetailEnhancer25): @classmethod def INPUT_TYPES(cls): return { "required": { "preset_step": (["Step 1", "Step 2", "Step 3", "Step 4"], {"default": "Step 1", "tooltip": "Choose the Step preset. Toggle Custom below to apply UI values; otherwise Step preset values are used."}), "custom": ("BOOLEAN", {"default": False, "tooltip": "Custom override: when enabled, your UI values override the selected Step for visible controls; hidden parameters still come from the Step preset."}), "model": ("MODEL", {}), "positive": ("CONDITIONING", {}), "negative": ("CONDITIONING", {}), "vae": ("VAE", {}), "latent": ("LATENT", {}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "control_after_generate": True, "tooltip": "Seed 0 = SmartSeed (Sobol + light probe). Non?zero = fixed seed (deterministic)."}), "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": "MGHybrid"}), }, "optional": { # Reference inputs must remain available in Easy "reference_image": ("IMAGE", {}), "clip_vision": ("CLIP_VISION", {}), # CLIPSeg prompt "clipseg_text": ("STRING", {"default": "", "multiline": False, "tooltip": "This field tells what the step should focus on (e.g., hand, feet, face). Separate with commas."}), } } # Easy outputs (hide steps/cfg/denoise) RETURN_TYPES = ("LATENT", "IMAGE", "IMAGE") RETURN_NAMES = ("LATENT", "IMAGE", "mask_preview") FUNCTION = "apply_easy" def apply_easy(self, preset_step, model, positive, negative, vae, latent, seed, steps, cfg, denoise, sampler_name, scheduler, clipseg_text="", reference_image=None, clip_vision=None, custom=False): lat, img, _s, _c, _d, mask = super().apply_cade2( model, vae, positive, negative, latent, int(seed), int(steps), float(cfg), float(denoise), str(sampler_name), str(scheduler), 0.0, preset_step=str(preset_step), custom_override=bool(custom), clipseg_text=str(clipseg_text), reference_image=reference_image, clip_vision=clip_vision, ) return lat, img, mask # Show simpler outputs in Easy variant RETURN_TYPES = ("LATENT", "IMAGE", "IMAGE") RETURN_NAMES = ("LATENT", "IMAGE", "mask_preview") FUNCTION = "apply_easy" def apply_easy(self, preset_step, model, positive, negative, vae, latent, seed, steps, cfg, denoise, sampler_name, scheduler, clipseg_text=""): lat, img, _s, _c, _d, mask = super().apply_cade2( model, vae, positive, negative, latent, int(seed), int(steps), float(cfg), float(denoise), str(sampler_name), str(scheduler), 0.0, preset_step=str(preset_step), custom_override=bool(custom), clipseg_text=str(clipseg_text), ) return lat, img, mask # === Smart seed helpers (Sobol/Halton + light probing) === def _splitmix64(x: int) -> int: x = (x + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF z = x z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9 & 0xFFFFFFFFFFFFFFFF z = (z ^ (z >> 27)) * 0x94D049BB133111EB & 0xFFFFFFFFFFFFFFFF z = z ^ (z >> 31) return z & 0xFFFFFFFFFFFFFFFF def _halton_single(index: int, base: int) -> float: f = 1.0 r = 0.0 i = index while i > 0: f = f / base r = r + f * (i % base) i //= base return r def _sobol_like_2d(n: int, anchor: int) -> tuple[float, float]: # lightweight 2D low-discrepancy via Halton(2,3) scrambled by anchor i = n + 1 + (anchor % 9973) return (_halton_single(i, 2), _halton_single(i, 3)) def _edge_density(img_bhwc: torch.Tensor) -> float: lum = (0.2126 * img_bhwc[..., 0] + 0.7152 * img_bhwc[..., 1] + 0.0722 * img_bhwc[..., 2]) kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=img_bhwc.device, dtype=img_bhwc.dtype).view(1,1,3,3) ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=img_bhwc.device, dtype=img_bhwc.dtype).view(1,1,3,3) gx = F.conv2d(lum.unsqueeze(1), kx, padding=1) gy = F.conv2d(lum.unsqueeze(1), ky, padding=1) g = torch.sqrt(gx*gx + gy*gy) return float(g.mean().item()) def _speckle_fraction(img_bhwc: torch.Tensor) -> float: # reuse S/V-based candidate mask from despeckle logic (no replacement) to estimate fraction 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 = 0.98 s_thr = 0.12 cand = (V > v_thr) & (S < s_thr) return float(cand.float().mean().item()) def _smart_seed_state_path() -> str: import os base = os.path.join(os.path.dirname(__file__), "..", "state") os.makedirs(base, exist_ok=True) return os.path.join(base, "smart_seed.json") def _smart_seed_counter(anchor: int) -> int: import os, json path = _smart_seed_state_path() try: with open(path, "r", encoding="utf-8") as f: data = json.load(f) except Exception: data = {} key = hex(anchor & 0xFFFFFFFFFFFFFFFF) n = int(data.get(key, 0)) data[key] = n + 1 try: with open(path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2) except Exception: pass return n def _smart_seed_select(model, vae, positive, negative, latent, sampler_name: str, scheduler: str, cfg: float, denoise: float, base_seed: int | None = None, k: int = 6, probe_steps: int = 6, clip_vision=None, reference_image=None, clipseg_text: str = "", step_tag: str | None = None, diversity: float = 0.0) -> int: # Log start of SmartSeed selection try: # cooperative cancel before any smart-seed work model_management.throw_exception_if_processing_interrupted() try: # Visual separation before SmartSeed block print("") print("") if step_tag: print(f"\x1b[34m==== {step_tag}, Smart_seed_random: Start (k={int(k)}, steps={int(probe_steps)}, div={float(diversity):.2f}) ====\x1b[0m") else: print(f"\x1b[34m==== Smart_seed_random: Start (k={int(k)}, steps={int(probe_steps)}, div={float(diversity):.2f}) ====\x1b[0m") except Exception: pass # Optional: precompute CLIP-Vision embedding of reference image ref_embed = None if (clip_vision is not None) and (reference_image is not None): try: ref_embed = _encode_clip_image(reference_image, clip_vision, target_res=224) except Exception: ref_embed = None # Anchor from latent shape + sampler/scheduler (+ cfg/denoise) sh = latent["samples"].shape if isinstance(latent, dict) and "samples" in latent else None anchor = 1469598103934665603 # FNV offset basis for v in (sh[2] if sh else 0, sh[3] if sh else 0, len(str(sampler_name)), len(str(scheduler)), int(cfg * 1000), int(denoise * 1000)): anchor = _splitmix64(anchor ^ int(v)) # Advance a persistent counter per anchor to vary indices between runs offset = _smart_seed_counter(anchor) * 7 # Build K candidate seeds from Halton(2,3) cands: list[int] = [] for i in range(k): u, v = _sobol_like_2d(offset + i, anchor) lo = int(u * (1 << 32)) & 0xFFFFFFFF hi = int(v * (1 << 32)) & 0xFFFFFFFF seed64 = _splitmix64((hi << 32) ^ lo ^ anchor) cands.append(seed64 & 0xFFFFFFFFFFFFFFFF) best_seed = cands[0] best_score = -1e9 for sd in cands: # allow user to cancel between candidates model_management.throw_exception_if_processing_interrupted() try: # quick KSampler preview at low steps lat_in = {"samples": latent["samples"].clone()} if isinstance(latent, dict) else latent lat_out, = _interruptible_ksampler( model, int(sd), int(probe_steps), float(cfg), str(sampler_name), str(scheduler), positive, negative, lat_in, denoise=float(min(denoise, 0.65)) ) img = safe_decode(vae, lat_out, to_fp32=bool(vae_decode_fp32)) # and again right after decode model_management.throw_exception_if_processing_interrupted() # Base score: edge density toward a target + low speckle + balanced exposure ed = _edge_density(img) speck = _speckle_fraction(img) lum = float(img.mean().item()) edge_target = 0.10 score = -abs(ed - edge_target) - 2.0 * speck - 0.5 * abs(lum - 0.5) # Deterministic jitter to avoid tie clusters (scaled by diversity) if float(diversity) > 0.0: try: rnd = (_splitmix64(int(sd) ^ int(anchor)) & 0xFFFFFFFF) / 4294967296.0 score += (rnd - 0.5) * float(diversity) except Exception: pass # Perceptual metrics: luminance std and Laplacian variance (downscaled) try: lum_t = (0.2126 * img[..., 0] + 0.7152 * img[..., 1] + 0.0722 * img[..., 2]) lstd = float(lum_t.std().item()) lch = lum_t.unsqueeze(1) lch_small = F.interpolate(lch, size=(128, 128), mode='bilinear', align_corners=False) lap_k = torch.tensor([[0.0, 1.0, 0.0], [1.0, -4.0, 1.0], [0.0, 1.0, 0.0]], device=lch_small.device, dtype=lch_small.dtype).view(1, 1, 3, 3) lap = F.conv2d(lch_small, lap_k, padding=1) lap_var = float(lap.var().item()) score += 0.15 * lstd + 0.10 * lap_var except Exception: pass # Semantic alignment via CLIP-Vision when available if ref_embed is not None and clip_vision is not None: try: cand_embed = _encode_clip_image(img, clip_vision, target_res=224) sim = float((cand_embed * ref_embed).sum(dim=-1).mean().clamp(-1.0, 1.0).item()) sim01 = 0.5 * (sim + 1.0) score += 0.75 * sim01 except Exception: pass # Focus coverage via CLIPSeg when text provided if isinstance(clipseg_text, str) and clipseg_text.strip() != "": try: cmask = _clipseg_build_mask(img, clipseg_text, preview=192, threshold=0.40, blur=5.0, dilate=2, gain=1.0) if cmask is not None: area = float(cmask.mean().item()) cov_target = 0.06 cov_score = 1.0 - min(1.0, abs(area - cov_target) / max(cov_target, 1e-3)) score += 0.30 * cov_score except Exception: pass if score > best_score: best_score = score best_seed = sd try: del img except Exception: pass try: del lat_out except Exception: pass try: del lat_in except Exception: pass try: del lch_small except Exception: pass try: del lap except Exception: pass try: del cand_embed except Exception: pass try: del cmask except Exception: pass except Exception as e: # do not swallow user interruption; also honour sentinel if isinstance(e, model_management.InterruptProcessingException) or globals().get("_MG_CANCEL_REQUESTED", False): globals()["_MG_CANCEL_REQUESTED"] = False raise continue # Log end with selected seed try: if step_tag: print(f"\x1b[34m==== {step_tag}, Smart_seed_random: End. Seed is: {int(best_seed & 0xFFFFFFFFFFFFFFFF)} ====\x1b[0m") else: print(f"\x1b[34m==== Smart_seed_random: End. Seed is: {int(best_seed & 0xFFFFFFFFFFFFFFFF)} ====\x1b[0m") except Exception: pass return int(best_seed & 0xFFFFFFFFFFFFFFFF) except Exception as e: if isinstance(e, model_management.InterruptProcessingException) or globals().get("_MG_CANCEL_REQUESTED", False): globals()["_MG_CANCEL_REQUESTED"] = False # propagate cancel to stop the whole prompt cleanly raise # Fallback to time-based random try: import time fallback_seed = int(_splitmix64(int(time.time_ns()))) except Exception: fallback_seed = int(base_seed or 0) try: if step_tag: print(f"\x1b[34m==== {step_tag}, Smart_seed_random: End. Seed is: {fallback_seed} ====\x1b[0m") else: print(f"\x1b[34m==== Smart_seed_random: End. Seed is: {fallback_seed} ====\x1b[0m") except Exception: pass return fallback_seed def _wrap_interruptible_callback(model, steps): base_cb = nodes.latent_preview.prepare_callback(model, int(steps)) def _cb(step, x0, x, total_steps): model_management.throw_exception_if_processing_interrupted() 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,)