fix(sampler): ask the value function for a gradient in its own space

Fitted on raw conditioning, it was handed H3's refined tensor and threw every
step; the run fell back to the fixed direction while reporting VF steering.
Take the gradient at the raw cond, lift it, cache it per scene.

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

CHANGELOG.md

@@ -39,6 +39,10 @@
   carried into the consumed space through the model's own text preprocessor, once per
   scene. DynaShift's prompt-similarity weighting was affected too — it compared the raw
   banked prompt against the refined one and weighted every negative equally.
+  The value function had the same fault one layer down: fitted on raw conditioning, it was
+  handed the refined tensor and threw every step, so the run steered on the fixed direction
+  while the report claimed the value function was driving. It is now asked for its gradient
+  in the space it was fitted in, once per scene rather than once per step.
 
 - **A resolution-changing second pass no longer fails on MiniMax H3.** A keyframe pin is
   packed as condition rows, so its token count belongs to the grid it was encoded on, and

samplers.py

@@ -4871,7 +4871,35 @@ class FunPackLTXAVSceneChainSampler:
         old_wrapper = model.model_options.get("model_function_wrapper")
         fixed_dir = torch.nn.functional.normalize(liked_dir.float(), dim=-1)
         resolve_dir = self._taste_direction_resolver(model, raw_cond, fixed_dir, "embed_guidance")
-        vf_noted = [False]
+        vf_state = {}
+
+        def resolve_vf(cond):
+            """Ask the value function for a gradient in the space it was FITTED in.
+
+            Its first layer is [raw_dim, 256], so handing it the refined conditioning the
+            DiT consumes is a matmul error, not a steering signal — the run then silently
+            fell back to the fixed direction while the report still claimed the value
+            function was driving. Take the gradient at the raw conditioning and carry the
+            answer across with the same bridge the fixed direction uses.
+
+            Cached: comfy builds the conds once per run, so within a scene this gradient is
+            the same every step — and identical across score_slider's +/- poles, which is
+            what makes the embed nudge cancel in (eps_+ - eps_-) as its docstring intends.
+            """
+            if "d" in vf_state:
+                return vf_state["d"]
+            vf_state["d"] = None
+            try:
+                source = raw_cond if isinstance(raw_cond, torch.Tensor) else cond
+                grad = value_fn.gradient(source)
+                unit = torch.nn.functional.normalize(grad.float(), dim=-1)
+                vf_state["d"] = self._direction_in_cond_space(model, raw_cond, unit, cond)
+                if vf_state["d"] is None:
+                    print("[FunPackSceneChain] embed_guidance: the value function gradient "
+                          "cannot be mapped onto this model's conditioning, using fixed direction")
+            except Exception as _e:
+                print(f"[FunPackSceneChain] embed_guidance: value function gradient failed ({_e}), using fixed direction")
+            return vf_state["d"]
 
         def _embed_wrapper(apply_fn, args, _ew=old_wrapper, _fixed=fixed_dir, _vf=value_fn, _s=strength):
             c = args.get("c") or {}
@@ -4886,15 +4914,9 @@ class FunPackLTXAVSceneChainSampler:
                 if scale > 0:
                     d = None
                     if _vf is not None:
-                        try:
-                            grad = _vf.gradient(cond)
-                            d = torch.nn.functional.normalize(grad.float(), dim=-1).to(cond.dtype)
-                        except Exception as _e:
-                            # Once per run: this fires on every step otherwise, and the
-                            # fixed-direction fallback below is the same each time.
-                            if not vf_noted[0]:
-                                vf_noted[0] = True
-                                print(f"[FunPackSceneChain] embed_guidance: value function gradient failed ({_e}), using fixed direction")
+                        _mapped_vf = resolve_vf(cond)
+                        if _mapped_vf is not None:
+                            d = _mapped_vf.to(cond.device, cond.dtype).expand_as(cond)
                     if d is None:
                         _mapped = resolve_dir(cond)
                         d = (_mapped.to(cond.device, cond.dtype).expand_as(cond)

tests/test_minimax_h3_sampler.py

@@ -824,3 +824,63 @@ def test_dynashift_weights_negatives_against_the_raw_prompt_not_the_refined_one(
             "c": {"c_crossattn": torch.randn(1, 12, HIDDEN)}}
     model.model_options["model_function_wrapper"](apply_fn, args)
     assert captured["ran"]
+
+
+class _RawSpaceValueFunction:
+    """A value function fitted on raw conditioning: it refuses anything else, the way a
+    [5120, 256] first layer refuses a 5376-wide tensor."""
+
+    def __init__(self):
+        self.seen = []
+
+    def gradient(self, conditioning):
+        self.seen.append(int(conditioning.shape[-1]))
+        if conditioning.shape[-1] != TEXT_DIM:
+            raise RuntimeError(
+                f"mat1 and mat2 shapes cannot be multiplied (1x{conditioning.shape[-1]} and {TEXT_DIM}x256)")
+        return torch.ones_like(conditioning)
+
+
+def _run_embed_wrapper_with_vf(model, raw_cond, cond, value_fn, steps=1):
+    node = FunPackLTXAVSceneChainSampler()
+    node._build_embed_guidance_wrapper(model, torch.randn(TEXT_DIM), 0.3,
+                                       value_fn=value_fn, raw_cond=raw_cond)
+    wrapper = model.model_options["model_function_wrapper"]
+    seen = {}
+
+    def apply_fn(x, t, **c):
+        seen["cond"] = c.get("c_crossattn")
+        return torch.zeros(1, 4)
+
+    for _ in range(steps):
+        wrapper(apply_fn, {"input": torch.zeros(1, 4), "timestep": torch.tensor([0.2]),
+                           "c": {"c_crossattn": cond}})
+    return seen
+
+
+def test_the_value_function_is_asked_for_a_gradient_in_the_space_it_was_fitted_in():
+    vf = _RawSpaceValueFunction()
+    cond = torch.randn(1, 12, HIDDEN)
+    seen = _run_embed_wrapper_with_vf(_h3ish_model(), torch.randn(1, 12, TEXT_DIM), cond, vf)
+    assert vf.seen == [TEXT_DIM]          # never handed the refined conditioning
+    assert not torch.equal(seen["cond"], cond)
+
+
+def test_the_value_function_gradient_is_taken_once_per_scene_not_once_per_step():
+    vf = _RawSpaceValueFunction()
+    _run_embed_wrapper_with_vf(_h3ish_model(), torch.randn(1, 12, TEXT_DIM),
+                               torch.randn(1, 12, HIDDEN), vf, steps=5)
+    assert len(vf.seen) == 1
+
+
+def test_an_unusable_value_function_falls_back_to_the_fixed_direction(capsys):
+    class Broken:
+        def gradient(self, conditioning):
+            raise RuntimeError("no")
+
+    cond = torch.randn(1, 12, HIDDEN)
+    seen = _run_embed_wrapper_with_vf(_h3ish_model(), torch.randn(1, 12, TEXT_DIM),
+                                      cond, Broken(), steps=3)
+    out = capsys.readouterr().out
+    assert out.count("value function gradient failed") == 1   # once per run, not per step
+    assert not torch.equal(seen["cond"], cond)                # fixed direction still steers