x-stride
fa673ce0d4
SDXL checkpoint merging with DARE/TIES deltas, bezier UNet profiles, stack visualizer, example workflow PNG, and API smoke tests.
415 lines
14 KiB
Python
415 lines
14 KiB
Python
"""
|
||
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}"
|