#!/usr/bin/env python3 """Unit test for H3 DyRoPE's layer-wise / sigma-faded time geometry. python tests/test_dyrope.py CPU only, no GPU, no files. Needs comfy on sys.path (the real comfy.ldm.minimax.model is the reference the compact grid is checked against, and the real rope_freqs is what the wrapper chains onto). Six properties, one per section: 1. COMPACT GRID: the node's rebuilt "stock" grid is exactly the core's _video_t_grid(n, origin) — DyRoPE must not invent its own idea of the geometry the model was trained on, 2. PHYSICAL SPANS: holds all-1 makes physical == compact (True Clock's own identity property, re-asserted here because DyRoPE's whole premise is that the two lists differ), 3. BLOCK MAP: physical_blocks 0..24 patches exactly {0..24} and nothing else; compact_blocks patches the same set with the inverse meaning, 4. FADE WEIGHTS: 1.0 at sigma_max, 0.0 at fade_end, 0.5 at the midpoint, clamped outside, 5. ROPE WRAPPER: on a hand-built position_ids (text + audio + video rows) the video segment is identified correctly, the alternate t column is the compact grid, and every non-video row is bit-identical, 6. IDENTITY ARMS: physical_all never requests the alternate table and registers no block patches; compact_all leaves _TRUE_CLOCK disarmed for the whole wrapped sample call. Exit code 0 = pass. """ import json import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.dirname(HERE)) import importlib.util spec = importlib.util.spec_from_file_location( "mainodes_motion", os.path.join(os.path.dirname(HERE), "motion.py")) motion = importlib.util.module_from_spec(spec) spec.loader.exec_module(motion) import torch import comfy.ldm.minimax.model as mm R = motion.ROPE_UNITS_PER_FRAME FAILS = [] # captured BEFORE anything in this file arms the node: section 3 installs the # rope_freqs chain as a side effect of wrap(), so "stock" has to be taken now STOCK_ROPE_FREQS = mm.MiniMaxH3Model.rope_freqs assert not getattr(STOCK_ROPE_FREQS, "_h3_dyrope", False) def check(name, ok, detail): print((" PASS " if ok else " FAIL ") + name + " " + detail) if not ok: FAILS.append(name) def reset_state(): motion._TRUE_CLOCK["spans"] = None motion._DYROPE.update({"active": False, "mode": None, "n_tokens": 0, "spans_phys": None, "spans_comp": None, "blocks": (), "fade_end": 0.5, "sigma_max": None, "sigma": None, "alt_angles": None, "alt_table": None, "alt_requested": 0, "last_weight": None}) # ---- stubs ----------------------------------------------------------------- class FakeBlock: pass class FakeDiffusionModel: def __init__(self, n=50): self.blocks = [FakeBlock() for _ in range(n)] class FakeInner: """Innermost model object; model.model.diffusion_model is what the node reads.""" def __init__(self, n=50): self.diffusion_model = FakeDiffusionModel(n) class FakePatcher: """Minimum of ModelPatcher the node touches: clone + the two registrars.""" def __init__(self, n=50): self.model = FakeInner(n) self.replaces = {} self.wrappers = {} self.model_options = {"transformer_options": {}} self.object_patches = {} self.clones = 0 def clone(self): c = FakePatcher.__new__(FakePatcher) c.model = self.model c.replaces = dict(self.replaces) c.wrappers = {k: dict(v) for k, v in self.wrappers.items()} c.model_options = self.model_options c.object_patches = dict(self.object_patches) c.clones = self.clones + 1 return c def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None): self.replaces[(name, block_name, number)] = patch def add_wrapper(self, wrapper_type, wrapper): self.wrappers.setdefault(wrapper_type, {}).setdefault(None, []).append(wrapper) class StubSampler: """Inner SAMPLER; records the module state visible DURING sample().""" def __init__(self): self.seen = None def max_denoise(self, model_wrap, sigmas): return False def sample(self, *args, **kwargs): self.seen = {"clock": motion._TRUE_CLOCK["spans"], "active": motion._DYROPE["active"], "mode": motion._DYROPE["mode"], "sigma_max": motion._DYROPE["sigma_max"]} return "SAMPLED" class StubRope(torch.nn.Module): """Just enough of MiniMaxH3Model for the real rope_freqs to run.""" def __init__(self, length=16): super().__init__() self.rope = torch.nn.Module() self.rope.register_buffer("inv_freq", torch.linspace(1.0, 0.01, length)) HOLDS = [4 if 132 <= f <= 150 else 1 for f in range(243)] # the pier cell's map HOLD_MAP = json.dumps({"holds": HOLDS}) # ---- 1) compact grid rebuild ---------------------------------------------- print("1) compact grid rebuild == comfy.ldm.minimax.model._video_t_grid(n, origin)") for n, origin in [(5, 0.0), (61, 123.0), (243, 7.5), (1, 0.0), (17, 512.0)]: mine = motion.dyrope_grid(motion.dyrope_stock_spans(n), origin) core = mm._video_t_grid(n, origin).tolist() worst = max(abs(a - b) for a, b in zip(mine, core)) if n else 0.0 check("n=%d origin=%s" % (n, origin), len(mine) == len(core) and worst == 0.0, "len %d/%d worst abs diff %.3e (first %.6f last %.6f)" % (len(mine), len(core), worst, mine[0], mine[-1])) # and the spans themselves against the core helper for n in (5, 61, 243): check("stock spans n=%d" % n, motion.dyrope_stock_spans(n) == mm._video_t_spans(n), "sum %.6f == %.6f" % (sum(motion.dyrope_stock_spans(n)), sum(mm._video_t_spans(n)))) # ---- 2) physical spans, holds all-1 --------------------------------------- print("2) true_clock_spans(all-1) == compact spans") # LEGAL lengths only (17k+5): at an illegal length _snap_holds pads the clip, # and pad frames are copies of the last world frame that cost no world time, so # physical is legitimately SHORTER than compact even with holds all 1. for length in (39, 73, 141, 243): phys = motion.true_clock_spans([1] * length) comp = motion.dyrope_stock_spans(len(phys)) worst = max(abs(a - b) for a, b in zip(phys, comp)) check("holds=[1]*%d" % length, len(phys) == len(comp) and worst < 1e-12, "%d tokens, worst abs diff %.3e, sums %.6f / %.6f" % (len(phys), worst, sum(phys), sum(comp))) phys = motion.true_clock_spans(HOLDS) comp = motion.dyrope_stock_spans(len(phys)) check("pier map (holds 4 on 132-150) really differs", phys != comp and sum(phys) < sum(comp), "%d tokens, physical sum %.6f (= 243 x 5/3 = %.6f), compact sum %.6f" % (len(phys), sum(phys), 243 * R, sum(comp))) # ---- 3) block map ---------------------------------------------------------- print("3) block map: which blocks get a double_block replacement") def wrap_once(mode, lo=0, hi=24, fade_end=0.5, model=None, sampler=None): reset_state() m = model or FakePatcher(50) s = sampler or StubSampler() out = motion.H3DyRoPE().wrap(m, s, HOLD_MAP, mode, lo, hi, fade_end) return m, s, out for mode, lo, hi in [("physical_blocks", 0, 24), ("compact_blocks", 0, 24), ("physical_blocks", 25, 49), ("physical_blocks", 40, 49)]: _, _, (out_model, _, report) = wrap_once(mode, lo, hi) keys = sorted(k[2] for k in out_model.replaces) want = list(range(lo, hi + 1)) unpatched = sorted(set(range(50)) - set(keys)) check("%s %d..%d" % (mode, lo, hi), keys == want and all(k[:2] == ("dit", "double_block") for k in out_model.replaces), "patched %d block(s) %d..%d, unpatched %d (%d..%d)" % (len(keys), keys[0], keys[-1], len(unpatched), unpatched[0], unpatched[-1])) _, _, (out_model, _, _) = wrap_once("physical_all") check("physical_all registers nothing", out_model.replaces == {} and out_model.clones == 0, "replaces=%r clones=%d (same object returned)" % (out_model.replaces, out_model.clones)) _, _, (out_model, _, _) = wrap_once("compact_all") check("compact_all registers nothing", out_model.replaces == {} and out_model.clones == 0, "replaces=%r clones=%d" % (out_model.replaces, out_model.clones)) _, _, (out_model, _, _) = wrap_once("fade_physical_to_compact") n_w = sum(len(v) for d in out_model.wrappers.values() for v in d.values()) check("fade registers no block patch, one diffusion_model wrapper", out_model.replaces == {} and n_w == 1, "replaces=%r wrappers=%r" % (out_model.replaces, list(out_model.wrappers))) # ---- 4) fade weights ------------------------------------------------------- print("4) fade weight: 1.0 at sigma_max, 0.0 at fade_end, 0.5 at the midpoint") S_MAX, F_END = 1.0, 0.5 cases = [(1.0, 1.0), (0.5, 0.0), (0.75, 0.5), (2.0, 1.0), (0.0, 0.0), (0.625, 0.25)] for sigma, want in cases: got = motion.dyrope_fade_weight(sigma, S_MAX, F_END) check("sigma=%.3f" % sigma, abs(got - want) < 1e-12, "w=%.6f (want %.6f)" % (got, want)) check("degenerate sigma_max == fade_end", motion.dyrope_fade_weight(0.4, 0.5, 0.5) == 0.0 and motion.dyrope_fade_weight(0.6, 0.5, 0.5) == 1.0, "below -> %.1f, above -> %.1f" % (motion.dyrope_fade_weight(0.4, 0.5, 0.5), motion.dyrope_fade_weight(0.6, 0.5, 0.5))) # ---- 5) rope wrapper on a hand-built position_ids -------------------------- print("5) rope_freqs wrapper: video-row detection, alternate t column, other rows untouched") TEXT_ROWS, AUDIO_T, ROWS_PER_FRAME = 11, 7, 6 spans_phys = motion.true_clock_spans(HOLDS) T_LAT = len(spans_phys) spans_comp = motion.dyrope_stock_spans(T_LAT) ORIGIN = float(TEXT_ROWS) pos = torch.zeros(TEXT_ROWS + AUDIO_T * 2 + T_LAT * ROWS_PER_FRAME, 3, dtype=torch.float64) pos[:TEXT_ROWS, 0] = torch.arange(TEXT_ROWS, dtype=torch.float64) a0 = TEXT_ROWS pos[a0:a0 + AUDIO_T * 2, 0] = (ORIGIN + torch.arange(AUDIO_T, dtype=torch.float64)).repeat(2) pos[a0:a0 + AUDIO_T, 2] = -1.5 pos[a0 + AUDIO_T:a0 + AUDIO_T * 2, 2] = 1.5 v0 = a0 + AUDIO_T * 2 g_phys = motion.dyrope_grid(spans_phys, ORIGIN) pos[v0:, 0] = torch.tensor(g_phys, dtype=torch.float64).repeat_interleave(ROWS_PER_FRAME) pos[v0:, 1] = torch.arange(T_LAT * ROWS_PER_FRAME, dtype=torch.float64) % 3.0 - 1.0 pos[v0:, 2] = torch.arange(T_LAT * ROWS_PER_FRAME, dtype=torch.float64) % 2.0 - 0.5 start, rpf = motion.dyrope_video_rows(pos, T_LAT) check("video rows located", (start, rpf) == (v0, ROWS_PER_FRAME), "start=%d rows_per_frame=%d (want %d/%d), seq_len=%d" % (start, rpf, v0, ROWS_PER_FRAME, pos.shape[0])) alt = motion.dyrope_retimed_position_ids(pos, start, rpf, motion.dyrope_grid(spans_comp, ORIGIN)) g_comp = motion.dyrope_grid(spans_comp, ORIGIN) got = alt[start::rpf, 0].tolist() check("alternate t column == compact grid", max(abs(a - b) for a, b in zip(got, g_comp)) < 1e-12, "first %.6f/%.6f last %.6f/%.6f" % (got[0], g_comp[0], got[-1], g_comp[-1])) check("non-video rows bit-identical", bool(torch.equal(alt[:start], pos[:start])), "%d rows (text %d + audio %d) compared with torch.equal" % (start, TEXT_ROWS, AUDIO_T * 2)) check("video h/w columns bit-identical", bool(torch.equal(alt[start:, 1:], pos[start:, 1:])), "%d video rows, columns 1:3" % (pos.shape[0] - start)) check("alternate t column is per-frame constant", bool(torch.equal(alt[start:, 0].view(T_LAT, rpf), alt[start::rpf, 0][:, None].expand(T_LAT, rpf))), "%d tokens x %d rows" % (T_LAT, rpf)) # now the installed wrapper, end to end stub = StubRope() stock_fn = STOCK_ROPE_FREQS # unpatched core method stock_angles = stock_fn(stub, pos, torch.device("cpu")) motion._install_dyrope_rope_patch() plain = mm.MiniMaxH3Model.rope_freqs # our chained version check("wrapper installed once (idempotent)", plain is not stock_fn and getattr(plain, "_h3_dyrope", False) and (motion._install_dyrope_rope_patch() or mm.MiniMaxH3Model.rope_freqs is plain), "rope_freqs marked _h3_dyrope, second install is a no-op") reset_state() un_armed = plain(stub, pos, torch.device("cpu")) check("disarmed wrapper is a bit-identical pass-through", bool(torch.equal(un_armed, stock_angles)), "shape %s, torch.equal vs the unpatched core method" % (tuple(un_armed.shape),)) motion._DYROPE.update({"active": True, "mode": "physical_blocks", "n_tokens": T_LAT, "spans_phys": spans_phys, "spans_comp": spans_comp}) default = plain(stub, pos, torch.device("cpu")) altang = motion._DYROPE["alt_angles"] check("physical_blocks: default table is COMPACT, alternate is PHYSICAL", altang is not None and torch.equal(altang, un_armed) and not torch.equal(default, un_armed), "alt == the physical angles: %s; default differs: %s; shapes %s" % (torch.equal(altang, un_armed), not torch.equal(default, un_armed), tuple(default.shape))) motion._DYROPE.update({"mode": "compact_blocks", "alt_angles": None}) default_c = plain(stub, pos, torch.device("cpu")) check("compact_blocks: default table is PHYSICAL, alternate is COMPACT", torch.equal(default_c, un_armed) and torch.equal(motion._DYROPE["alt_angles"], default), "default == physical angles: %s; alt == the compact angles: %s" % (torch.equal(default_c, un_armed), torch.equal(motion._DYROPE["alt_angles"], default))) motion._DYROPE.update({"mode": "fade_physical_to_compact", "alt_angles": None, "sigma_max": 1.0, "fade_end": 0.5, "sigma": 1.0}) f_hi = plain(stub, pos, torch.device("cpu")) motion._DYROPE["sigma"] = 0.5 motion._DYROPE["alt_angles"] = None f_lo = plain(stub, pos, torch.device("cpu")) check("fade at sigma_max == physical, at fade_end == compact, no alternate", torch.equal(f_hi, un_armed) and torch.equal(f_lo, default) and motion._DYROPE["alt_angles"] is None, "w(1.0)=%.3f w(0.5)=%.3f" % (motion.dyrope_fade_weight(1.0, 1.0, 0.5), motion.dyrope_fade_weight(0.5, 1.0, 0.5))) motion._DYROPE["sigma"] = 0.75 motion._DYROPE["alt_angles"] = None plain(stub, pos, torch.device("cpu")) mid = [0.5 * a + 0.5 * b for a, b in zip(spans_phys, spans_comp)] check("fade midpoint interpolates SPANS (monotone grid)", abs(motion._DYROPE["last_weight"] - 0.5) < 1e-12 and all(s > 0 for s in mid) and all(b > a for a, b in zip(motion.dyrope_grid(mid), motion.dyrope_grid(mid)[1:])), "w=%.3f, %d mixed spans all positive, grid strictly increasing" % (motion._DYROPE["last_weight"], len(mid))) # combo: fade_physical_blocks = faded grid to the named blocks, compact default motion._DYROPE.update({"mode": "fade_physical_blocks", "alt_angles": None, "sigma_max": 1.0, "fade_end": 0.5, "sigma": 1.0}) c_hi = plain(stub, pos, torch.device("cpu")) c_hi_alt = motion._DYROPE["alt_angles"] check("combo at sigma_max: default COMPACT, alternate == PHYSICAL", torch.equal(c_hi, default) and c_hi_alt is not None and torch.equal(c_hi_alt, un_armed), "default == compact angles: %s; alt == physical angles: %s" % (torch.equal(c_hi, default), torch.equal(c_hi_alt, un_armed))) motion._DYROPE["sigma"] = 0.5 motion._DYROPE["alt_angles"] = None c_lo = plain(stub, pos, torch.device("cpu")) check("combo at fade_end: alternate collapses to COMPACT (whole model stock)", torch.equal(c_lo, default) and torch.equal(motion._DYROPE["alt_angles"], default), "default and alt both == compact angles at w=0") motion._DYROPE["sigma"] = 0.75 motion._DYROPE["alt_angles"] = None plain(stub, pos, torch.device("cpu")) c_mid = motion._DYROPE["alt_angles"] check("combo midpoint: alternate strictly between the two geometries", c_mid is not None and not torch.equal(c_mid, un_armed) and not torch.equal(c_mid, default) and abs(motion._DYROPE["last_weight"] - 0.5) < 1e-12, "w=%.3f, alt differs from both pure tables" % motion._DYROPE["last_weight"]) # the mis-rotation guard bad = pos.clone() bad[v0 + 3 * rpf, 0] += 0.25 motion._DYROPE.update({"mode": "physical_blocks", "alt_angles": None, "sigma": None}) try: plain(stub, bad, torch.device("cpu")) check("mis-rotation guard raises", False, "no exception") except RuntimeError as e: check("mis-rotation guard raises", "do not carry the armed t-grid" in str(e), str(e).split("\n")[0][:110]) reset_state() # ---- 6) identity arms ------------------------------------------------------ print("6) identity arms: physical_all == True Clock, compact_all == no node") m, s, (out_model, wrapped, report) = wrap_once("physical_all") check("physical_all arms the True Clock spans at node time", motion._TRUE_CLOCK["spans"] == motion.true_clock_spans(HOLDS) and motion._DYROPE["active"] is False, "spans len %d, _DYROPE active=%s" % (len(motion._TRUE_CLOCK["spans"]), motion._DYROPE["active"])) sig = torch.linspace(0.7, 0.0, 7) res = wrapped.sample(object(), sig) check("physical_all: True Clock armed during sample, no alternate requested", res == "SAMPLED" and s.seen["clock"] == motion.true_clock_spans(HOLDS) and s.seen["active"] is False and motion._DYROPE["alt_requested"] == 0, "seen clock len=%d active=%s alt_requested=%d" % (len(s.seen["clock"]), s.seen["active"], motion._DYROPE["alt_requested"])) check("physical_all: disarmed after sample", motion._TRUE_CLOCK["spans"] is None and motion._DYROPE["active"] is False, "_TRUE_CLOCK[spans]=%r _DYROPE[active]=%r" % (motion._TRUE_CLOCK["spans"], motion._DYROPE["active"])) m, s, (out_model, wrapped, report) = wrap_once("compact_all") check("compact_all leaves True Clock disarmed at node time", motion._TRUE_CLOCK["spans"] is None and motion._DYROPE["active"] is False, "_TRUE_CLOCK[spans]=%r" % (motion._TRUE_CLOCK["spans"],)) wrapped.sample(object(), sig) check("compact_all: _TRUE_CLOCK[spans] is None DURING the sample call", s.seen["clock"] is None and s.seen["active"] is False and motion._DYROPE["alt_requested"] == 0, "seen clock=%r active=%r alt_requested=%d" % (s.seen["clock"], s.seen["active"], motion._DYROPE["alt_requested"])) m, s, (out_model, wrapped, report) = wrap_once("fade_compact_to_physical") wrapped.sample(object(), sig) check("fade: sigma_max read off the schedule's first sigma", abs(s.seen["sigma_max"] - float(sig[0])) < 1e-9 and s.seen["active"] is True, "sigma_max=%.6f (sigmas[0]=%.6f) active=%s" % (s.seen["sigma_max"], float(sig[0]), s.seen["active"])) reset_state() print("7) report string") _, _, (_, _, report) = wrap_once("physical_blocks", 0, 24) print(" " + report.replace("\n", "\n ")) check("report names mode, tokens, both sums and the block split", "physical_blocks" in report and "tokens=%d" % len(spans_phys) in report and "blocks 0-24 -> physical" in report and "%.6f" % (243 * R) in report, "%d chars" % len(report)) reset_state() # ---- presets --------------------------------------------------------------- print("presets: override the knobs; custom leaves them alone") import motion as _m check("preset table has custom + three measured settings", list(_m.DYROPE_PRESETS)[0] == "custom" and len(_m.DYROPE_PRESETS) == 4, ", ".join(_m.DYROPE_PRESETS)) opt = _m.H3DyRoPE.INPUT_TYPES().get("optional", {}) check("preset is an OPTIONAL input appended after every knob (old graphs unaffected)", list(opt) == ["preset"] and opt["preset"][1]["default"] == "custom", str(list(opt))) req = list(_m.H3DyRoPE.INPUT_TYPES()["required"]) check("required knobs unchanged in order", req == ["model", "sampler", "hold_map", "mode", "block_lo", "block_hi", "fade_end"], str(req)) print() if FAILS: print("FAILED: " + ", ".join(FAILS)) sys.exit(1) print("all dyrope checks passed")