feat: DynaShift steers late steps away from disliked clips, banked on the taste key
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
53cbc29948
commit
b831372879
@@ -34,6 +34,7 @@ from .latent_upscaler import load_upscaler as _load_upscaler
|
||||
from . import prompt_layout as _prompt_layout
|
||||
from .rows import target_rows as _target_rows
|
||||
from .rope import query_rotation as _query_rotation
|
||||
from .streams import video_stream as _video_stream
|
||||
|
||||
has_block = _traits.has_block
|
||||
|
||||
@@ -193,4 +194,5 @@ PROVIDES = {"empty_latent": empty_latent, "decode": decode, "detect": detect,
|
||||
"text_tokenizer": _prompt_layout.text_tokenizer,
|
||||
"prompt_rows": _prompt_layout.prompt_rows,
|
||||
"video_time_rows": _prompt_layout.video_time_rows,
|
||||
"query_rotation": _query_rotation}
|
||||
"query_rotation": _query_rotation,
|
||||
"video_stream": _video_stream}
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Where H3's video lives in a sampling-time latent: the first of [video, audio]."""
|
||||
|
||||
|
||||
def video_stream(x):
|
||||
"""-> (video [B, C, T, H, W], rebuild(new_video) -> x), or None when `x` is
|
||||
not H3's two-branch latent."""
|
||||
if not getattr(x, "is_nested", False) or len(x.tensors) != 2 or x.tensors[0].dim() != 5:
|
||||
return None
|
||||
from comfy.nested_tensor import NestedTensor
|
||||
rest = x.tensors[1:]
|
||||
return x.tensors[0], lambda video: NestedTensor([video, *rest])
|
||||
@@ -0,0 +1,225 @@
|
||||
"""DynaShift: steer each late step's predicted picture away from clips you disliked.
|
||||
|
||||
A negative prompt for a model that ignores negative prompts, made of your own
|
||||
disliked clips instead of words. Every run banks its final picture latent on the
|
||||
taste key; rating it decides which side of the bank it lands on. Then, over the
|
||||
last half of the steps:
|
||||
|
||||
1. each frame of the current prediction is matched (cosine, on a pooled
|
||||
fingerprint) against every banked disliked frame -- by content, not time, so
|
||||
scene position and length don't matter;
|
||||
2. a disliked clip made from a different prompt counts less;
|
||||
3. a frame that matches above the threshold has the matched frame's direction
|
||||
taken out of it, harder the closer the match. Once the unwanted thing is gone
|
||||
the match drops and it stops by itself.
|
||||
|
||||
With 2+ liked and 2+ disliked clips at this resolution it also nudges every
|
||||
frame along liked-minus-disliked: an average direction, never toward one clip
|
||||
(there is one way to BE a past clip and many ways to not resemble one).
|
||||
|
||||
No extra model pass. Sound untouched by construction.
|
||||
|
||||
User-confirmed on H3 (2026-09-20) at threshold 0.6. Every earlier "does
|
||||
nothing" verdict was measured at 0.95, which is the practically-off end.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from comfy.patcher_extension import WrappersMP
|
||||
|
||||
from ..._core import dit_hooks, log, registry
|
||||
|
||||
ID = "dynashift"
|
||||
TITLE = "DynaShift"
|
||||
MOUNT = "generation.sampling"
|
||||
STAGE = "guidance"
|
||||
CATEGORY = "guidance"
|
||||
STATUS = "proven"
|
||||
REQUIRES = ["temporal_latent"]
|
||||
|
||||
SETTINGS = {
|
||||
"enabled": {
|
||||
"type": "bool", "default": False,
|
||||
"label": "Avoid what I disliked",
|
||||
"hint": "Steers away from clips you rated down, and toward ones you liked.",
|
||||
},
|
||||
"strength": {
|
||||
"type": "float", "default": 0.3, "min": 0.0, "max": 2.0, "step": 0.05,
|
||||
"label": "Strength", "ui": "slider",
|
||||
"hint": "0 = only learn.",
|
||||
"when": {"enabled": True},
|
||||
},
|
||||
"threshold": {
|
||||
"type": "float", "default": 0.6, "min": 0.3, "max": 0.95, "step": 0.05,
|
||||
"label": "How alike before it acts", "ui": "slider",
|
||||
"hint": "Lower acts on looser resemblance. 0.95 almost never acts.",
|
||||
"when": {"enabled": True},
|
||||
},
|
||||
}
|
||||
|
||||
KIND = "dynashift"
|
||||
BANK = 8 # per side: v4's ring size
|
||||
MIN_FOR_PULL = 2
|
||||
_DESC = 512
|
||||
|
||||
|
||||
def _say(message):
|
||||
log.once(f"{ID}:{message}", log.ALERT, "FunPack DynaShift", message)
|
||||
|
||||
|
||||
def _video_of(x):
|
||||
split = registry.current().ask("video_stream", x)
|
||||
if split is not None:
|
||||
return split
|
||||
if torch.is_tensor(x) and x.dim() == 5:
|
||||
return x, lambda video: video
|
||||
return None
|
||||
|
||||
|
||||
def _pooled(cond):
|
||||
if not torch.is_tensor(cond):
|
||||
return None
|
||||
p = cond.detach().float()
|
||||
while p.dim() > 1:
|
||||
p = p.mean(0)
|
||||
return p
|
||||
|
||||
|
||||
def _frames(latent):
|
||||
"""[C, T, H, W] -> [T, C*H*W] float."""
|
||||
return latent.permute(1, 0, 2, 3).reshape(latent.shape[1], -1).float()
|
||||
|
||||
|
||||
def _fingerprint(frames):
|
||||
return F.normalize(F.adaptive_avg_pool1d(frames.unsqueeze(1), _DESC).squeeze(1), dim=-1)
|
||||
|
||||
|
||||
def gate(transformer_options):
|
||||
"""0 over the first half of the steps, rising to 1 at the last: the late-
|
||||
step gate every learning feature shares, by step position so it holds on
|
||||
any schedule (v4's sigma-based one was 0 of 12 steps on H3)."""
|
||||
where = dit_hooks.current_step(transformer_options)
|
||||
if where is None:
|
||||
return 0.0
|
||||
index, total = where
|
||||
return max(0.0, 2.0 * index / max(1, total) - 1.0)
|
||||
|
||||
|
||||
class Bank:
|
||||
"""The key's banked clips, prepared once per (device, C, H, W)."""
|
||||
|
||||
def __init__(self, rows):
|
||||
rated = [r for r in rows if torch.is_tensor(r["rows"].get("latent"))
|
||||
and r["rows"]["latent"].dim() == 4]
|
||||
self.negatives = [r["rows"] for r in rated if r["reward"] < 0][-BANK:]
|
||||
self.positives = [r["rows"] for r in rated if r["reward"] > 0][-BANK:]
|
||||
self._ready = {}
|
||||
|
||||
def at(self, device, c, h, w):
|
||||
key = (str(device), c, h, w)
|
||||
if key not in self._ready:
|
||||
self._ready[key] = self._prepare(device, c, h, w)
|
||||
return self._ready[key]
|
||||
|
||||
def _prepare(self, device, c, h, w):
|
||||
def fits(entry):
|
||||
lat = entry["latent"]
|
||||
return lat.shape[0] == c and lat.shape[2] == h and lat.shape[3] == w
|
||||
|
||||
negs = [e for e in self.negatives if fits(e)]
|
||||
poss = [e for e in self.positives if fits(e)]
|
||||
if len(negs) < len(self.negatives) or len(poss) < len(self.positives):
|
||||
_say(f"{len(self.negatives) - len(negs) + len(self.positives) - len(poss)} banked "
|
||||
"clip(s) skipped: made at a different resolution than this run")
|
||||
if not negs:
|
||||
return None
|
||||
units, descs, owner, conds = [], [], [], []
|
||||
for i, e in enumerate(negs):
|
||||
f = _frames(e["latent"].to(device))
|
||||
descs.append(_fingerprint(f))
|
||||
units.append(F.normalize(f, dim=-1).half())
|
||||
owner += [i] * f.shape[0]
|
||||
conds.append(e.get("cond"))
|
||||
units = torch.cat(units)
|
||||
pull = None
|
||||
if len(negs) >= MIN_FOR_PULL and len(poss) >= MIN_FOR_PULL:
|
||||
liked = torch.cat([F.normalize(_frames(e["latent"].to(device)), dim=-1) for e in poss])
|
||||
diff = liked.mean(0) - units.float().mean(0)
|
||||
n = diff.norm()
|
||||
if torch.isfinite(n) and n > 1e-8:
|
||||
pull = diff / n
|
||||
return torch.cat(descs), units, torch.tensor(owner, device=device), conds, pull
|
||||
|
||||
|
||||
def shift(video, bank, strength, threshold, cond=None):
|
||||
"""One step's x0 video [B, C, T, H, W] -> the shifted one, or None when
|
||||
nothing changed."""
|
||||
b, c, t, h, w = video.shape
|
||||
ready = bank.at(video.device, c, h, w)
|
||||
if ready is None:
|
||||
return None
|
||||
desc, units, owner, conds, pull = ready
|
||||
current = _pooled(cond)
|
||||
weights = torch.ones(len(conds), device=video.device)
|
||||
if current is not None:
|
||||
for i, stored in enumerate(conds):
|
||||
if torch.is_tensor(stored) and stored.numel() == current.numel():
|
||||
weights[i] = F.cosine_similarity(current.to(video.device), stored.to(video.device).float(),
|
||||
dim=0).clamp(0.0, 1.0)
|
||||
out = []
|
||||
for sample in video:
|
||||
f = _frames(sample)
|
||||
best, idx = ((_fingerprint(f) @ desc.T) * weights[owner]).max(dim=1)
|
||||
g = ((best - threshold) / max(1e-6, 1.0 - threshold)).clamp(0.0, 1.0) * strength
|
||||
matched = units[idx].float()
|
||||
coef = (f * matched).sum(-1).clamp(min=0.0) # only frames that point AT it
|
||||
f = f - (g * coef).unsqueeze(1) * matched
|
||||
if pull is not None:
|
||||
f = f + strength * f.norm(dim=-1, keepdim=True) * pull.unsqueeze(0)
|
||||
out.append(f.reshape(t, c, h, w).permute(1, 0, 2, 3))
|
||||
return torch.stack(out).to(video.dtype)
|
||||
|
||||
|
||||
def install(patcher, values, key):
|
||||
if not values.get("enabled"):
|
||||
return None
|
||||
taste = registry.current().ask("taste_store", patcher)
|
||||
if taste is None:
|
||||
_say("off: set a Taste key first -- it is what your ratings teach")
|
||||
return None
|
||||
strength = float(values.get("strength", 0.3))
|
||||
threshold = float(values.get("threshold", 0.6))
|
||||
bank = Bank(taste.rows(KIND))
|
||||
captured = taste.collect(patcher, key, KIND, keep=4 * BANK)
|
||||
|
||||
def apply_model(executor, x, t, *args, **kwargs):
|
||||
out = executor(x, t, *args, **kwargs)
|
||||
to = kwargs.get("transformer_options")
|
||||
split = _video_of(out)
|
||||
if split is None:
|
||||
_say("off this run: could not find the picture in this model's latent")
|
||||
return out
|
||||
video, rebuild = split
|
||||
if dit_hooks.last_step(to):
|
||||
captured["latent"] = video[0].detach().half()
|
||||
cond = _pooled(kwargs.get("c_crossattn"))
|
||||
if cond is not None:
|
||||
captured["cond"] = cond
|
||||
amount = strength * gate(to)
|
||||
if amount <= 0.0:
|
||||
return out
|
||||
shifted = shift(video, bank, amount, threshold, kwargs.get("c_crossattn"))
|
||||
return out if shifted is None else rebuild(shifted)
|
||||
|
||||
patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model)
|
||||
if strength <= 0.0:
|
||||
return "strength 0: learning only"
|
||||
if not bank.negatives:
|
||||
return "learning: nothing disliked banked yet"
|
||||
pull = ("with pull toward liked" if len(bank.positives) >= MIN_FOR_PULL
|
||||
and len(bank.negatives) >= MIN_FOR_PULL else "no pull yet (needs 2 liked + 2 disliked)")
|
||||
return (f"{len(bank.negatives)} disliked, {len(bank.positives)} liked banked; "
|
||||
f"strength {strength:g}, threshold {threshold:g}, last half of steps; {pull}")
|
||||
|
||||
|
||||
PROVIDES = {"modifier": install}
|
||||
@@ -0,0 +1,83 @@
|
||||
"""DynaShift: the step gate, the shift itself, and banking through a taste key."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _needs_comfy(comfyui):
|
||||
"""Imports comfy."""
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def root(tmp_path, monkeypatch):
|
||||
from modules.system.taste import store
|
||||
monkeypatch.setattr(store, "ROOT", tmp_path / "taste")
|
||||
|
||||
|
||||
def _to(index, steps=4):
|
||||
sched = torch.linspace(1.0, 0.0, steps + 1)
|
||||
return {"sample_sigmas": sched, "sigmas": sched[index:index + 1]}
|
||||
|
||||
|
||||
def test_the_gate_opens_over_the_last_half_only():
|
||||
from modules.sampling.dynashift import gate
|
||||
assert [gate(_to(i)) for i in range(4)] == [0.0, 0.0, 0.0, 0.5]
|
||||
assert gate({}) == 0.0
|
||||
|
||||
|
||||
def _bank(neg, pos=()):
|
||||
from modules.sampling.dynashift import Bank
|
||||
rows = [{"reward": -1.0, "rows": {"latent": n}} for n in neg]
|
||||
rows += [{"reward": 1.0, "rows": {"latent": p}} for p in pos]
|
||||
return Bank(rows)
|
||||
|
||||
|
||||
def test_a_frame_like_a_disliked_one_loses_that_likeness_and_others_are_left_alone():
|
||||
from modules.sampling.dynashift import shift
|
||||
torch.manual_seed(0)
|
||||
bad = torch.randn(4, 2, 8, 8)
|
||||
other = torch.randn(4, 1, 8, 8)
|
||||
video = torch.cat([bad[:, :1] + 0.1 * torch.randn(4, 1, 8, 8), other], dim=1)[None]
|
||||
out = shift(video, _bank([bad]), strength=1.0, threshold=0.6)
|
||||
cos = lambda a, b: torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0)
|
||||
assert cos(out[0, :, 0], bad[:, 0]) < 0.5 * cos(video[0, :, 0], bad[:, 0])
|
||||
assert torch.equal(out[0, :, 1], video[0, :, 1]) # never matched: untouched
|
||||
|
||||
|
||||
def test_nothing_banked_at_this_resolution_changes_nothing():
|
||||
from modules.sampling.dynashift import shift
|
||||
assert shift(torch.randn(1, 4, 2, 8, 8), _bank([torch.randn(4, 2, 16, 16)]), 1.0, 0.6) is None
|
||||
|
||||
|
||||
def test_it_banks_the_last_step_and_steers_once_rated(tiny_h3, monkeypatch):
|
||||
from comfy.nested_tensor import NestedTensor
|
||||
from comfy.patcher_extension import WrappersMP
|
||||
from modules.sampling.modifiers.nodes import FunPackLoadModifiers
|
||||
from modules.system.taste import store
|
||||
|
||||
def load():
|
||||
patched, status = FunPackLoadModifiers.execute(tiny_h3.patcher, {
|
||||
"taste": {"key": "fox"}, "dynashift": {"enabled": True, "strength": 1.0}}).result
|
||||
wrap = [w for ws in patched.wrappers[WrappersMP.APPLY_MODEL].values() for w in ws][-1]
|
||||
outer = [w for ws in patched.wrappers[WrappersMP.OUTER_SAMPLE].values() for w in ws]
|
||||
return wrap, outer, status
|
||||
|
||||
torch.manual_seed(1)
|
||||
bad = torch.randn(1, 4, 3, 8, 8)
|
||||
audio = torch.randn(1, 8, 5)
|
||||
x0 = NestedTensor([bad, audio])
|
||||
for i, rating in enumerate(["disliked", "disliked"]):
|
||||
monkeypatch.setattr(store, "current_prompt_id", lambda i=i: f"run-{i}")
|
||||
wrap, outer, status = load()
|
||||
for o in outer:
|
||||
o(lambda: wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(3)))
|
||||
assert store.rate(f"run-{i}", rating)["recorded"] == ["dynashift"]
|
||||
|
||||
wrap, _outer, status = load()
|
||||
assert "2 disliked, 0 liked banked" in status
|
||||
early = wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(1))
|
||||
assert early is x0 # first half: untouched
|
||||
late = wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(3))
|
||||
assert not torch.allclose(late.tensors[0], bad)
|
||||
assert late.tensors[1] is audio # sound untouched
|
||||
Reference in New Issue
Block a user