From 3e6cb7fbb399ee06aab7f0b2b321c4bff2b6bc1a Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:40:32 +0300 Subject: [PATCH] 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 --- CHANGELOG.md | 4 +++ samplers.py | 42 ++++++++++++++++------ tests/test_minimax_h3_sampler.py | 60 ++++++++++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ccb796f8..39327435 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/samplers.py b/samplers.py index 88536590..8c6494da 100644 --- a/samplers.py +++ b/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) diff --git a/tests/test_minimax_h3_sampler.py b/tests/test_minimax_h3_sampler.py index a6c90a48..ccaf2621 100644 --- a/tests/test_minimax_h3_sampler.py +++ b/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