From 3c187bda8ae7a378bc3b75a7656f6347eb196d43 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 27 Sep 2026 19:08:02 +0300 Subject: [PATCH] feat: first-step seed search; ALG passes the sampler call on instead of skipping inner wrappers Co-Authored-By: Claude Opus 5.5 --- core/dit_hooks.py | 15 +- modules/sampling/alg/__init__.py | 17 +-- .../tests/test_alg_through_a_real_sampler.py | 39 ++++- .../sampling/explore_first_step/__init__.py | 134 ++++++++++++++++++ .../tests/test_explore_first_step.py | 98 +++++++++++++ modules/sampling/score_slider/__init__.py | 3 +- .../sampling/trajectory_guidance/__init__.py | 2 + 7 files changed, 295 insertions(+), 13 deletions(-) create mode 100644 modules/sampling/explore_first_step/__init__.py create mode 100644 modules/sampling/explore_first_step/tests/test_explore_first_step.py diff --git a/core/dit_hooks.py b/core/dit_hooks.py index 430a0b51..dc0ebdfd 100644 --- a/core/dit_hooks.py +++ b/core/dit_hooks.py @@ -167,9 +167,22 @@ def current_step(transformer_options): return int(hit[0]), int(sched.shape[0]) - 1 +PROBE = "funpack_probe" + + +def probing(transformer_options) -> bool: + """A throwaway model call (a candidate scored and discarded, e.g. the + first-step seed search). Nothing may be learned from one: a feature that + captures from it banks a step the clip never took.""" + return bool((transformer_options or {}).get(PROBE)) + + def last_step(transformer_options) -> bool: """Whether this denoise call is the schedule's final step -- True when that - cannot be told, so a capture taken "on the last step" still happens.""" + cannot be told, so a capture taken "on the last step" still happens. Never + during a probe.""" + if probing(transformer_options): + return False where = current_step(transformer_options) return where is None or where[0] >= where[1] - 1 diff --git a/modules/sampling/alg/__init__.py b/modules/sampling/alg/__init__.py index e5583845..d1b4a6d3 100644 --- a/modules/sampling/alg/__init__.py +++ b/modules/sampling/alg/__init__.py @@ -150,14 +150,15 @@ def install(patcher, values, key): # inherit a blurred one. model_k.latent_image = latent_image - import comfy.samplers - replacement = comfy.samplers.KSAMPLER( - alg_sampler_function, - extra_options=getattr(sampler, "extra_options", {}), - inpaint_options=getattr(sampler, "inpaint_options", {}), - ) - return replacement.sample(model_wrap, sigmas, extra_args, callback, noise, - latent_image, denoise_mask, disable_pbar) + # Swapped on the sampler for this call and passed ON, not run by a + # replacement sampler: a replacement would skip every SAMPLER_SAMPLE + # wrapper inside this one (the seed search, anything later). + sampler.sampler_function = alg_sampler_function + try: + return executor(model_wrap, sigmas, extra_args, callback, noise, + latent_image, denoise_mask, disable_pbar) + finally: + sampler.sampler_function = inner_fn patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, wrapper) return f"loosening the starting image (amount {strength}, until sigma {threshold})" diff --git a/modules/sampling/alg/tests/test_alg_through_a_real_sampler.py b/modules/sampling/alg/tests/test_alg_through_a_real_sampler.py index 2d9569ea..47dbc846 100644 --- a/modules/sampling/alg/tests/test_alg_through_a_real_sampler.py +++ b/modules/sampling/alg/tests/test_alg_through_a_real_sampler.py @@ -57,6 +57,10 @@ class _Executor: self.fell_back = False def __call__(self, *args, **kwargs): + if self._fallback == "sample": + # The end of a real chain: the sampler's own sample(), with whatever + # sampler_function the wrappers above left on it. + return self.class_obj.sample(*args, **kwargs) self.fell_back = True return self._fallback @@ -111,7 +115,7 @@ def test_the_anchor_the_denoiser_sees_changes_with_sigma(): wrapper, _note, _kind, _mp = _install_alg() real = comfy.samplers.KSAMPLER(sampler_function) - executor = _Executor(real, fallback="unused") + executor = _Executor(real, fallback="sample") wrapper(executor, _Guider(), torch.tensor([0.9, 0.2]), {}, None, torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True) @@ -138,7 +142,7 @@ def test_the_real_denoiser_is_left_holding_the_sharp_anchor(): return x wrapper, _n, _k, _mp = _install_alg() - wrapper(_Executor(comfy.samplers.KSAMPLER(sampler_function), "unused"), + wrapper(_Executor(comfy.samplers.KSAMPLER(sampler_function), "sample"), _Guider(), torch.tensor([0.9]), {}, None, torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True) @@ -160,7 +164,7 @@ def test_a_sampler_error_still_restores_the_sharp_anchor(): wrapper, _n, _k, _mp = _install_alg() with pytest.raises(RuntimeError, match="interrupted"): - wrapper(_Executor(comfy.samplers.KSAMPLER(sampler_function), "unused"), + wrapper(_Executor(comfy.samplers.KSAMPLER(sampler_function), "sample"), _Guider(), torch.tensor([0.9]), {}, None, torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True) @@ -202,3 +206,32 @@ def test_an_image_latent_stands_aside_rather_than_reshaping_a_guess(): torch.ones_like(image_latent), True) assert out == "untouched" and executor.fell_back + + +def test_alg_passes_the_call_on_and_leaves_the_sampler_as_it_found_it(): + """A replacement sampler would skip every SAMPLER_SAMPLE wrapper beneath ALG + (the seed search). And the sampler object is shared across runs: the swap + must not outlive this call.""" + import torch + import comfy.samplers + + sharp = torch.randn(1, 3, 4, 16, 16) + ran = [] + + def sampler_function(model_k, x, sigmas, **kwargs): + ran.append(True) + return x + + real = comfy.samplers.KSAMPLER(sampler_function) + below = [] + + class _Beneath(_Executor): + def __call__(self, *args, **kwargs): + below.append(True) # a wrapper installed after ALG + return self.class_obj.sample(*args, **kwargs) + + wrapper, _n, _k, _mp = _install_alg() + wrapper(_Beneath(real, None), _Guider(), torch.tensor([0.9]), {}, None, + torch.zeros_like(sharp), sharp, torch.ones_like(sharp), True) + assert below and ran + assert real.sampler_function is sampler_function diff --git a/modules/sampling/explore_first_step/__init__.py b/modules/sampling/explore_first_step/__init__.py new file mode 100644 index 00000000..d5968c7b --- /dev/null +++ b/modules/sampling/explore_first_step/__init__.py @@ -0,0 +1,134 @@ +"""First-step seed search: try a few seeds for one step, keep the one your taste likes. + +Before the real run, each candidate seed gets exactly ONE denoising step, its +predicted picture is scored by a judge trained on how your rated clips looked +at their first step, and the run continues with the best seed. The first step +already carries most of what a rating reacts to (v4's probe: ~88% of the late +half's signal). + +User verdict (v4): a win on simple prompts; no help on prompts that depend on +word trickery, which the first step can't read yet. + +Costs one extra model pass per candidate. Candidate steps are marked as probes, +so no learning feature records a step the clip never took. +""" + +import torch +from comfy.patcher_extension import WrappersMP + +from ..._core import dit_hooks, log, registry, streams + +ID = "explore_first_step" +TITLE = "First-step seed search" +MOUNT = "generation.sampling" +STAGE = "sampling" +CATEGORY = "sampling" +STATUS = "proven" +REQUIRES = ["temporal_latent"] + +SETTINGS = { + "enabled": { + "type": "bool", "default": False, + "label": "Seed search", + "hint": "After 10+ ratings, tries a few seeds for one step each and keeps the one your taste likes.", + }, + "candidates": { + "type": "int", "default": 3, "min": 2, "max": 8, + "label": "Seeds to try", + "hint": "Each costs one extra step. Your seed is always one of them.", + "when": {"enabled": True}, + }, +} + +KIND = "first_step" +NAME = "first" + + +def _say(message): + log.once(f"{ID}:{message}", log.ALERT, "FunPack Seed search", message) + + +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 + count = int(values.get("candidates", 3)) + live = {} + + def fresh(): + live["judge"] = taste.judge(KIND, NAME) + liked, disliked = taste.counts(KIND, NAME) + if live["judge"] is None: + log.once(f"{ID}:state", log.INFO, "FunPack Seed search", + f"key {taste.key!r}: learning ({liked} liked / {disliked} disliked; " + "searches from 10 rated clips with both kinds)") + + captured = taste.collect(patcher, key, KIND, fresh=fresh) + + # What the real run's first step looked like: what the judge learns from. + def apply_model(executor, x, t, *args, **kwargs): + out = executor(x, t, *args, **kwargs) + named = streams.model_args(args, kwargs) + to = named.get("transformer_options") + where = dit_hooks.current_step(to) + if dit_hooks.probing(to) or where is None or where[0] != 0: + return out + split = streams.video_of(out, named) + if split is not None and split[0].shape[0] == 1: + captured[NAME] = taste.describe(split[0].detach()) + return out + + patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model) + + def sampler_sample(executor, model_wrap, sigmas, extra_args, callback, noise, + latent_image=None, denoise_mask=None, disable_pbar=False): + run = lambda n: executor(model_wrap, sigmas, extra_args, callback, n, + latent_image, denoise_mask, disable_pbar) + judge = live.get("judge") + sampler = getattr(executor, "class_obj", None) + if judge is None or noise.shape[0] != 1 or not hasattr(sampler, "max_denoise"): + return run(noise) + best = pick(judge, model_wrap, sigmas, extra_args, noise, latent_image, denoise_mask, + sampler.max_denoise(model_wrap, sigmas)) + return run(best) + + def pick(judge, model_wrap, sigmas, extra_args, noise, latent_image, denoise_mask, max_denoise): + from comfy.samplers import KSamplerX0Inpaint + options = dict(extra_args.get("model_options") or {}) + options["transformer_options"] = {**options.get("transformer_options", {}), + dit_hooks.PROBE: True} + named = {"latent_shapes": getattr(model_wrap.inner_model, "latent_shapes", None)} + seed = int(extra_args.get("seed") or 0) + sigma = sigmas[0] * noise.new_ones([noise.shape[0]]) + scored = [] + for k in range(count): + if k == 0: + candidate = noise # the seed you asked for + else: + gen = torch.Generator().manual_seed(seed + 7919 * k) + candidate = torch.randn(noise.shape, generator=gen).to(noise) + model_k = KSamplerX0Inpaint(model_wrap, sigmas) + model_k.latent_image, model_k.noise = latent_image, candidate + x = model_wrap.inner_model.model_sampling.noise_scaling( + sigmas[0], candidate, latent_image, max_denoise) + x0 = model_k(x, sigma, denoise_mask=denoise_mask, model_options=options, + seed=extra_args.get("seed")) + split = streams.video_of(x0, named) + if split is None: + _say("off this run: could not find the picture to score") + return noise + scored.append((judge.score(split[0]), k, candidate)) + score, k, best = max(scored, key=lambda s: s[0]) + log.info("FunPack Seed search", + ("kept your seed" if k == 0 else f"switched to candidate {k + 1}") + + f" of {count} (scores {', '.join(f'{s:.3f}' for s, _k, _c in scored)})") + return best + + patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, sampler_sample) + return f"{count} seeds, one step each; learns every run, searches from 10 rated clips" + + +PROVIDES = {"modifier": install} diff --git a/modules/sampling/explore_first_step/tests/test_explore_first_step.py b/modules/sampling/explore_first_step/tests/test_explore_first_step.py new file mode 100644 index 00000000..165b490d --- /dev/null +++ b/modules/sampling/explore_first_step/tests/test_explore_first_step.py @@ -0,0 +1,98 @@ +"""Seed search: scores one real first step per candidate, keeps the best, learns nothing from probes.""" + +import types + +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") + + +SHAPES = [torch.Size([1, 4, 2, 4, 4]), torch.Size([1, 3, 5])] +N_VIDEO = 4 * 2 * 4 * 4 + + +class Model: + """What SAMPLER_SAMPLE is handed: x0 = x, packed like H3's.""" + + def __init__(self): + self.options = [] + self.inner_model = types.SimpleNamespace( + latent_shapes=SHAPES, + model_sampling=types.SimpleNamespace(sigma_max=1.0, noise_scaling=lambda s, n, l, m: n)) + + def __call__(self, x, sigma, model_options=None, seed=None): + self.options.append(model_options) + return x + + +class Executor: + def __init__(self): + import comfy.samplers + self.class_obj = comfy.samplers.KSAMPLER(lambda *a, **k: None) + self.noise = None + + def __call__(self, model_wrap, sigmas, extra_args, callback, noise, *rest): + self.noise = noise + return noise + + +def _load(tiny): + from comfy.patcher_extension import WrappersMP + from modules.sampling.modifiers.nodes import FunPackLoadModifiers + patched, _ = FunPackLoadModifiers.execute(tiny.patcher, { + "taste": {"key": "fox"}, "explore_first_step": {"enabled": True, "candidates": 4}}).result + sample = [w for ws in patched.wrappers[WrappersMP.SAMPLER_SAMPLE].values() for w in ws][-1] + for ws in patched.wrappers[WrappersMP.OUTER_SAMPLE].values(): + for o in ws: + o(lambda: None) + return sample + + +def _teach(): + """Liked first steps were bright, disliked dark.""" + from modules.system.taste import store, value + for i in range(12): + sign = 1.0 if i % 2 else -1.0 + store.capture("fox", "first_step", + {"first": value.describe(sign + 0.2 * torch.randn(1, 4, 2, 4, 4))}, + prompt_id=f"t{i}") + store.rate(f"t{i}", "liked" if sign > 0 else "disliked") + + +def test_until_it_has_learned_the_seed_is_untouched_and_nothing_extra_runs(tiny_h3): + sample = _load(tiny_h3) + model, ex = Model(), Executor() + noise = torch.randn(1, 1, N_VIDEO + 15) + sample(ex, model, torch.tensor([1.0, 0.0]), {"model_options": {}, "seed": 5}, None, noise, + torch.zeros_like(noise), None, True) + assert ex.noise is noise and not model.options + + +def test_it_keeps_the_candidate_the_judge_scores_best_and_marks_every_try_a_probe(tiny_h3): + from core import dit_hooks + from modules.system.taste import value + _teach() + sample = _load(tiny_h3) + model, ex = Model(), Executor() + noise = torch.randn(1, 1, N_VIDEO + 15) + sample(ex, model, torch.tensor([1.0, 0.0]), {"model_options": {}, "seed": 5}, None, noise, + torch.zeros_like(noise), None, True) + assert len(model.options) == 4 + assert all(dit_hooks.probing(o["transformer_options"]) for o in model.options) + + from modules.system.taste import Handle + judge = Handle("fox").judge("first_step", "first") + score = lambda n: judge.score(n[..., :N_VIDEO].reshape(SHAPES[0])) + candidates = [noise] + [torch.randn(noise.shape, generator=torch.Generator().manual_seed(5 + 7919 * k)) + for k in range(1, 4)] + assert torch.equal(ex.noise, max(candidates, key=score)) diff --git a/modules/sampling/score_slider/__init__.py b/modules/sampling/score_slider/__init__.py index c672a074..c825d9b9 100644 --- a/modules/sampling/score_slider/__init__.py +++ b/modules/sampling/score_slider/__init__.py @@ -109,7 +109,8 @@ def install(patcher, values, key): words = torch.ones(c.shape[1], dtype=torch.bool, device=c.device) words = words.to(c.device) pooled = c[0][words].float().mean(0).detach() - captured[NAME] = pooled + if not dit_hooks.probing(named.get("transformer_options")): + captured[NAME] = pooled if "dir" not in live: live["dir"], how = direction(taste.rows(KIND), pooled, similar) log.once(f"{ID}:state", log.INFO, "FunPack Taste slider", f"key {taste.key!r}: {how}") diff --git a/modules/sampling/trajectory_guidance/__init__.py b/modules/sampling/trajectory_guidance/__init__.py index 31d74bbf..1a2f1bc3 100644 --- a/modules/sampling/trajectory_guidance/__init__.py +++ b/modules/sampling/trajectory_guidance/__init__.py @@ -79,6 +79,8 @@ def install(patcher, values, key): def apply_model(executor, x, t, *args, **kwargs): out = executor(x, t, *args, **kwargs) named = streams.model_args(args, kwargs) + if dit_hooks.probing(named.get("transformer_options")): + return out # a discarded candidate: learn and steer nothing q = quarter(named.get("transformer_options")) split = streams.video_of(out, named) if q is None or split is None: