fix: H3 picture found in the packed model output; latent steering was off on every real run

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

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)

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:

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

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

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

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

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

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

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)
 

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)

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

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"]