""" Stack visualizer: effective UNet weights per checkpoint × SDXL block (21 columns). Matches Merge semantics: eff[i,k] = model_weight × block_weights[k], missing profile = flat 1.0. """ from __future__ import annotations import colorsys from typing import Any import numpy as np import torch from PIL import Image, ImageDraw, ImageFont from .dare_ties_apply import normalize_weights from .dare_ties_stack import STACK_TYPE N_BLOCKS = 21 SECTION_COLORS = ( [(46, 134, 255)] * 9 # IN — align with profile JS #2e86ff + [(243, 156, 18)] * 3 # MID — #f39c12 + [(39, 174, 96)] * 9 # OUT — #27ae60 ) # Stacked-segment colors (colorsys.hsv_to_rgb): H = golden-angle per child index. # Vivid (hue_per_child): S=0.68; V = VB_SCALE_VIVID * bright_mul, bright_mul = 0.52+0.48*rel, clamp ~0.10..0.95. # Muted: lower S but higher than before for contrast; wider bright_mul spread so UNet depth read clearer. SAT_VIVID = 0.68 SAT_MUTED = 0.58 # was 0.42 — less gray, still softer than vivid VB_SCALE_VIVID = 0.58 VB_SCALE_MUTED = 0.64 VB_CLIP_LO, VB_CLIP_HI = 0.10, 0.95 # bright_mul when brightness_from_unet: rel = block / row max in [0,1] BRIGHT_LO_VIVID, BRIGHT_HI_VIVID = 0.52, 0.48 # 0.52 + 0.48*rel → 0.52..1.0 BRIGHT_LO_MUTED, BRIGHT_HI_MUTED = 0.32, 0.66 # 0.32..0.98 — stronger light/dark split # Legend swatch (fixed V for readability) LEGEND_V_VIVID = 0.62 LEGEND_V_MUTED = 0.70 def _hsv_segment( *, checkpoint_color_mode: str, child_index: int, rel: float, brightness_from_unet: bool, ) -> tuple[float, float, float]: hue = (child_index * 0.618033988749895) % 1.0 if checkpoint_color_mode == "muted": sat = SAT_MUTED vb_scale = VB_SCALE_MUTED if brightness_from_unet: m = BRIGHT_LO_MUTED + BRIGHT_HI_MUTED * rel else: m = 1.0 else: sat = SAT_VIVID vb_scale = VB_SCALE_VIVID if brightness_from_unet: m = BRIGHT_LO_VIVID + BRIGHT_HI_VIVID * rel else: m = 1.0 vb = min(VB_CLIP_HI, max(VB_CLIP_LO, vb_scale * m)) return hue, sat, vb def _hsv_legend_swatch(checkpoint_color_mode: str, child_index: int) -> tuple[float, float, float]: hue = (child_index * 0.618033988749895) % 1.0 if checkpoint_color_mode == "muted": return hue, SAT_MUTED, LEGEND_V_MUTED return hue, SAT_VIVID, LEGEND_V_VIVID def _block_weights_row(entry: dict[str, Any]) -> list[float]: """21 floats per merge rules; missing profile or bad list ≡ flat 1.0.""" prof = entry.get("profile") if not isinstance(prof, dict): return [1.0] * N_BLOCKS bw = prof.get("block_weights") if not isinstance(bw, list): return [1.0] * N_BLOCKS out: list[float] = [] for k in range(N_BLOCKS): if k < len(bw): try: out.append(float(bw[k])) except (TypeError, ValueError): out.append(1.0) else: out.append(1.0) return out def _effective_matrix(entries: list[dict[str, Any]], *, normalize_like_merge: bool) -> np.ndarray: """Shape (n, 21), float64.""" n = len(entries) if normalize_like_merge: weights = normalize_weights(entries, "model_weight") else: weights = [max(0.0, float(e.get("model_weight", 0.0))) for e in entries] m = np.zeros((n, N_BLOCKS), dtype=np.float64) for i in range(n): row_b = _block_weights_row(entries[i]) for k in range(N_BLOCKS): m[i, k] = weights[i] * row_b[k] return m def _stack_fingerprint(entries: list[dict[str, Any]]) -> str: parts: list[str] = [] for e in entries: bw_t = tuple(_block_weights_row(e)) parts.append( f"{e.get('ckpt_name', '')}|{float(e.get('model_weight', 0))}|{bw_t}" ) return "|".join(parts) def _legend_labels_max_width(entries: list[dict[str, Any]], font) -> int: """Pixel width of the widest legend line `[i] full_ckpt_name`.""" d = ImageDraw.Draw(Image.new("RGB", (4, 4))) m = 0 for i, e in enumerate(entries): label = f"[{i}] {e.get('ckpt_name', '')}" bbox = d.textbbox((0, 0), label, font=font) m = max(m, bbox[2] - bbox[0]) return m def _render_preview( entries: list[dict[str, Any]], mat: np.ndarray, *, column_scale: str, checkpoint_color_mode: str, brightness_from_unet: bool, width: int, height: int, ) -> Image.Image: n = len(entries) w_user = max(320, min(8192, int(width))) h = max(200, min(1200, int(height))) font_small = ImageFont.load_default() sw = 14 label_pad = 8 + sw + 6 # left margin + swatch + gap legend_w = label_pad + _legend_labels_max_width(entries, font_small) + 12 legend_w = max(100, legend_w) min_chart_px = 280 w = max(w_user, legend_w + min_chart_px) w = min(w, 8192) chart_w = w - legend_w margin_top = 8 margin_bottom = 28 chart_h = h - margin_top - margin_bottom img = Image.new("RGB", (w, h), (26, 26, 38)) arr = np.asarray(img, dtype=np.float32) / 255.0 x0_chart = legend_w col_px = chart_w / N_BLOCKS # Vertical IN/MID/OUT tint (~8% blend), chart area only alpha_band = 0.085 for k in range(N_BLOCKS): cx0 = int(x0_chart + k * col_px) cx1 = int(x0_chart + (k + 1) * col_px) cx1 = max(cx1, cx0 + 1) br, bg, bb = SECTION_COLORS[k][0] / 255.0, SECTION_COLORS[k][1] / 255.0, SECTION_COLORS[k][2] / 255.0 arr[margin_top : margin_top + chart_h, cx0:cx1, 0] *= 1.0 - alpha_band arr[margin_top : margin_top + chart_h, cx0:cx1, 1] *= 1.0 - alpha_band arr[margin_top : margin_top + chart_h, cx0:cx1, 2] *= 1.0 - alpha_band arr[margin_top : margin_top + chart_h, cx0:cx1, 0] += br * alpha_band arr[margin_top : margin_top + chart_h, cx0:cx1, 1] += bg * alpha_band arr[margin_top : margin_top + chart_h, cx0:cx1, 2] += bb * alpha_band img = Image.fromarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8)) draw = ImageDraw.Draw(img) row_max = np.maximum(mat.max(axis=1), 1e-12) col_sums = mat.sum(axis=0) global_max = float(col_sums.max()) if col_sums.size else 1.0 if global_max <= 0.0: global_max = 1e-12 bar_top = margin_top usable_h = chart_h - 4 y_base = bar_top + usable_h # Baseline and y=1 reference (like profile bars) draw.line([(legend_w, y_base), (w, y_base)], fill=(140, 140, 160), width=1) y_one = bar_top + int(usable_h * 0.25) # crude mid reference within stack area draw.line([(legend_w, y_one), (w, y_one)], fill=(90, 90, 110), width=1) stroke_segments = n > 4 for k in range(N_BLOCKS): cx0 = int(x0_chart + k * col_px) + 1 cw = max(1, int(round(col_px)) - 2) cx1 = cx0 + cw if column_scale == "stacked_share": total_k = float(col_sums[k]) if total_k <= 1e-18: continue scale_k = usable_h / total_k else: total_k = float(col_sums[k]) scale_k = usable_h / global_max y_cursor = float(y_base) for i in range(n): val = float(mat[i, k]) if val <= 0.0: continue pix_h = max(1.0, val * scale_k) y_top_f = y_cursor - pix_h y_top = int(round(y_top_f)) y_bot = int(round(y_cursor)) rel = float(mat[i, k] / row_max[i]) if brightness_from_unet else 1.0 rel = max(0.0, min(1.0, rel)) hue, sat, vb = _hsv_segment( checkpoint_color_mode=checkpoint_color_mode, child_index=i, rel=rel, brightness_from_unet=bool(brightness_from_unet), ) r, g, b = colorsys.hsv_to_rgb(hue, sat, vb) fill = (int(r * 255), int(g * 255), int(b * 255)) draw.rectangle([cx0, y_top, cx1, y_bot], fill=fill) if stroke_segments: draw.line([(cx0, y_top), (cx1, y_top)], fill=(34, 34, 42), width=1) y_cursor = y_top_f # Labels IN/MID/OUT draw.text((legend_w + 6, y_base + 4), "IN", fill=(180, 190, 210), font=font_small) draw.text((legend_w + chart_w // 2 - 10, y_base + 4), "MID", fill=(180, 190, 210), font=font_small) draw.text((w - 44, y_base + 4), "OUT", fill=(180, 190, 210), font=font_small) # Legend (full checkpoint names — canvas width grows via `width` / auto min width above) ty = margin_top + 4 max_rows = max(1, min(n, (chart_h - 8) // 18)) for i in range(min(n, max_rows)): hue, sat, vb = _hsv_legend_swatch(checkpoint_color_mode, i) r, g, b = colorsys.hsv_to_rgb(hue, sat, vb) fc = (int(r * 255), int(g * 255), int(b * 255)) draw.rectangle([8, ty, 8 + sw, ty + sw], fill=fc, outline=(60, 60, 72)) name = str(entries[i].get("ckpt_name", "")) draw.text((8 + sw + 6, ty + 1), f"[{i}] {name}", fill=(210, 210, 220), font=font_small) ty += sw + 6 if n > max_rows: draw.text((8, ty), f"... +{n - max_rows} more", fill=(150, 150, 165), font=font_small) cap = ( f"scale={column_scale} | color={checkpoint_color_mode} | " f"bright_unet={brightness_from_unet} | children={n}" ) draw.text((8, h - 22), cap[:120], fill=(130, 130, 145), font=font_small) return img class DARETIESStackVisualizer: """ Renders stacked effective UNet coefficients (model_weight × profile block curve, flat if no profile). """ DESCRIPTION = ( "Preview-only **IMAGE** of how stack entries combine across **21 SDXL UNet blocks** " "(9 IN + 3 MID + 9 OUT).\n" "\n" "Each column stacks checkpoints bottom→top in stack order; segment height reflects " "**model_weight × block_weights[k]** (missing profile ≡ flat **1.0** per block, same as Merge).\n" "\n" "**normalize_like_merge**: when True, model weights are normalized to sum to 1 before multiplying " "by profiles — matches **DARE TIES Merge Models** with **normalize_model_weights** enabled.\n" "\n" "Does **not** simulate DARE, TIES trim, delta_strength, or CLIP merge." ) @classmethod def INPUT_TYPES(cls): return { "required": { "model_stack": ( STACK_TYPE, {"tooltip": "Output from DARE TIES Model Stack (chainable)."}, ), "normalize_like_merge": ( "BOOLEAN", { "default": False, "tooltip": "If True, normalize model_weight across entries like Merge with normalize_model_weights.", }, ), "column_scale": ( ["stacked_share", "absolute"], { "default": "stacked_share", "tooltip": ( "stacked_share: each column sums to full bar height (mix proportions). " "absolute: heights use one global scale from the largest column total (magnitudes)." ), }, ), "checkpoint_color_mode": ( ["hue_per_child", "muted"], { "default": "hue_per_child", "tooltip": ( "hue_per_child: S=0.68, legend V≈0.62. " "muted: S=0.58, wider block brightness swing + legend V≈0.70 — softer than vivid but higher contrast than legacy muted (S=0.42)." ), }, ), "brightness_from_unet": ( "BOOLEAN", { "default": True, "tooltip": "Within each checkpoint hue, brighten blocks where that child's effective weight is high relative to its own row max.", }, ), "width": ( "INT", { "default": 960, "min": 320, "max": 8192, "step": 8, "tooltip": "Minimum image width; expanded if needed so long checkpoint names in the legend fit.", }, ), "height": ( "INT", {"default": 320, "min": 200, "max": 1200, "step": 8, "tooltip": "Output image height."}, ), }, } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("preview", "caption") FUNCTION = "visualize" CATEGORY = "advanced/model_merging/dare_ties" def visualize( self, model_stack: Any, normalize_like_merge: bool, column_scale: str, checkpoint_color_mode: str, brightness_from_unet: bool, width: int, height: int, ): entries = list(model_stack) if len(entries) == 0: raise ValueError("DARE TIES Stack Visualizer: model_stack is empty.") mat = _effective_matrix(entries, normalize_like_merge=bool(normalize_like_merge)) pil = _render_preview( entries, mat, column_scale=column_scale, checkpoint_color_mode=checkpoint_color_mode, brightness_from_unet=bool(brightness_from_unet), width=int(width), height=int(height), ) rgb = np.asarray(pil.convert("RGB"), dtype=np.float32) / 255.0 # Comfy IMAGE: B H W C image = torch.from_numpy(rgb)[None,] cap = ( f"{len(entries)} children; scale={column_scale}; " f"norm_weights={normalize_like_merge}; color={checkpoint_color_mode}" ) return (image, cap) @classmethod def IS_CHANGED( cls, model_stack: Any, normalize_like_merge: Any, column_scale: Any, checkpoint_color_mode: Any, brightness_from_unet: Any, width: Any, height: Any, **kwargs: Any, ) -> str: try: entries = list(model_stack) except (TypeError, ValueError): return "bad_stack" fp = _stack_fingerprint(entries) meta = "|".join( [ str(bool(normalize_like_merge)), str(column_scale), str(checkpoint_color_mode), str(bool(brightness_from_unet)), str(int(width)), str(int(height)), ] ) return f"{len(entries)}::{fp}::{meta}"