From b831372879b15b2375db45bef1f5fb5dde1d65ef Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 27 Sep 2026 18:34:16 +0300 Subject: [PATCH] feat: DynaShift steers late steps away from disliked clips, banked on the taste key Co-Authored-By: Claude Opus 5.5 --- modules/models/minimax_h3/__init__.py | 4 +- modules/models/minimax_h3/streams.py | 11 + modules/sampling/dynashift/__init__.py | 225 ++++++++++++++++++ .../dynashift/tests/test_dynashift.py | 83 +++++++ 4 files changed, 322 insertions(+), 1 deletion(-) create mode 100644 modules/models/minimax_h3/streams.py create mode 100644 modules/sampling/dynashift/__init__.py create mode 100644 modules/sampling/dynashift/tests/test_dynashift.py diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index 14c813c7..619312f2 100644 --- a/modules/models/minimax_h3/__init__.py +++ b/modules/models/minimax_h3/__init__.py @@ -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} diff --git a/modules/models/minimax_h3/streams.py b/modules/models/minimax_h3/streams.py new file mode 100644 index 00000000..2a918337 --- /dev/null +++ b/modules/models/minimax_h3/streams.py @@ -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]) diff --git a/modules/sampling/dynashift/__init__.py b/modules/sampling/dynashift/__init__.py new file mode 100644 index 00000000..b609e6e9 --- /dev/null +++ b/modules/sampling/dynashift/__init__.py @@ -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} diff --git a/modules/sampling/dynashift/tests/test_dynashift.py b/modules/sampling/dynashift/tests/test_dynashift.py new file mode 100644 index 00000000..8567fafe --- /dev/null +++ b/modules/sampling/dynashift/tests/test_dynashift.py @@ -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