feat: first-step seed search; ALG passes the sampler call on instead of skipping inner wrappers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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
 

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})"

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

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}

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))

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}")

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: