fix: shot memory blends a CPU-stored shot into GPU noise; learners also capture on the last step; multi-call clears the capture

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

core/input_steer.py

@@ -61,7 +61,12 @@ class Steer:
 
     def felt(self):
         """Whether this run's pushes were big enough to be felt (and so worth learning from)."""
-        return self._delivered > 0 and self._reach >= FLOOR
+        return not self._multi and self._delivered > 0 and self._reach >= FLOOR
+
+    @property
+    def multi(self):
+        """The model is called more than once per step: nothing here was steered as designed."""
+        return self._multi
 
     def effect(self):
         """How hard this run's pushes landed (largest, as a share of the input they joined),

modules/sampling/decisiveness/__init__.py

@@ -150,6 +150,15 @@ def install(patcher, values, key):
 
         patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, start)
 
+    def learn(k):
+        """Only a k whose push reached the model teaches. Checked on every call, the last
+        included: on a few-step schedule the biggest push lands on the final step."""
+        if steer.felt():
+            captured["logk"] = torch.tensor(math.log(k))
+            captured["steps"] = torch.tensor(live["steps"])
+        elif steer.multi:
+            captured.clear()
+
     def apply_model(executor, x, t, *args, **kwargs):
         named = streams.model_args(args, kwargs)
         to = named.get("transformer_options")
@@ -170,6 +179,8 @@ def install(patcher, values, key):
                      f"a step's noise, under {INERT_SHARE:.0%}); more steps or a larger k would "
                      "let it be felt" + ("" if manual else "; this run teaches nothing"))
         k = live["k"]
+        if not manual:
+            learn(k)
         if not live["acts"] or step.final or k == 1.0:
             return out
         split_x, split_out = streams.video_of(step.x, named), streams.video_of(out, named)
@@ -178,11 +189,7 @@ def install(patcher, values, key):
             return out
         video, rebuild = split_out
         sigma = float(t.max())
-        kept = step.keep(out, rebuild(rescale_x0(split_x[0], video, sigma, k).to(video.dtype)))
-        if not manual and steer.felt():
-            captured["logk"] = torch.tensor(math.log(k))     # only a k that reached the model teaches
-            captured["steps"] = torch.tensor(live["steps"])
-        return kept
+        return step.keep(out, rebuild(rescale_x0(split_x[0], video, sigma, k).to(video.dtype)))
 
     patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model)
     return ("manual k" if manual else "learned k") + ", picture only, carried into the next step"

modules/sampling/decisiveness/tests/test_decisiveness.py

@@ -164,3 +164,16 @@ def test_ratings_from_another_schedule_length_are_left_out_of_the_centre(tiny_h3
     log.new_run()
     _run(outer, lambda: wrap(lambda *a, **k: x0, x0, _sched(8)[0:1], None, None, None, _to(0), latent_shapes=shapes))
     assert any("from 0 rating(s) at 8 steps" in e["message"] for e in log.history())
+
+
+def test_on_a_four_step_schedule_the_biggest_push_lands_on_the_last_step_and_still_teaches(tiny_h3, monkeypatch):
+    import random
+    from modules.system.taste import store
+    monkeypatch.setattr(random, "gauss", lambda *a: 0.5)         # clamps to the largest k
+    wrap, outer, _ = _load(tiny_h3)
+    x0, shapes = _packed()
+    sched = torch.tensor([1.0, 0.95, 0.89, 0.80, 0.0])
+    to = lambda i: {"sample_sigmas": sched, "sigmas": sched[i:i + 1]}
+    _run(outer, lambda: [wrap(lambda *a, **k: x0, x0, sched[i:i + 1], None, None, None, to(i),
+                              latent_shapes=shapes) for i in range(4)])
+    assert store.rate("run-1", "liked")["recorded"] == ["decisiveness"]

modules/sampling/late_branch/__init__.py

@@ -145,6 +145,13 @@ def install(patcher, values, key):
             normal = executor(step.x, t, *args, **kwargs)
         finally:
             live["want"] = False
+        if not manual:
+            if steer.felt():       # every call, the last included: the biggest push lands there
+                captured["v"] = torch.tensor(live["w"])
+                captured["e"] = torch.tensor(steer.effect())
+                captured["b"] = torch.tensor(branch)   # a strength means something per branch block
+            elif steer.multi:
+                captured.clear()
         if not branching:
             if step.steering and not step.final and not cond_only:
                 _say("a negative-prompt call (CFG above 1) is left unguided: guidance is for "
@@ -164,12 +171,7 @@ def install(patcher, values, key):
             _say("off this run: could not find the picture in this model's latent")
             return normal
         video, rebuild = split_n
-        guided = step.keep(normal, rebuild(mix(video, split_w[0], live["w"]).to(video.dtype)))
-        if not manual and steer.felt():          # a push too small to feel teaches nothing
-            captured["v"] = torch.tensor(live["w"])
-            captured["e"] = torch.tensor(steer.effect())
-            captured["b"] = torch.tensor(branch)       # a strength means something per branch block
-        return guided
+        return step.keep(normal, rebuild(mix(video, split_w[0], live["w"]).to(video.dtype)))
 
     patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model)
     return (("manual" if manual else "learned") + f" strength, weak copy left out at block {branch}, "

modules/sampling/shot_memory/__init__.py

@@ -106,10 +106,11 @@ def with_coarse(video, coarse):
     return out
 
 
-def fit(coarse, h, w):
-    """A stored coarse map on this run's grid. Another resolution is resampled and
-    re-standardised: the layout survives, exact Gaussianity does not."""
-    coarse = coarse.float()
+def fit(coarse, h, w, device=None):
+    """A stored coarse map on this run's grid and device (stored maps live on the CPU; the
+    noise is on the GPU). Another resolution is resampled and re-standardised: the layout
+    survives, exact Gaussianity does not."""
+    coarse = coarse.float().to(device) if device is not None else coarse.float()
     if tuple(coarse.shape[-2:]) == (h, w):
         return coarse
     x = F.interpolate(coarse[None], size=(h, w), mode="bilinear", align_corners=False)[0]
@@ -182,7 +183,7 @@ def choose_parent(rows, cond, channels, rng=random):
 
 def blend(own, parent_coarse, amount):
     """The layout to start from: the liked shot's, `amount` of the way."""
