digital-garbageandClaude Sonnet 5.5
3f8607b343
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
93 lines
4.2 KiB
Python
93 lines
4.2 KiB
Python
"""The math: take a negative prompt's direction out of the positive's words. Pure torch, no ComfyUI.
|
|
|
|
H3 samples at CFG 1, where the negative branch is never run, so a negative prompt does nothing. This
|
|
gives it a job: pool it to one unit direction and remove that component from each positive word.
|
|
|
|
Projection, not subtraction: `h - a*(h.n)n` touches only the part of a word that lies along the
|
|
negative (a word at right angles is untouched, and a=1 leaves exactly none of it). Whether a plain
|
|
sentence has a clean linear presence in the encoder is unproven: nouns ("a hat") should work better
|
|
than adjectives ("ugly").
|
|
"""
|
|
|
|
import torch
|
|
|
|
MODES = ("project", "subtract")
|
|
MAX_GAIN = 2.0 # how far a word may be scaled back up: without a cap a word that WAS the unwanted
|
|
# concept returns as amplified rounding error
|
|
|
|
|
|
def text_rows(meta, length):
|
|
"""Indices of the text rows (H3 tags each row's kind, 1 = text), or None for 'every row'."""
|
|
tags = (meta or {}).get("minimax_token_tags")
|
|
if tags is None:
|
|
return None
|
|
flat = tags.reshape(-1)
|
|
if int(flat.numel()) != int(length):
|
|
return None
|
|
idx = (flat == 1).nonzero(as_tuple=False).reshape(-1)
|
|
return idx if idx.numel() else None
|
|
|
|
|
|
def direction(tensor, meta=None):
|
|
"""The unit vector the negative occupies (its text rows pooled), or None."""
|
|
if tensor is None or not hasattr(tensor, "dim"):
|
|
return None
|
|
t = tensor.detach().float()
|
|
t = t[0] if t.dim() == 3 else t
|
|
if t.dim() != 2 or t.shape[0] == 0:
|
|
return None
|
|
idx = text_rows(meta, t.shape[0])
|
|
if idx is not None:
|
|
t = t.index_select(0, idx.to(t.device))
|
|
v = t.mean(dim=0)
|
|
n = float(v.norm())
|
|
return v / n if n and n == n else None
|
|
|
|
|
|
def erase(tensor, unit, strength, meta=None, mode="project", renorm=True):
|
|
"""`tensor` with `unit` taken out of its text rows; the same tensor when nothing can be done."""
|
|
if tensor is None or unit is None or not strength:
|
|
return tensor
|
|
out = tensor.detach().float().clone()
|
|
squeeze = out.dim() == 3
|
|
work = out[0] if squeeze else out
|
|
if work.dim() != 2 or work.shape[-1] != int(unit.shape[-1]):
|
|
return tensor
|
|
u = unit.to(device=work.device, dtype=work.dtype)
|
|
idx = text_rows(meta, work.shape[0])
|
|
rows = work if idx is None else work.index_select(0, idx.to(work.device))
|
|
before = rows.norm(dim=-1, keepdim=True)
|
|
if mode == "subtract":
|
|
rows = rows - float(strength) * before * u # a fraction of the row's own size
|
|
else:
|
|
rows = rows - float(strength) * (rows @ u).unsqueeze(-1) * u
|
|
if renorm:
|
|
rows = rows * (before / rows.norm(dim=-1, keepdim=True).clamp_min(1e-12)).clamp(max=MAX_GAIN)
|
|
if not bool(torch.isfinite(rows).all()):
|
|
return tensor # never hand a non-finite conditioning on
|
|
work = rows if idx is None else work.index_copy(0, idx.to(work.device), rows)
|
|
return (work.unsqueeze(0) if squeeze else work).to(dtype=tensor.dtype, device=tensor.device)
|
|
|
|
|
|
def apply(positive, negative, strength, mode="project", renorm=True):
|
|
"""(conditioning, what happened) for a CONDITIONING list. All entries change or none do."""
|
|
if not strength or not positive:
|
|
return positive, "off"
|
|
if not negative:
|
|
return positive, "there is no negative conditioning to take out"
|
|
try:
|
|
unit = direction(negative[0][0], negative[0][1] if len(negative[0]) > 1 else {})
|
|
except Exception as exc: # noqa: BLE001
|
|
return positive, f"could not read the negative conditioning ({exc})"
|
|
if unit is None:
|
|
return positive, "the negative prompt is empty or has no usable direction, so nothing was taken out"
|
|
out, changed = [], 0
|
|
for entry in positive:
|
|
meta = entry[1] if len(entry) > 1 else {}
|
|
new = erase(entry[0], unit, strength, meta, mode if mode in MODES else "project", renorm)
|
|
changed += new is not entry[0]
|
|
out.append([new, *entry[1:]])
|
|
if not changed:
|
|
return positive, "the conditioning could not be modified (size mismatch with the negative?)"
|
|
return out, f"{mode} {float(strength):.2f} on {changed} conditioning entr{'y' if changed == 1 else 'ies'}"
|