From 0bb7e3a01a860e85f60ed2ecb1f6dc045577fa30 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 27 Sep 2026 19:01:13 +0300 Subject: [PATCH] fix: H3 picture found in the packed model output; latent steering was off on every real run Co-Authored-By: Claude Opus 5.5 --- conftest.py | 13 ++++++ core/streams.py | 8 ++-- modules/models/minimax_h3/streams.py | 41 +++++++++++++++---- .../models/minimax_h3/tests/test_streams.py | 36 ++++++++++++++++ modules/sampling/dynashift/__init__.py | 2 +- .../dynashift/tests/test_dynashift.py | 14 +++---- modules/sampling/output_guidance/__init__.py | 2 +- .../tests/test_output_guidance.py | 17 ++++---- modules/sampling/score_slider/__init__.py | 6 +-- .../score_slider/tests/test_score_slider.py | 15 +++---- .../sampling/trajectory_guidance/__init__.py | 2 +- .../tests/test_trajectory_guidance.py | 11 ++--- 12 files changed, 124 insertions(+), 43 deletions(-) create mode 100644 modules/models/minimax_h3/tests/test_streams.py diff --git a/conftest.py b/conftest.py index 0ccd4508..dadd643b 100644 --- a/conftest.py +++ b/conftest.py @@ -145,3 +145,16 @@ class TinyH3: @pytest.fixture def tiny_h3(comfyui): return TinyH3() + + +def packed_av(video, audio): + """[video, audio] exactly as H3's model call hands them back: comfy's + `_apply_model` packs the pair flat with `pack_latents`, and the shapes that + undo it travel with the call as `latent_shapes`. -> (x, latent_shapes).""" + import comfy.utils + return comfy.utils.pack_latents([video, audio]) + + +def unpacked(x, latent_shapes): + import comfy.utils + return comfy.utils.unpack_latents(x, latent_shapes) diff --git a/core/streams.py b/core/streams.py index c55f852d..d8cafdcf 100644 --- a/core/streams.py +++ b/core/streams.py @@ -1,14 +1,16 @@ """The picture inside a sampling-time latent, whatever the model packs around it.""" -def video_of(x): +def video_of(x, named=None): """-> (video [B, C, T, H, W], rebuild(new_video) -> x), or None. A model module that packs streams answers `video_stream`; a plain 5-D - latent is its own picture. Anything else is not guessed at. + latent is its own picture. Anything else is not guessed at. `named` is the + model call's arguments (see model_args) -- a packed result is only + unpackable with the `latent_shapes` that travel with it. """ from . import registry - split = registry.current().ask("video_stream", x) + split = registry.current().ask("video_stream", x, named or {}) if split is not None: return split if getattr(x, "dim", None) and not getattr(x, "is_nested", False) and x.dim() == 5: diff --git a/modules/models/minimax_h3/streams.py b/modules/models/minimax_h3/streams.py index 2a918337..fd93a28a 100644 --- a/modules/models/minimax_h3/streams.py +++ b/modules/models/minimax_h3/streams.py @@ -1,11 +1,38 @@ -"""Where H3's video lives in a sampling-time latent: the first of [video, audio].""" +"""Where H3's video lives in a sampling-time latent. + +Two shapes reach a feature, and both are real: + +* the model call's result (APPLY_MODEL) is PACKED: comfy's `_apply_model` + flattens [video, audio] into one [B, 1, N] tensor with `pack_latents`, and the + shapes to undo that ride along as `latent_shapes`; +* elsewhere (callbacks, the finished latent) it is a two-part NestedTensor. +""" + +import math -def video_stream(x): +def video_stream(x, named=None): """-> (video [B, C, T, H, W], rebuild(new_video) -> x), or None when `x` is - not H3's two-branch latent.""" - if not getattr(x, "is_nested", False) or len(x.tensors) != 2 or x.tensors[0].dim() != 5: + neither of H3's two shapes.""" + if getattr(x, "is_nested", False): + if len(x.tensors) != 2 or x.tensors[0].dim() != 5: + return None + from comfy.nested_tensor import NestedTensor + rest = x.tensors[1:] + return x.tensors[0], lambda video: NestedTensor([video, *rest]) + shapes = (named or {}).get("latent_shapes") + if not hasattr(x, "dim") or x.dim() != 3 or not shapes or len(shapes) != 2 \ + or len(shapes[0]) != 5: return None - from comfy.nested_tensor import NestedTensor - rest = x.tensors[1:] - return x.tensors[0], lambda video: NestedTensor([video, *rest]) + n = math.prod(shapes[0][1:]) + if x.shape[-1] != n + math.prod(shapes[1][1:]): + return None + b = x.shape[0] + video = x[..., :n].reshape(b, *shapes[0][1:]) + rest = x[..., n:] + + def rebuild(new_video): + import torch + return torch.cat([new_video.reshape(b, 1, n).to(rest.dtype), rest], dim=-1) + + return video, rebuild diff --git a/modules/models/minimax_h3/tests/test_streams.py b/modules/models/minimax_h3/tests/test_streams.py new file mode 100644 index 00000000..db5bbea9 --- /dev/null +++ b/modules/models/minimax_h3/tests/test_streams.py @@ -0,0 +1,36 @@ +"""H3's picture, found in both shapes it reaches features in.""" + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +def test_packed_model_output_splits_and_rebuilds_exactly(): + from conftest import packed_av, unpacked + from modules.models.minimax_h3.streams import video_stream + video, audio = torch.randn(2, 4, 3, 6, 6), torch.randn(2, 8, 5) + x, shapes = packed_av(video, audio) + got, rebuild = video_stream(x, {"latent_shapes": shapes}) + assert torch.equal(got, video) + out = rebuild(got * 2) + assert torch.equal(unpacked(out, shapes)[0], video * 2) + assert torch.equal(unpacked(out, shapes)[1], audio) + + +def test_packed_without_its_shapes_is_not_guessed(): + from conftest import packed_av + from modules.models.minimax_h3.streams import video_stream + x, _ = packed_av(torch.randn(1, 4, 3, 6, 6), torch.randn(1, 8, 5)) + assert video_stream(x, {}) is None + + +def test_nested_still_works(): + from comfy.nested_tensor import NestedTensor + from modules.models.minimax_h3.streams import video_stream + video, audio = torch.randn(1, 4, 3, 6, 6), torch.randn(1, 8, 5) + got, rebuild = video_stream(NestedTensor([video, audio])) + assert got is video and rebuild(video).tensors[1] is audio diff --git a/modules/sampling/dynashift/__init__.py b/modules/sampling/dynashift/__init__.py index acd087b3..eba0315e 100644 --- a/modules/sampling/dynashift/__init__.py +++ b/modules/sampling/dynashift/__init__.py @@ -188,7 +188,7 @@ def install(patcher, values, key): out = executor(x, t, *args, **kwargs) named = streams.model_args(args, kwargs) to = named.get("transformer_options") - split = streams.video_of(out) + split = streams.video_of(out, named) if split is None: _say("off this run: could not find the picture in this model's latent") return out diff --git a/modules/sampling/dynashift/tests/test_dynashift.py b/modules/sampling/dynashift/tests/test_dynashift.py index 7871c6e5..96a41acb 100644 --- a/modules/sampling/dynashift/tests/test_dynashift.py +++ b/modules/sampling/dynashift/tests/test_dynashift.py @@ -51,7 +51,6 @@ def test_nothing_banked_at_this_resolution_changes_nothing(): def test_it_banks_the_last_step_and_steers_once_rated(tiny_h3, monkeypatch): - from comfy.nested_tensor import NestedTensor from comfy.patcher_extension import WrappersMP from modules.sampling.modifiers.nodes import FunPackLoadModifiers from modules.system.taste import store @@ -66,12 +65,13 @@ def test_it_banks_the_last_step_and_steers_once_rated(tiny_h3, monkeypatch): torch.manual_seed(1) bad = torch.randn(1, 4, 3, 8, 8) audio = torch.randn(1, 8, 5) - x0 = NestedTensor([bad, audio]) + from conftest import packed_av, unpacked + x0, shapes = packed_av(bad, audio) for i, rating in enumerate(["disliked", "disliked"]): monkeypatch.setattr(store, "current_prompt_id", lambda i=i: f"run-{i}") wrap, outer, status = load() for o in outer: - o(lambda: wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3))) + o(lambda: wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3), latent_shapes=shapes)) assert store.rate(f"run-{i}", rating)["recorded"] == ["dynashift"] wrap, outer, _status = load() @@ -80,8 +80,8 @@ def test_it_banks_the_last_step_and_steers_once_rated(tiny_h3, monkeypatch): for o in outer: o(lambda: None) # a run starts: the bank is read assert any("2 disliked, 0 liked banked" in e["message"] for e in log.history()) - early = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(1)) + early = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(1), latent_shapes=shapes) assert early is x0 # first half: untouched - late = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3)) - assert not torch.allclose(late.tensors[0], bad) - assert late.tensors[1] is audio # sound untouched + late = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3), latent_shapes=shapes) + assert not torch.allclose(unpacked(late, shapes)[0], bad) + assert torch.equal(unpacked(late, shapes)[1], audio) # sound untouched diff --git a/modules/sampling/output_guidance/__init__.py b/modules/sampling/output_guidance/__init__.py index 2f7674dc..93ba8078 100644 --- a/modules/sampling/output_guidance/__init__.py +++ b/modules/sampling/output_guidance/__init__.py @@ -74,7 +74,7 @@ def install(patcher, values, key): out = executor(x, t, *args, **kwargs) named = streams.model_args(args, kwargs) to = named.get("transformer_options") - split = streams.video_of(out) + split = streams.video_of(out, named) if split is None: _say("off this run: could not find the picture in this model's latent") return out diff --git a/modules/sampling/output_guidance/tests/test_output_guidance.py b/modules/sampling/output_guidance/tests/test_output_guidance.py index 4e498411..7b178578 100644 --- a/modules/sampling/output_guidance/tests/test_output_guidance.py +++ b/modules/sampling/output_guidance/tests/test_output_guidance.py @@ -42,9 +42,9 @@ def _teach(kind, names): def _x0(): - from comfy.nested_tensor import NestedTensor + from conftest import packed_av torch.manual_seed(3) - return NestedTensor([torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)]) + return packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)) def _start(outer): @@ -55,16 +55,17 @@ def _start(outer): def test_late_guidance_banks_the_final_picture_and_steers_only_late(tiny_h3, monkeypatch): from modules.system.taste import store monkeypatch.setattr(store, "current_prompt_id", lambda: "run-1") - x0 = _x0() + x0, shapes = _x0() wrap, outer = _load(tiny_h3, "output_guidance") for o in outer: - o(lambda: wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3))) + o(lambda: wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3), latent_shapes=shapes)) assert store.rate("run-1", "liked")["recorded"] == ["x0_final"] _teach("x0_final", ["final"]) wrap, outer = _load(tiny_h3, "output_guidance") _start(outer) - assert wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(1)) is x0 - late = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3)) - assert not torch.allclose(late.tensors[0], x0.tensors[0]) - assert late.tensors[1] is x0.tensors[1] + from conftest import unpacked + assert wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(1), latent_shapes=shapes) is x0 + late = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3), latent_shapes=shapes) + assert not torch.allclose(unpacked(late, shapes)[0], unpacked(x0, shapes)[0]) + assert torch.equal(unpacked(late, shapes)[1], unpacked(x0, shapes)[1]) diff --git a/modules/sampling/score_slider/__init__.py b/modules/sampling/score_slider/__init__.py index b4fe46e0..c672a074 100644 --- a/modules/sampling/score_slider/__init__.py +++ b/modules/sampling/score_slider/__init__.py @@ -120,12 +120,12 @@ def install(patcher, values, key): return base size = torch.linalg.vector_norm(c[:, words], dim=-1, dtype=torch.float32).mean() step = (d.to(c.device) * STEP * size).to(c.dtype) * words.view(1, -1, 1) - split = streams.video_of(base) + split = streams.video_of(base, named) if split is None: _say("off this run: could not find the picture in this model's latent") return base - plus = streams.video_of(executor(x, t, **{**named, "c_crossattn": c + step})) - minus = streams.video_of(executor(x, t, **{**named, "c_crossattn": c - step})) + plus = streams.video_of(executor(x, t, **{**named, "c_crossattn": c + step}), named) + minus = streams.video_of(executor(x, t, **{**named, "c_crossattn": c - step}), named) video, rebuild = split return rebuild(video + (plus[0] - minus[0]) * amount) diff --git a/modules/sampling/score_slider/tests/test_score_slider.py b/modules/sampling/score_slider/tests/test_score_slider.py index 13efdd1b..44db7ecf 100644 --- a/modules/sampling/score_slider/tests/test_score_slider.py +++ b/modules/sampling/score_slider/tests/test_score_slider.py @@ -44,7 +44,7 @@ def test_similar_prompts_pick_their_own_direction(): def test_late_steps_combine_three_passes_on_the_picture_words_only(tiny_h3): - from comfy.nested_tensor import NestedTensor + from conftest import packed_av, unpacked from comfy.patcher_extension import WrappersMP from modules.sampling.modifiers.nodes import FunPackLoadModifiers from modules.system.taste import store @@ -68,19 +68,20 @@ def test_late_steps_combine_three_passes_on_the_picture_words_only(tiny_h3): def executor(x, t, c_concat=None, c_crossattn=None, control=None, transformer_options=None, **kw): seen.append(c_crossattn.clone()) - return NestedTensor([torch.full((1, 2, 1, 2, 2), float(c_crossattn[0, :, 0].sum())), audio]) + return packed_av(torch.full((1, 2, 1, 2, 2), float(c_crossattn[0, :, 0].sum())), audio)[0] c = torch.ones(1, 4, dim) tags = torch.tensor([1, 1, 0, 1]) # row 2: a reference image's token x0 = executor(None, None, c_crossattn=c) + shapes = packed_av(torch.zeros(1, 2, 1, 2, 2), audio)[1] seen.clear() - early = wrap(executor, x0, None, None, c, None, _to(0), minimax_payload={"text_token_tags": tags}) - assert len(seen) == 1 and torch.equal(early.tensors[0], x0.tensors[0]) + early = wrap(executor, x0, None, None, c, None, _to(0), minimax_payload={"text_token_tags": tags}, latent_shapes=shapes) + assert len(seen) == 1 and torch.equal(early, x0) seen.clear() - late = wrap(executor, x0, None, None, c, None, _to(3), minimax_payload={"text_token_tags": tags}) + late = wrap(executor, x0, None, None, c, None, _to(3), minimax_payload={"text_token_tags": tags}, latent_shapes=shapes) assert len(seen) == 3 plus = seen[1] assert torch.equal(plus[0, 2], c[0, 2]) and not torch.equal(plus[0, 0], c[0, 0]) - assert not torch.equal(late.tensors[0], x0.tensors[0]) - assert late.tensors[1] is audio + assert not torch.equal(unpacked(late, shapes)[0], unpacked(x0, shapes)[0]) + assert torch.equal(unpacked(late, shapes)[1], audio) diff --git a/modules/sampling/trajectory_guidance/__init__.py b/modules/sampling/trajectory_guidance/__init__.py index 6f1b97df..31d74bbf 100644 --- a/modules/sampling/trajectory_guidance/__init__.py +++ b/modules/sampling/trajectory_guidance/__init__.py @@ -80,7 +80,7 @@ def install(patcher, values, key): out = executor(x, t, *args, **kwargs) named = streams.model_args(args, kwargs) q = quarter(named.get("transformer_options")) - split = streams.video_of(out) + split = streams.video_of(out, named) if q is None or split is None: _say("off this run: could not place this step on the schedule or find the picture") return out diff --git a/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py b/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py index ab960ede..d9869649 100644 --- a/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py +++ b/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py @@ -42,9 +42,9 @@ def _teach(kind, names): def _x0(): - from comfy.nested_tensor import NestedTensor + from conftest import packed_av torch.manual_seed(3) - return NestedTensor([torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)]) + return packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)) def _start(outer): @@ -56,17 +56,18 @@ def test_early_guidance_steers_in_the_first_quarter_and_banks_every_quarter(tiny from modules.system.taste import store _teach("x0_quarters", ["q0", "q1", "q2", "q3"]) monkeypatch.setattr(store, "current_prompt_id", lambda: "run-2") - x0 = _x0() + x0, shapes = _x0() wrap, outer = _load(tiny_h3, "trajectory_guidance") def run(): - return [wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(i)) for i in range(4)] + return [wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(i), latent_shapes=shapes) for i in range(4)] results = {} for o in outer: o(lambda: results.setdefault("steps", run())) first = results["steps"][0] - assert not torch.allclose(first.tensors[0], x0.tensors[0]) + assert not torch.allclose(first[..., :4 * 3 * 8 * 8], x0[..., :4 * 3 * 8 * 8]) + assert torch.equal(first[..., 4 * 3 * 8 * 8:], x0[..., 4 * 3 * 8 * 8:]) pending = torch.load(store.ROOT / "fox" / "x0_quarters.pending.pt")["rows"] assert sorted(pending) == ["q0", "q1", "q2", "q3"]