diff --git a/minimax_h3.py b/minimax_h3.py index 977f4183..01a34038 100644 --- a/minimax_h3.py +++ b/minimax_h3.py @@ -506,15 +506,152 @@ def checkpoint_mode_note(has_keyframes: bool, has_refs: bool): return None -def keyframe_indices_supported(frame_index, frame_count): - """H3's packed layout only places a pin at the FIRST or LAST pixel frame. +# ── interior keyframe pins ──────────────────────────────────────────────────── +# +# A pin is a block of condition rows carrying ONE time coordinate (``cond_t``) on the +# same spatial grid as the target video, never denoised (``img_update=False``). Upstream +# ``PackedLayout`` computes that coordinate in two hard-coded branches and raises for +# anything else: +# +# pixel_index == 0 -> cond_t = text_len +# pixel_index == frame_count - 1 -> cond_t = text_len + sum(_video_t_spans(latent_t)) +# - FRAME_RESCALE +# +# Those are not two special cases. They are the two ENDPOINTS of one straight line: the +# video time axis advances exactly ``FRAME_RESCALE`` per PIXEL frame, so +# +# cond_t = text_len + FRAME_RESCALE * pixel_index +# +# reproduces both, exactly, at every clip length — including the irregular +# ``FRAME_PER_TOKEN = (1, 4, 4, 4, 4)`` latent grouping, because the grouping changes how +# many pixel frames a LATENT frame covers without changing the per-pixel-frame rate. +# :func:`_linear_rule_matches_upstream` re-derives that equality from upstream's own +# constants before anything is patched, so an upstream change to the grid turns this +# feature off instead of silently mis-placing pins. +# +# What this does NOT establish is that the WEIGHTS were trained on interior pins. fl2va +# saw condition rows at t=0 and t=end and nowhere between. MM-RoPE is continuous in t, so +# an interior coordinate is representable rather than undefined, and the pin reaches the +# target through ordinary self-attention with no special-cased path — but "representable" +# is not "learned". Treat an interior pin as experimental until a run shows the frame +# actually landing. +KEYFRAME_T_PER_FRAME = 5.0 / 3.0 # upstream FRAME_RESCALE, re-checked before patching - ``PackedLayout`` raises for anything else, so callers must check rather than - discover it mid-sample. Mid-clip guidance has to go through a reference block - (identity/style, not frame-pinned) instead. +_INTERIOR_PINS = {"state": None} # None = not tried, True/False = patch verdict + + +def keyframe_cond_t(text_len, pixel_index): + """The packed-sequence time coordinate of a pin at ``pixel_index``.""" + return float(text_len) + KEYFRAME_T_PER_FRAME * float(int(pixel_index)) + + +def _linear_rule_matches_upstream(): + """True when :func:`keyframe_cond_t` reproduces upstream's own two branches. + + Checked against ``_video_t_spans`` and ``FRAME_RESCALE`` themselves rather than + against remembered numbers, so a changed frame grid is caught here instead of + becoming pins placed at the wrong instant. + """ + from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE, _video_t_spans + + if abs(float(FRAME_RESCALE) - KEYFRAME_T_PER_FRAME) > 1e-9: + return False + text_len = 7 # any value; it cancels + for latent_t in (1, 2, 5, 6, 11, 24, 25, 26): + spans = _video_t_spans(latent_t) + frame_count = sum(FRAME_PER_TOKEN[k % 5] for k in range(latent_t)) + first = float(text_len) + last = float(text_len) + sum(spans) - FRAME_RESCALE + if abs(first - keyframe_cond_t(text_len, 0)) > 1e-9: + return False + if abs(last - keyframe_cond_t(text_len, frame_count - 1)) > 1e-9: + return False + return True + + +def install_interior_keyframes(): + """Let a keyframe pin sit at any pixel frame, not only the first or the last. + + Deliberately NOT a re-implementation of ``PackedLayout.__init__``. Upstream builds the + layout exactly as it always has — with every interior pin declared at index 0, which is + always legal — and this rewrites the ``t`` column of the resulting condition rows to the + coordinate the pin actually wants. Everything else about the layout (row order, latent + placement, ``img_update``, the segment table, refs, audio) is upstream's, so upstream + changes are inherited rather than frozen into a copy here. + + A run whose pins are all endpoints never enters the rewrite at all: it calls straight + through to the original, so existing behaviour is bit-identical. + + Returns True when interior pins are available. + """ + if _INTERIOR_PINS["state"] is not None: + return _INTERIOR_PINS["state"] + try: + from comfy.ldm.minimax.model import PackedLayout + except Exception as error: # no H3 in this ComfyUI + _INTERIOR_PINS["state"] = False + print(f"[FunPack H3] interior keyframe pins unavailable ({error}).") + return False + if not _linear_rule_matches_upstream(): + _INTERIOR_PINS["state"] = False + print("[FunPack H3] interior keyframe pins DISABLED: this ComfyUI's video time grid " + "no longer matches the rule FunPack derives pin coordinates from. Pins stay " + "first/last only. (Upstream comfy/ldm/minimax/model.py changed.)") + return False + if getattr(PackedLayout.__init__, "_funpack_interior_pins", False): + _INTERIOR_PINS["state"] = True + return True + + original = PackedLayout.__init__ + + def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, + keyframes=None, refs=None, frame_count=None): + pins = list(keyframes or ()) + interior = [kf for kf in pins + if not keyframe_is_endpoint(kf.get("resolved_frame_index"), frame_count)] + if not interior: + original(self, text_len, latent_t, latent_h, latent_w, audio_t, + keyframes=keyframes, refs=refs, frame_count=frame_count) + return + # Declared at 0 so upstream's branch accepts them; the coordinate is corrected below. + safe = [dict(kf, resolved_frame_index=0) for kf in pins] + original(self, text_len, latent_t, latent_h, latent_w, audio_t, + keyframes=safe, refs=refs, frame_count=frame_count) + # Condition segments are emitted in `keyframes` order, immediately after text. + spans = [(a, b) for a, b, kind in self.segments if kind == "cond"] + for (start, stop), kf in zip(spans, pins): + self.position_ids[start:stop, 0] = keyframe_cond_t( + text_len, kf.get("resolved_frame_index", 0)) + + __init__._funpack_interior_pins = True + PackedLayout.__init__ = __init__ + _INTERIOR_PINS["state"] = True + return True + + +def keyframe_is_endpoint(frame_index, frame_count): + """True for the first or last pixel frame — the two pins upstream places unaided.""" + if frame_index is None: + return False + idx = int(frame_index) + if idx == 0: + return True + return frame_count is not None and idx == int(frame_count) - 1 + + +def keyframe_indices_supported(frame_index, frame_count): + """Can a pin be placed at this pixel frame? + + The first and last frames always can. Anything between them needs + :func:`install_interior_keyframes`, which is attempted here rather than assumed, so a + ComfyUI whose frame grid has moved refuses the pin instead of mis-placing it. """ idx = int(frame_index) - return idx == 0 or idx == int(frame_count) - 1 + if idx < 0 or idx > int(frame_count) - 1: + return False + if keyframe_is_endpoint(idx, frame_count): + return True + return install_interior_keyframes() def encode_keyframe(vae, image, width, height, frame_index, crop="disabled"): diff --git a/samplers.py b/samplers.py index 6b6e3d41..b73f8c6b 100644 --- a/samplers.py +++ b/samplers.py @@ -5501,22 +5501,36 @@ class FunPackLTXAVSceneChainSampler: condition rows into the sequence itself — they are never denoised and never rendered, so there is no latent to append and no tail to crop (the returned tail is 0). - The layout only accepts a pin at the FIRST or LAST pixel frame; ``PackedLayout`` raises - for anything else. A mid-clip request is refused here, loudly, rather than crashing - several seconds into the sample. + Stock ``PackedLayout`` places a pin only at the FIRST or LAST pixel frame and raises + for anything else. Those two are the endpoints of ONE straight line — the packed + sequence's time axis advances at a fixed rate per pixel frame — so + ``install_interior_keyframes`` (attempted by ``keyframe_indices_supported``) extends + the same rule to the frames between them. When it cannot install, a mid-clip request + is refused here, loudly, rather than crashing several seconds into the sample. + + An interior pin is EXPERIMENTAL in a way the endpoints are not: fl2va was trained with + condition rows at the two ends and nowhere between, so the coordinate is representable + (MM-RoPE is continuous in t) without being something the weights have seen. It costs no + extra model call either way. """ try: - from .minimax_h3 import keyframe_indices_supported + from .minimax_h3 import keyframe_indices_supported, keyframe_is_endpoint except ImportError: - from minimax_h3 import keyframe_indices_supported + from minimax_h3 import keyframe_indices_supported, keyframe_is_endpoint frame_count = max(1, int(self._h3_frame_count)) at = self._resolve_frame_index(frame_count, int(apply_at)) if not keyframe_indices_supported(at, frame_count): - print(f"[FunPackSceneChain] H3: guide at pixel frame {at} skipped — the packed layout " + print(f"[FunPackSceneChain] H3: guide at pixel frame {at} skipped — this ComfyUI " f"pins only the first (0) or last ({frame_count - 1}) frame. Use a reference " f"image (ref2va) for mid-clip guidance instead.") return positive, negative, 0 + if not keyframe_is_endpoint(at, frame_count): + _log.note_on_change( + "h3:interior_pin", + f"[FunPackSceneChain] H3: guide pinned mid-clip at pixel frame {at} of " + f"{frame_count}. EXPERIMENTAL — the checkpoint was trained on first/last pins " + f"only; if the frame does not land, move the guide to 0 or {frame_count - 1}.") pins = [{"resolved_frame_index": int(at), "latent": guide_frame}] # strength maps onto the DiT's condition noise augmentation: 1.0 pins the clean latent, diff --git a/tests/test_h3_interior_pins.py b/tests/test_h3_interior_pins.py new file mode 100644 index 00000000..fb69e9e1 --- /dev/null +++ b/tests/test_h3_interior_pins.py @@ -0,0 +1,241 @@ +"""A keyframe pin can sit at any pixel frame, not only the first or the last. + +Upstream ``PackedLayout`` computes a pin's time coordinate in two branches and raises for +anything between them. Those two branches are the endpoints of one straight line — the video +time axis advances ``FRAME_RESCALE`` per PIXEL frame — so an interior pin is a different point +on an axis the model already uses, not a new kind of conditioning. + +What is tested here: + +* the linear rule reproduces upstream's own two branches (against the REAL ComfyUI when one is + present, in a subprocess, because the suite stubs ``comfy``); +* the patch rewrites ONLY the condition rows' time column; +* an endpoint-only run never enters the rewrite, so existing behaviour is untouched; +* a changed upstream frame grid turns the feature OFF rather than mis-placing pins. +""" +import os +import subprocess +import sys +import types +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import minimax_h3 as h3 + +FRAME_PER_TOKEN = (1, 4, 4, 4, 4) +FRAME_RESCALE = 5.0 / 3.0 + + +def _spans(n): + return [FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] for k in range(n)] + + +def _frames_for(latent_t): + return sum(FRAME_PER_TOKEN[k % 5] for k in range(latent_t)) + + +class _FakePackedLayout: + """Upstream's keyframe branch, faithfully — including the raise.""" + + ROWS = 6 # frame_rows, one latent frame's spatial grid + + def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, + keyframes=None, refs=None, frame_count=None): + segments = [("text", text_len)] + pos = [torch.zeros(text_len, 3, dtype=torch.float64)] + for kf in (keyframes or ()): + pixel_index = kf["resolved_frame_index"] + if pixel_index == 0: + cond_t = float(text_len) + elif frame_count is not None and pixel_index == frame_count - 1: + cond_t = float(text_len) + sum(_spans(latent_t)) - FRAME_RESCALE + else: + raise ValueError("only first/last keyframe anchors are supported") + g = torch.zeros(self.ROWS, 3, dtype=torch.float64) + g[:, 0] = cond_t + g[:, 1] = 1.0 # a spatial column, to prove it is left alone + segments.append(("cond", self.ROWS)) + pos.append(g) + segments.append(("video", latent_t * self.ROWS)) + pos.append(torch.full((latent_t * self.ROWS, 3), 9.0, dtype=torch.float64)) + + self.position_ids = torch.cat(pos) + self.signature = (text_len, latent_t, latent_h, latent_w, audio_t) + seg_abs, off = [], 0 + for kind, n in segments: + seg_abs.append((off, off + n, kind)) + off += n + self.segments = seg_abs + + +@pytest.fixture +def upstream(monkeypatch): + """Install a fake comfy.ldm.minimax.model and reset the patch cache.""" + mod = types.ModuleType("comfy.ldm.minimax.model") + mod.PackedLayout = type("PackedLayout", (_FakePackedLayout,), {}) + mod.FRAME_PER_TOKEN = FRAME_PER_TOKEN + mod.FRAME_RESCALE = FRAME_RESCALE + mod._video_t_spans = _spans + for name in ("comfy.ldm.minimax", "comfy.ldm.minimax.model"): + monkeypatch.setitem(sys.modules, name, mod) + monkeypatch.setitem(h3._INTERIOR_PINS, "state", None) + return mod + + +def _build(mod, keyframes, text_len=7, latent_t=6): + return mod.PackedLayout(text_len, latent_t, 4, 4, 3, + keyframes=keyframes, frame_count=_frames_for(latent_t)) + + +def _cond_t(layout): + return [float(layout.position_ids[a, 0]) for a, _b, kind in layout.segments if kind == "cond"] + + +# ── the rule ──────────────────────────────────────────────────────────────── + +def test_the_rule_reproduces_both_upstream_branches(upstream): + """The whole feature rests on this: first/last are two points on one line.""" + assert h3._linear_rule_matches_upstream() + + +def test_the_rule_is_linear_in_the_pixel_index(): + step = h3.keyframe_cond_t(7, 1) - h3.keyframe_cond_t(7, 0) + assert step == pytest.approx(FRAME_RESCALE) + assert h3.keyframe_cond_t(7, 12) == pytest.approx(7.0 + FRAME_RESCALE * 12) + + +def test_a_changed_frame_grid_disables_the_feature(upstream, capsys): + """Drift must turn pins off, not move them. FunPack patches ComfyUI by name and has + been silently broken by an upstream grid change before.""" + upstream.FRAME_RESCALE = 2.0 + assert h3.install_interior_keyframes() is False + assert h3.keyframe_indices_supported(8, 18) is False + assert "DISABLED" in capsys.readouterr().out + + +def test_a_comfy_without_h3_declines_quietly(monkeypatch, capsys): + monkeypatch.setitem(h3._INTERIOR_PINS, "state", None) + monkeypatch.setitem(sys.modules, "comfy.ldm.minimax.model", None) + assert h3.install_interior_keyframes() is False + assert "unavailable" in capsys.readouterr().out + + +# ── the patch ─────────────────────────────────────────────────────────────── + +def test_an_interior_pin_is_refused_before_the_patch(upstream): + with pytest.raises(ValueError): + _build(upstream, [{"resolved_frame_index": 8, "latent": None}]) + + +def test_an_interior_pin_lands_on_its_own_coordinate(upstream): + assert h3.install_interior_keyframes() is True + layout = _build(upstream, [{"resolved_frame_index": 8, "latent": None}]) + assert _cond_t(layout) == [pytest.approx(h3.keyframe_cond_t(7, 8))] + + +def test_the_pin_sits_between_the_two_endpoints(upstream): + h3.install_interior_keyframes() + first = _cond_t(_build(upstream, [{"resolved_frame_index": 0}]))[0] + mid = _cond_t(_build(upstream, [{"resolved_frame_index": 8}]))[0] + last = _cond_t(_build(upstream, [{"resolved_frame_index": 17}]))[0] + assert first < mid < last + + +def test_only_the_time_column_is_rewritten(upstream): + """The pin keeps the target's spatial grid — that is how it addresses a frame at all.""" + h3.install_interior_keyframes() + layout = _build(upstream, [{"resolved_frame_index": 8}]) + start, stop = [(a, b) for a, b, k in layout.segments if k == "cond"][0] + assert torch.all(layout.position_ids[start:stop, 1] == 1.0) + assert torch.all(layout.position_ids[start:stop, 2] == 0.0) + + +def test_rows_outside_the_pin_are_untouched(upstream): + h3.install_interior_keyframes() + layout = _build(upstream, [{"resolved_frame_index": 8}]) + video = [(a, b) for a, b, k in layout.segments if k == "video"][0] + assert torch.all(layout.position_ids[video[0]:video[1]] == 9.0) + assert torch.all(layout.position_ids[:7] == 0.0) + + +def test_several_pins_each_get_their_own_coordinate(upstream): + """Segments are emitted in `keyframes` order, so the k-th cond span is the k-th pin.""" + h3.install_interior_keyframes() + layout = _build(upstream, [{"resolved_frame_index": 0}, + {"resolved_frame_index": 8}, + {"resolved_frame_index": 17}]) + assert _cond_t(layout) == [pytest.approx(h3.keyframe_cond_t(7, i)) for i in (0, 8, 17)] + + +def test_an_endpoint_only_run_is_bit_identical(upstream): + """The patch must be invisible to every run that does not use it.""" + before = _build(upstream, [{"resolved_frame_index": 0}, + {"resolved_frame_index": 17}]).position_ids.clone() + h3.install_interior_keyframes() + after = _build(upstream, [{"resolved_frame_index": 0}, + {"resolved_frame_index": 17}]).position_ids + assert torch.equal(before, after) + + +def test_a_pinless_run_is_bit_identical(upstream): + before = _build(upstream, []).position_ids.clone() + h3.install_interior_keyframes() + assert torch.equal(before, _build(upstream, []).position_ids) + + +def test_patching_twice_does_not_stack(upstream): + assert h3.install_interior_keyframes() is True + first = upstream.PackedLayout.__init__ + h3._INTERIOR_PINS["state"] = None + assert h3.install_interior_keyframes() is True + assert upstream.PackedLayout.__init__ is first + + +# ── the gate callers use ──────────────────────────────────────────────────── + +def test_endpoints_never_need_the_patch(monkeypatch): + """A first/last pin must keep working on a ComfyUI where the patch cannot install.""" + monkeypatch.setitem(h3._INTERIOR_PINS, "state", False) + assert h3.keyframe_indices_supported(0, 18) is True + assert h3.keyframe_indices_supported(17, 18) is True + assert h3.keyframe_indices_supported(8, 18) is False + + +def test_out_of_range_indices_are_still_refused(upstream): + h3.install_interior_keyframes() + assert h3.keyframe_indices_supported(-1, 18) is False + assert h3.keyframe_indices_supported(18, 18) is False + assert h3.keyframe_indices_supported(17, 18) is True + + +# ── against the real ComfyUI, if one is here ──────────────────────────────── + +REAL_COMFY = Path.home() / "Documents" / "ComfyUI" + + +@pytest.mark.skipif(not (REAL_COMFY / "comfy" / "ldm" / "minimax" / "model.py").exists(), + reason="no local ComfyUI with MiniMax H3 to check against") +def test_the_rule_matches_the_real_upstream_layout(): + """The suite stubs `comfy`, so this re-derives the equality in a clean subprocess from + upstream's own constants. This is the test that catches a real ComfyUI update.""" + script = ( + "import sys\n" + "sys.path.insert(0, %r)\n" + "from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE, _video_t_spans\n" + "RATE = 5.0 / 3.0\n" + "assert abs(float(FRAME_RESCALE) - RATE) < 1e-9, 'FRAME_RESCALE moved'\n" + "for latent_t in range(1, 40):\n" + " spans = _video_t_spans(latent_t)\n" + " frames = sum(FRAME_PER_TOKEN[k %% 5] for k in range(latent_t))\n" + " last = 7.0 + sum(spans) - FRAME_RESCALE\n" + " assert abs(7.0 - (7.0 + RATE * 0)) < 1e-9\n" + " assert abs(last - (7.0 + RATE * (frames - 1))) < 1e-9, (latent_t, last)\n" + "print('OK')\n" % str(REAL_COMFY)) + env = dict(os.environ, PYTHONPATH="") + out = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, + env=env, cwd=str(REAL_COMFY)) + assert "OK" in out.stdout, out.stderr[-2000:]