import torch def _get_samples(latent: dict) -> torch.Tensor: x = latent["samples"] if getattr(x, "is_nested", False): x = x.unbind()[0] return x def _fmt(x: float) -> str: if abs(x) >= 1000: return f"{x:.3e}" return f"{x:.6f}" def _latent_stats_string( x: torch.Tensor, name: str = "latent", top_k: int = 12, ) -> str: """ x: [B,C,H,W] Prints extra diagnostics to help design latent-space ops: - channels with largest std (most active) - channels with largest |mean| (biased) - correlation with an "energy" map (structure/detail proxy) - top correlated channel pairs (good candidates for rotate/mix) """ with torch.no_grad(): B, C, H, W = x.shape device = str(x.device) dtype = str(x.dtype) # Global stats g_min = x.min().item() g_max = x.max().item() g_mean = x.mean().item() g_std = x.std(unbiased=False).item() # Quantiles (global per-batch-item, then avg across batch) flat = x.reshape(B, -1) if flat.shape[1] > 1_000_000: idx = torch.randint( 0, flat.shape[1], (B, 1_000_000), device=x.device ) flat_s = flat.gather(1, idx) else: flat_s = flat qs = torch.tensor( [0.001, 0.01, 0.5, 0.99, 0.999], device=x.device, dtype=x.dtype ) q = torch.quantile(flat_s, qs, dim=1).mean(dim=1) # [len(qs)] # Per-channel mean/std over (B,H,W) ch_mean = x.mean(dim=(0, 2, 3)) ch_std = x.std(dim=(0, 2, 3), unbiased=False) # Top-K channels by std k = min(int(top_k), C) std_vals, std_idx = torch.topk(ch_std, k=k, largest=True) # Top-K channels by |mean| absmean_vals, absmean_idx = torch.topk(ch_mean.abs(), k=k, largest=True) # "Energy map" proxy: per-pixel magnitude averaged over channels # This is a decent proxy for "structure/detail strength" energy = x.abs().mean(dim=1) # [B,H,W] energy_flat = energy.reshape(B, -1) # Correlation of each channel with energy (avg over batch) # corr(c, energy) = cov / (std_c * std_e) # We compute per batch item then average. x_ch = x.reshape(B, C, -1) # [B,C,N] x_ch_mean = x_ch.mean(dim=2, keepdim=True) e_mean = energy_flat.mean(dim=1, keepdim=True) x0 = x_ch - x_ch_mean e0 = energy_flat - e_mean # [B,N] cov = (x0 * e0[:, None, :]).mean(dim=2) # [B,C] x_std_b = x_ch.std(dim=2, unbiased=False).clamp_min(1e-6) # [B,C] e_std_b = energy_flat.std(dim=1, unbiased=False).clamp_min(1e-6) # [B] corr = (cov / (x_std_b * e_std_b[:, None])).mean(dim=0) # [C] corr_abs = corr.abs() corr_vals, corr_idx = torch.topk(corr_abs, k=k, largest=True) # Top correlated channel pairs (for rotate/mix candidates) # Do this on a reduced subset for speed: use the union of top std and top |corr| subset = torch.unique(torch.cat([std_idx, corr_idx], dim=0)) subC = subset.numel() pair_lines = [] if subC >= 2: # Build correlation matrix on subset using flattened vectors (averaged over batch) # We standardize per (B) and then average correlations across batch xs = x_ch[:, subset, :] # [B,subC,N] xs0 = xs - xs.mean(dim=2, keepdim=True) xs_std = xs.std(dim=2, unbiased=False, keepdim=True).clamp_min(1e-6) zs = xs0 / xs_std # [B,subC,N] # corr matrix per batch: (Z @ Z^T) / N # We'll average over B # Result: [subC, subC] corr_mat = torch.einsum("bcn,bdn->bcd", zs, zs).mean(dim=0) / float( zs.shape[2] ) # Take upper triangle pairs, get top few by |corr| iu, ju = torch.triu_indices(subC, subC, offset=1, device=x.device) pair_scores = corr_mat[iu, ju].abs() top_pairs = min(8, pair_scores.numel()) ps, pidx = torch.topk(pair_scores, k=top_pairs, largest=True) for rank in range(top_pairs): i = int(iu[int(pidx[rank])].item()) j = int(ju[int(pidx[rank])].item()) ci = int(subset[i].item()) cj = int(subset[j].item()) raw_corr = float(corr_mat[i, j].item()) pair_lines.append(f"({ci},{cj}): corr={_fmt(raw_corr)}") else: pair_lines.append("(not enough channels in subset)") # Helper format lists def _idxvals(idx_t, val_t, label): parts = [] for i in range(idx_t.numel()): parts.append( f"{int(idx_t[i].item())}:{_fmt(float(val_t[i].item()))}" ) return f"{label}: " + ", ".join(parts) ch_mean_list = ", ".join(_fmt(v.item()) for v in ch_mean) ch_std_list = ", ".join(_fmt(v.item()) for v in ch_std) # Recommendations rec_pairs = [] if pair_lines and "(not enough" not in pair_lines[0]: # first couple pairs are good rotate candidates rec_pairs = pair_lines[:3] s = ( f"{name}: shape=[{B},{C},{H},{W}] dtype={dtype} device={device}\n" f"global: min={_fmt(g_min)} max={_fmt(g_max)} mean={_fmt(g_mean)} std={_fmt(g_std)}\n" f"quantiles(avg batch): q0.1%={_fmt(q[0].item())} q1%={_fmt(q[1].item())} " f"q50%={_fmt(q[2].item())} q99%={_fmt(q[3].item())} q99.9%={_fmt(q[4].item())}\n" f"{_idxvals(std_idx, std_vals, f'top{k}_ch_by_std')}\n" f"{_idxvals(absmean_idx, ch_mean.abs()[absmean_idx], f'top{k}_ch_by_|mean|')}\n" f"{_idxvals(corr_idx, corr_abs[corr_idx], f'top{k}_ch_by_|corr_with_energy|')}\n" f"top_correlated_pairs_in_subset: " + "; ".join(pair_lines) + "\n" f"recommended_rotate_pairs: " + ("; ".join(rec_pairs) if rec_pairs else "(none)") + "\n" f"channel mean: [{ch_mean_list}]\n" f"channel std : [{ch_std_list}]" ) return s class SuperLatentStats: @classmethod def INPUT_TYPES(cls): return { "required": { "latent": ("LATENT",), "label": ("STRING", {"default": "latent"}), "top_k": ( "INT", {"default": 12, "min": 3, "max": 64, "step": 1}, ), } } RETURN_TYPES = ("STRING",) FUNCTION = "run" CATEGORY = "SuperNodes/debug" def run(self, latent, label="latent", top_k=12): x = _get_samples(latent) return (_latent_stats_string(x, name=str(label), top_k=int(top_k)),) class SuperStopExecution: @classmethod def INPUT_TYPES(cls): return { "required": { "message": ( "STRING", { "multiline": False, "default": "Execution stopped by condition.", }, ), "trigger": ( "BOOLEAN", { "default": True, "tooltip": "If True, execution halts. If False, nothing happens.", }, ), } } RETURN_TYPES = () FUNCTION = "halt_execution" OUTPUT_NODE = True CATEGORY = "SuperNodes/debug" def halt_execution(self, message, trigger): if trigger: alert = str(message) raise Exception(f"{alert}") return ()