-    coarse = amount * fit(parent_coarse, *own.shape[-2:]) + math.sqrt(1.0 - amount ** 2) * own
+    coarse = amount * fit(parent_coarse, *own.shape[-2:], device=own.device) + math.sqrt(1.0 - amount ** 2) * own
     # Unit variance only holds when the two layouts are unrelated. Rerunning the liked
     # shot's own seed makes them the same pattern and the blend stacks it (1.41x at 0.7):
     # back to a normal-strength layout, per channel. Only ever down: a few-cell grid's

modules/sampling/shot_memory/tests/test_shot_memory.py

@@ -42,6 +42,14 @@ def test_one_cell_grids_are_left_alone():
     assert coarse_of(torch.randn(1, 4, 2, 4, 4)) is None and coarse_of(torch.randn(2, 4, 2, 16, 16)) is None
 
 
+def test_a_stored_cpu_shot_blends_into_noise_on_another_device():
+    from modules.sampling.shot_memory import blend, fit
+    device = "mps" if torch.backends.mps.is_available() else "meta"
+    own = torch.randn(4, 2, 2, device=device)
+    assert fit(torch.randn(4, 2, 2).half(), 2, 2, device=own.device).device == own.device
+    assert blend(own, torch.randn(4, 2, 2).half(), 0.7).device == own.device
+
+
 def test_the_blend_never_raises_the_layout_strength_even_for_the_same_seed():
     from modules.sampling.shot_memory import blend
     c = torch.randn(4, 8, 8)

modules/sampling/stas/__init__.py

@@ -17,6 +17,9 @@ Whether H3 has massive channels at the chosen block is unknown: the run says whi
 found, or that it steered nothing. v4's user saw motion get "much more fluid and
 detailed" with this stack on; STAS is the prime suspect, not isolated.
 
+Keep this block different from late-branch guidance's branch block: at the same block the
+weak copy lacks this edit and guidance would amplify it.
+
 Switches ComfyUI's model compiler off for the run (an in-place edit of one block's output).
 """