Files
olivv-cs-ComfyUI-FunPack/tests/test_h3_interior_pins.py
T
digital-garbage 54fa8e4b43 feat(h3): region-locked anchors, and the taste direction lands after the refiner
Region lock: a keyframe pin can cover part of a frame. Transparent areas of an
i2v anchor are dropped from the packed sequence entirely — not blanked, not
noised, absent — so the model invents them. Row deletion re-derives the layout
from upstream's own output; two patches share one keep-mask so they cannot
disagree. Surfaced through the anchor's alpha, which needs no new wire and works
in the Editor. Excluded pixels are filled with the kept area's mean before the
VAE so no hard edge bleeds in. EXPERIMENTAL: fl2va saw whole condition frames.

Taste bias: the liked direction Studio already learns is now also added AFTER
H3's token refiner, where it lands as sent instead of being mixed across tokens
and normalized away. Relative to the span's own row norm, so one strength means
the same thing on every prompt. How far to push is a fifth learned gain,
centred on 0 and signed; h3_taste_bias is the manual override. Needs 3 liked
runs.

Gain state generalized to per-key centres, bounds and exploration width.

Tests: 1608 passed (+46). Untested on a rental.
2026-08-25 21:48:28 +03:00

245 lines
10 KiB
Python

"""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)
monkeypatch.setitem(h3._REGION_LOCKS, "state", None)
monkeypatch.setitem(h3._LAYOUT_PATCH, "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(h3._LAYOUT_PATCH, "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:]