Download ComfyUI-MiniMaxH3-Context-Loop/tests/_mock_harness.py from dzmdd/Fast_deploy: direct link, hf CLI and curl.
- Browser
- Download file 31.8 kB
-
https://huggingface.co/dzmdd/Fast_deploy/resolve/main/ComfyUI-MiniMaxH3-Context-Loop/tests/_mock_harness.py
- Command line
-
hf download hf://dzmdd/Fast_deploy/ComfyUI-MiniMaxH3-Context-Loop/tests/_mock_harness.py
-
curl -L -o _mock_harness.py https://huggingface.co/dzmdd/Fast_deploy/resolve/main/ComfyUI-MiniMaxH3-Context-Loop/tests/_mock_harness.py
31.8 kB
| """Standalone harness for patch_layout, no ComfyUI and no GPU needed. | |
| Fakes `comfy.ldm.minimax.model` with a PackedLayout that reproduces the | |
| structure of the real one in comfy/ldm/minimax/model.py: | |
| segment order text, then keyframe cond rows, then reference blocks, | |
| then target audio, then target video. Reference blocks | |
| come AFTER the keyframe rows, and the target rows are | |
| always the last two segments. | |
| segment kinds text / cond / ref_img / ref_audio / audio / video | |
| position_ids [S, 3] float64, columns (t, h, w) | |
| cond rows one segment per keyframe, at STOCK coordinates computed | |
| from text_len and NOT from the reference cursor, which | |
| is the uncompensated behaviour the patch exists to fix, | |
| and rejecting interior anchors | |
| references laid out from a cursor starting at text_len: | |
| image one ref_img segment of (h//2)*(w//2) | |
| rows, all on the cursor; cursor += 1 | |
| audio one ref_audio segment of rt*2 rows, | |
| channel-major stereo, coordinates | |
| cursor..cursor+rt-1 once per channel; | |
| cursor += rt. Skipped entirely when rt | |
| is 0, which still advances. | |
| video, ref_audio then ref_img, both sharing the | |
| video_audio cursor origin; cursor advances by | |
| max(rt, sum of the video spans) | |
| target audio rows from the final cursor, then video rows whose | |
| first row sits exactly on it | |
| An earlier version of this file put references before the cond rows, used | |
| four position columns and labelled every reference block `ref`. All three | |
| were wrong, and stopped mattering only because nothing read the segment | |
| table. Anything that does read it needs this to be faithful. | |
| Checks: | |
| 1. apply_patch succeeds against the faithful layout | |
| 2. interior anchors build, and reference compensation lands | |
| 3. the marked audio block moves onto the target timeline while a | |
| Ref2VA graph's own references stay put | |
| 4. the mixed stock/MC keyframe guard trips | |
| 5. a layout emitting the wrong number of rows for an audio reference is | |
| caught rather than silently half-moved | |
| 6. a layout whose reference cursor advances differently still produces | |
| correctly aligned output, because nothing here recomputes it | |
| 7. a second copy of the patch, as another pack vendoring the same file | |
| would load, stands down instead of wrapping the first copy | |
| 8. a DIFFERENT copy (a fork, or a version predating the marker) is also | |
| recognised and stood down for | |
| 9. a wrapper from an unrelated pack solving the same problem its own | |
| way is recognised too, and refused rather than stacked on. Several | |
| H3 packs lift the same restriction independently and they cannot | |
| both own the constructor. | |
| 10. SolAttn's H3 Morton layout observer composes when it loads first. | |
| 11. When this patch loads first and SolAttn wraps it, a second vendored copy | |
| sees the matching patch through SolAttn and stands down. | |
| 12. The CUDA PR117 checkout's path-valued wrapper identity composes in both | |
| load orders, but only with the same audited ``original_init`` closure. | |
| 13. Renamed/path-loaded copies of the read-only SolAttn observer are | |
| recognised structurally, including defensive getattr access and a nested | |
| observer stack, while a partial fingerprint remains refused. | |
| """ | |
| import importlib | |
| import importlib.util | |
| import os | |
| import sys | |
| import types | |
| import numpy as np | |
| FRAME_RESCALE = 5.0 / 3.0 | |
| FRAME_PER_TOKEN = (1, 4, 4, 4, 4) | |
| def _video_t_spans(latent_t): | |
| return [FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] for k in range(latent_t)] | |
| def _frame_grid(h, w): | |
| """[n, 2] (h, w) coordinates for one latent frame's 2x2-patch rows. | |
| The values only have to be deterministic and shaped right; the patch | |
| never writes these columns, it only has to leave them alone. | |
| """ | |
| n_h, n_w = h // 2, w // 2 | |
| hh, ww = np.meshgrid(np.arange(n_h, dtype=np.float64), | |
| np.arange(n_w, dtype=np.float64), indexing="ij") | |
| return np.stack([hh.reshape(-1), ww.reshape(-1)], axis=-1) | |
| def _audio_grid(cursor, t, rows_per_step=2): | |
| g = np.zeros((t * rows_per_step, 3), dtype=np.float64) | |
| g[:, 0] = np.tile(cursor + np.arange(t, dtype=np.float64), rows_per_step) | |
| if rows_per_step == 2: | |
| g[:t, 2] = -1.0 | |
| g[t:, 2] = 1.0 | |
| return g | |
| def _video_t_grid(n, origin): | |
| spans = np.array(_video_t_spans(n), dtype=np.float64) | |
| head = np.cumsum(spans[:-1]) if n > 1 else np.zeros(0, dtype=np.float64) | |
| return float(origin) + np.concatenate(([0.0], head)) | |
| def _video_grid(vt, frame, cursor): | |
| g = np.empty((vt, frame.shape[0], 3), dtype=np.float64) | |
| g[:, :, 0] = _video_t_grid(vt, cursor)[:, None] | |
| g[:, :, 1:] = frame[None] | |
| return g.reshape(-1, 3) | |
| def make_mm(ref_advance_factor=1.0, audio_rows_per_step=2): | |
| """Build a fake comfy.ldm.minimax.model. | |
| ref_advance_factor != 1 simulates an upstream change to how an audio | |
| reference advances the cursor. audio_rows_per_step != 2 simulates a | |
| change to how many rows one audio latent step occupies. | |
| """ | |
| mm = types.ModuleType("comfy.ldm.minimax.model") | |
| mm.FRAME_RESCALE = FRAME_RESCALE | |
| mm.FRAME_PER_TOKEN = FRAME_PER_TOKEN | |
| mm._video_t_spans = _video_t_spans | |
| class PackedLayout: | |
| def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, | |
| keyframes=None, refs=None, frame_count=None): | |
| frame = _frame_grid(latent_h, latent_w) | |
| frame_rows = frame.shape[0] | |
| segs, blocks = [], [] | |
| def emit(kind, g): | |
| segs.append((kind, g.shape[0])) | |
| blocks.append(g) | |
| g = np.zeros((text_len, 3), dtype=np.float64) | |
| g[:, 0] = np.arange(text_len, dtype=np.float64) | |
| emit("text", g) | |
| spans = _video_t_spans(latent_t) | |
| for kf in (keyframes or []): | |
| p = kf["resolved_frame_index"] | |
| # faithful stock: computed from text_len, NOT the cursor | |
| if p == 0: | |
| t = float(text_len) | |
| elif frame_count is not None and p == frame_count - 1: | |
| t = float(text_len) + sum(spans) - FRAME_RESCALE | |
| else: | |
| raise ValueError( | |
| "only first/last keyframe anchors are supported") | |
| g = np.empty((frame_rows, 3), dtype=np.float64) | |
| g[:, 0] = t | |
| g[:, 1:] = frame | |
| emit("cond", g) | |
| cursor = float(text_len) | |
| for blk in (refs or []): | |
| kind = blk["kind"] | |
| if kind == "image": | |
| r_frame = _frame_grid(blk["latent_h"], blk["latent_w"]) | |
| g = np.empty((r_frame.shape[0], 3), dtype=np.float64) | |
| g[:, 0] = cursor | |
| g[:, 1:] = r_frame | |
| emit("ref_img", g) | |
| cursor += 1.0 | |
| elif kind == "audio": | |
| rt = int(blk["ref_audio_t"]) | |
| if rt > 0: | |
| emit("ref_audio", | |
| _audio_grid(cursor, rt, audio_rows_per_step)) | |
| cursor += float(rt) * ref_advance_factor | |
| elif kind in ("video", "video_audio"): | |
| rt = int(blk["ref_audio_t"]) | |
| vt = int(blk["latent_t"]) | |
| r_frame = _frame_grid(blk["latent_h"], blk["latent_w"]) | |
| if rt > 0: | |
| emit("ref_audio", | |
| _audio_grid(cursor, rt, audio_rows_per_step)) | |
| emit("ref_img", _video_grid(vt, r_frame, cursor)) | |
| cursor += max(float(rt), sum(_video_t_spans(vt))) | |
| else: | |
| raise ValueError("mock: unsupported ref kind %r" % kind) | |
| # target audio then target video, always the last two segments | |
| emit("audio", _audio_grid(cursor, audio_t)) | |
| emit("video", _video_grid(latent_t, frame, cursor)) | |
| seg_abs, off = [], 0 | |
| for kind, n in segs: | |
| seg_abs.append((off, off + n, kind)) | |
| off += n | |
| self.segments = seg_abs | |
| self.seq_len = off | |
| self.position_ids = np.concatenate(blocks) | |
| mm.PackedLayout = PackedLayout | |
| return mm | |
| def make_torch(): | |
| t = types.ModuleType("torch") | |
| t.equal = lambda a, b: a.shape == b.shape and bool(np.array_equal(a, b)) | |
| t.zeros = lambda shape: np.zeros(shape, dtype=np.float32) | |
| t.ones = lambda shape: np.ones(shape, dtype=np.float32) | |
| return t | |
| def load_patch(mm): | |
| for name in ("comfy", "comfy.ldm", "comfy.ldm.minimax"): | |
| sys.modules.setdefault(name, types.ModuleType(name)) | |
| sys.modules["comfy.ldm.minimax.model"] = mm | |
| sys.modules["comfy"].ldm = sys.modules["comfy.ldm"] | |
| sys.modules["comfy.ldm"].minimax = sys.modules["comfy.ldm.minimax"] | |
| sys.modules["comfy.ldm.minimax"].model = mm | |
| sys.modules["torch"] = make_torch() | |
| sys.modules.pop("patch_layout", None) | |
| return importlib.import_module("patch_layout") | |
| def load_patch_named(mm, name): | |
| for module_name in ("comfy", "comfy.ldm", "comfy.ldm.minimax"): | |
| sys.modules.setdefault(module_name, types.ModuleType(module_name)) | |
| sys.modules["comfy.ldm.minimax.model"] = mm | |
| sys.modules["comfy"].ldm = sys.modules["comfy.ldm"] | |
| sys.modules["comfy.ldm"].minimax = sys.modules["comfy.ldm.minimax"] | |
| sys.modules["comfy.ldm.minimax"].model = mm | |
| sys.modules["torch"] = make_torch() | |
| spec = importlib.util.spec_from_file_location( | |
| name, os.path.join(os.path.dirname(__file__), "..", "patch_layout.py")) | |
| module = importlib.util.module_from_spec(spec) | |
| sys.modules[name] = module | |
| spec.loader.exec_module(module) | |
| return module | |
| def install_structural_solattn(mm, module_name): | |
| """Install the defensive-getattr observer under an arbitrary name.""" | |
| namespace = {"_SPANS": {}} | |
| exec(compile( | |
| '''def _video_span(layout): | |
| return next(((a, b) for a, b, kind in layout.segments | |
| if kind == "video"), None) | |
| def install(layout_cls, module_name): | |
| original_init = layout_cls.__init__ | |
| def __init__(self, *args, **kwargs): | |
| original_init(self, *args, **kwargs) | |
| span = _video_span(self) | |
| bounds = next(((a, b) for a, b, kind in | |
| getattr(self, "segments", []) or [] | |
| if kind == "video"), None) | |
| if getattr(self, "position_ids", None) is not None and bounds is not None: | |
| _SPANS[id(self.position_ids)] = (self, bounds, span) | |
| __init__.__module__ = module_name | |
| layout_cls.__init__ = __init__ | |
| ''', "<structural-solattn-observer>", "exec"), namespace) | |
| namespace["install"](mm.PackedLayout, module_name) | |
| return namespace | |
| def install_partial_solattn_lookalike(mm, module_name): | |
| """Install a closure with only half the audited layout fingerprint.""" | |
| namespace = {"_SPANS": {}, "_video_span": lambda _layout: None} | |
| exec(compile( | |
| '''def install(layout_cls, module_name): | |
| original_init = layout_cls.__init__ | |
| def __init__(self, *args, **kwargs): | |
| original_init(self, *args, **kwargs) | |
| span = _video_span(self) | |
| position_ids = getattr(self, "position_ids", None) | |
| if position_ids is not None: | |
| _SPANS[id(position_ids)] = (self, span) | |
| __init__.__module__ = module_name | |
| layout_cls.__init__ = __init__ | |
| ''', "<partial-solattn-lookalike>", "exec"), namespace) | |
| namespace["install"](mm.PackedLayout, module_name) | |
| return namespace | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T = 7, 7, 22, 38, 16 | |
| FC = sum(FRAME_PER_TOKEN[k % 5] for k in range(LATENT_T)) | |
| def _cond_ts(layout): | |
| return [float(layout.position_ids[a, 0]) | |
| for a, _, k in layout.segments if k == "cond"] | |
| def main(): | |
| # 1. faithful stock: patch must apply | |
| mm = make_mm() | |
| pl = load_patch(mm) | |
| assert pl.apply_patch(), "self-test failed against faithful stock" | |
| assert pl.is_applied() | |
| print("1. self-test passes against the faithful stock layout") | |
| # 2. interior anchors build, reference compensation lands | |
| run = [{"resolved_frame_index": 0, pl.MC_KEY: i} for i in range(4)] | |
| lay = mm.PackedLayout(TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, frame_count=FC) | |
| ts = _cond_ts(lay) | |
| exp = [TEXT_LEN + FRAME_RESCALE * i for i in range(4)] | |
| assert np.allclose(ts, exp), (ts, exp) | |
| ref = [{"kind": "audio", "ref_audio_t": 8}] | |
| lay2 = mm.PackedLayout(TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=ref, frame_count=FC) | |
| ts2 = _cond_ts(lay2) | |
| assert np.allclose(ts2, [t + 8.0 for t in ts]), (ts, ts2) | |
| print("2. interior run at", [round(t, 4) for t in ts], | |
| "-> with an 8-step reference", [round(t, 4) for t in ts2]) | |
| # 3. Ref2VA: the graph's own references survive the audio move intact. | |
| # The marked block sits third of four, not last, because locating it | |
| # by segment does not depend on where in the list it is. | |
| rt, end_frame = 8, 22 | |
| head = [ | |
| {"kind": "image", "latent_h": 8, "latent_w": 12}, | |
| {"kind": "video_audio", "latent_h": 8, "latent_w": 12, | |
| "latent_t": 3, "ref_audio_t": 5}, | |
| ] | |
| tail = [{"kind": "audio", "ref_audio_t": 3}] | |
| marked = {"kind": "audio", "ref_audio_t": rt, pl.MC_AUDIO_KEY: end_frame} | |
| plain = {"kind": "audio", "ref_audio_t": rt} | |
| refs_plain = head + [plain] + tail | |
| refs_marked = head + [marked] + tail | |
| base = mm.PackedLayout(TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=refs_plain, frame_count=FC) | |
| after = mm.PackedLayout(TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=refs_marked, frame_count=FC) | |
| a, b = pl._ref_segment_map(base, refs_plain)[2]["ref_audio"] | |
| diff = set(i for i in range(base.position_ids.shape[0]) | |
| if float(base.position_ids[i, 0]) | |
| != float(after.position_ids[i, 0])) | |
| assert diff == set(range(a, b)), (len(diff), b - a) | |
| origin = pl._target_origin(after) | |
| got_end = float(after.position_ids[a, 0]) + rt | |
| assert abs(got_end - (origin + FRAME_RESCALE * end_frame)) < 1e-9 | |
| ts3 = _cond_ts(after) | |
| assert np.allclose( | |
| ts3, [origin + FRAME_RESCALE * i for i in range(4)]), ts3 | |
| print("3. marked block (3rd of 4 references) ends at target frame %d; " | |
| "the image, video and audio references either side of it did not " | |
| "move" % end_frame) | |
| # 4. mixed stock/MC keyframes under a reference are rejected | |
| mixed = [{"resolved_frame_index": 0}, | |
| {"resolved_frame_index": 0, pl.MC_KEY: 2}] | |
| try: | |
| mm.PackedLayout(TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=mixed, refs=ref, frame_count=FC) | |
| except RuntimeError as e: | |
| assert "mixed" in str(e) | |
| print("4. mixed stock/MC keyframes under a reference rejected loudly") | |
| else: | |
| raise AssertionError("mixed keyframes under a reference not rejected") | |
| # 5. an audio block with an unexpected row count must be caught, not | |
| # half-moved. The count is asserted exactly, so 3 rows per step trips. | |
| bad_rows = make_mm(audio_rows_per_step=3) | |
| pl_bad = load_patch(bad_rows) | |
| assert not pl_bad.apply_patch(), \ | |
| "self-test passed against a changed audio row count" | |
| assert bad_rows.PackedLayout.__init__.__name__ != "_patched_init" | |
| print("5. changed audio row count caught, patch refused") | |
| # 6. a changed reference cursor advance must NOT break anything. The | |
| # target origin is read back from the layout rather than recomputed, | |
| # so the anchors and the audio window follow it wherever it goes. The | |
| # previous implementation kept its own copy of that arithmetic and | |
| # had no choice but to refuse. | |
| drift = make_mm(ref_advance_factor=2.0) | |
| pl_drift = load_patch(drift) | |
| assert pl_drift.apply_patch(), \ | |
| "a changed reference advance should no longer defeat the patch" | |
| only = [dict(marked)] | |
| lay_d = drift.PackedLayout(TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=only, frame_count=FC) | |
| origin_d = pl_drift._target_origin(lay_d) | |
| ts_d = _cond_ts(lay_d) | |
| assert np.allclose( | |
| ts_d, [origin_d + FRAME_RESCALE * i for i in range(4)]), ts_d | |
| a_d, _ = pl_drift._ref_segment_map(lay_d, only)[0]["ref_audio"] | |
| end_d = float(lay_d.position_ids[a_d, 0]) + rt | |
| assert abs(end_d - (origin_d + FRAME_RESCALE * end_frame)) < 1e-9 | |
| print("6. doubled reference advance: anchors and audio window still " | |
| "land correctly, because the origin is read and not recomputed") | |
| # 7. two packs vendoring this file: the second must stand down, and | |
| # the first must still own the constructor and still work. | |
| mm7 = make_mm() | |
| first = load_patch(mm7) | |
| assert first.apply_patch() | |
| installed = mm7.PackedLayout.__init__ | |
| sys.modules.pop("patch_layout", None) | |
| second = importlib.import_module("patch_layout") | |
| assert second is not first, "second load returned the same module object" | |
| assert second.apply_patch(), \ | |
| "standing down should report success, the patch IS active" | |
| assert second.is_applied(), \ | |
| "is_applied must be true when another pack owns the patch, or the "\ | |
| "second pack's nodes would refuse to run" | |
| assert mm7.PackedLayout.__init__ is installed, \ | |
| "the second copy wrapped the first instead of standing down" | |
| assert second._orig_init is None, "the second copy captured an original" | |
| lay7 = mm7.PackedLayout(TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| origin7 = first._target_origin(lay7) | |
| ts7 = _cond_ts(lay7) | |
| assert np.allclose( | |
| ts7, [origin7 + FRAME_RESCALE * i for i in range(4)]), ts7 | |
| a7, _ = first._ref_segment_map(lay7, [marked])[0]["ref_audio"] | |
| end7 = float(lay7.position_ids[a7, 0]) + rt | |
| assert abs(end7 - (origin7 + FRAME_RESCALE * end_frame)) < 1e-9 | |
| print("7. second copy of the patch stands down, first still owns the " | |
| "constructor and still places anchors and audio correctly") | |
| # 8. a wrapper from a different copy carries no marker, only the | |
| # name. It must still be recognised, or the second copy self-tests | |
| # through the first one's behaviour and refuses over a limitation | |
| # that may not exist in its own code. | |
| mm8 = make_mm() | |
| other = load_patch(mm8) | |
| assert other.apply_patch() | |
| wrapper = mm8.PackedLayout.__init__ | |
| try: | |
| delattr(wrapper, other.PATCH_MARKER) # simulate an older copy | |
| except AttributeError: | |
| raise AssertionError("the wrapper was never marked") | |
| assert not getattr(wrapper, other.PATCH_MARKER, False) | |
| assert wrapper.__name__ == "_patched_init" | |
| sys.modules.pop("patch_layout", None) | |
| mine = importlib.import_module("patch_layout") | |
| assert mine._already_patched() == "other", mine._already_patched() | |
| assert mine.apply_patch(), "an unmarked copy must still be stood down for" | |
| assert mine.is_applied() | |
| assert mm8.PackedLayout.__init__ is wrapper, "wrapped a foreign copy" | |
| assert mine._orig_init is None | |
| print("8. an unmarked copy of the patch (older version or fork) is " | |
| "recognised by name and stood down for, not wrapped") | |
| # 9. an unrelated pack's wrapper: no marker, a different name, and | |
| # defined somewhere other than the module PackedLayout lives in. | |
| mm9 = make_mm() | |
| stock_init = mm9.PackedLayout.__init__ | |
| assert stock_init.__module__ == mm9.PackedLayout.__module__ | |
| def _some_other_pack_init(self, *args, **kwargs): | |
| stock_init(self, *args, **kwargs) | |
| _some_other_pack_init.__module__ = "some_other_h3_pack.patches" | |
| mm9.PackedLayout.__init__ = _some_other_pack_init | |
| ours = load_patch(mm9) | |
| assert ours._already_patched() == "foreign", ours._already_patched() | |
| assert not ours.apply_patch(), \ | |
| "stacked on an unrelated pack's wrapper instead of refusing" | |
| assert not ours.is_applied() | |
| assert mm9.PackedLayout.__init__ is _some_other_pack_init, \ | |
| "replaced another pack's wrapper" | |
| # functools.wraps copies __module__ and __name__ across, so the | |
| # module check alone would miss it; __wrapped__ is the giveaway | |
| import functools | |
| mm10 = make_mm() | |
| orig10 = mm10.PackedLayout.__init__ | |
| def _wrapped_init(self, *args, **kwargs): | |
| orig10(self, *args, **kwargs) | |
| mm10.PackedLayout.__init__ = _wrapped_init | |
| assert _wrapped_init.__module__ == mm10.PackedLayout.__module__ | |
| ours2 = load_patch(mm10) | |
| assert ours2._already_patched() == "foreign", ours2._already_patched() | |
| assert not ours2.apply_patch() | |
| print("9. an unrelated pack's wrapper is detected and refused, whether " | |
| "it hides behind functools.wraps or not") | |
| # 10. Kijai's SolAttn Morton hook wraps the constructor only to observe | |
| # and retain the final video span. Our position fix can safely wrap it: | |
| # SolAttn receives the same layout object and sees our in-place changes. | |
| mm11 = make_mm() | |
| stock11 = mm11.PackedLayout.__init__ | |
| observed11 = {} | |
| def install_solattn_first(): | |
| original_init = stock11 | |
| def __init__(self, *args, **kwargs): | |
| original_init(self, *args, **kwargs) | |
| observed11[id(self.position_ids)] = self | |
| __init__.__module__ = "ComfyUI-SolAttn_triton._morton_h3" | |
| mm11.PackedLayout.__init__ = __init__ | |
| install_solattn_first() | |
| solattn11 = mm11.PackedLayout.__init__ | |
| ours3 = load_patch(mm11) | |
| assert ours3._already_patched() == "solattn", ours3._already_patched() | |
| assert ours3.apply_patch(), "failed to compose outside SolAttn" | |
| assert mm11.PackedLayout.__init__ is not solattn11 | |
| lay11 = mm11.PackedLayout( | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| assert observed11[id(lay11.position_ids)] is lay11 | |
| origin11 = ours3._target_origin(lay11) | |
| assert np.allclose( | |
| _cond_ts(lay11), | |
| [origin11 + FRAME_RESCALE * i for i in range(4)]) | |
| print("10. SolAttn first: its observer and the interior-anchor patch " | |
| "compose and see the same corrected layout") | |
| # 11. Reverse order. This is also the case where another node pack owns | |
| # the shared H3 patch before SolAttn loads: look through the known wrapper | |
| # and stand down instead of wrapping both a second time. | |
| mm12 = make_mm() | |
| first12 = load_patch(mm12) | |
| assert first12.apply_patch() | |
| patched12 = mm12.PackedLayout.__init__ | |
| observed12 = {} | |
| def install_solattn_second(): | |
| original_init = patched12 | |
| def __init__(self, *args, **kwargs): | |
| original_init(self, *args, **kwargs) | |
| observed12[id(self.position_ids)] = self | |
| __init__.__module__ = "ComfyUI-SolAttn_triton._morton_h3" | |
| mm12.PackedLayout.__init__ = __init__ | |
| install_solattn_second() | |
| solattn12 = mm12.PackedLayout.__init__ | |
| sys.modules.pop("patch_layout", None) | |
| second12 = importlib.import_module("patch_layout") | |
| assert second12._already_patched() == "same", second12._already_patched() | |
| assert second12.apply_patch() | |
| assert second12.is_applied() | |
| assert mm12.PackedLayout.__init__ is solattn12 | |
| lay12 = mm12.PackedLayout( | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| assert observed12[id(lay12.position_ids)] is lay12 | |
| origin12 = first12._target_origin(lay12) | |
| assert np.allclose( | |
| _cond_ts(lay12), | |
| [origin12 + FRAME_RESCALE * i for i in range(4)]) | |
| print("11. interior-anchor patch first: SolAttn composes outside it and a " | |
| "second vendored copy stands down through the known wrapper") | |
| # 12. The CUDA PR117 checkout loads the same read-only observer from its | |
| # pack root, so its nested wrapper receives a path-valued module identity | |
| # instead of the upstream ._morton_h3 suffix. Exercise both load orders and | |
| # keep the exact original_init closure requirement. | |
| cuda_module = "/basedir/custom_nodes/ComfyUI-SolAttn-CUDA-PR117" | |
| mm_cuda_first = make_mm() | |
| stock_cuda_first = mm_cuda_first.PackedLayout.__init__ | |
| observed_cuda_first = {} | |
| def install_cuda_first(): | |
| original_init = stock_cuda_first | |
| def __init__(self, *args, **kwargs): | |
| original_init(self, *args, **kwargs) | |
| observed_cuda_first[id(self.position_ids)] = self | |
| __init__.__module__ = cuda_module | |
| mm_cuda_first.PackedLayout.__init__ = __init__ | |
| install_cuda_first() | |
| cuda_wrapper_first = mm_cuda_first.PackedLayout.__init__ | |
| ours_cuda_first = load_patch(mm_cuda_first) | |
| assert ours_cuda_first._already_patched() == "solattn", \ | |
| ours_cuda_first._already_patched() | |
| assert ours_cuda_first.apply_patch() | |
| lay_cuda_first = mm_cuda_first.PackedLayout( | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| assert observed_cuda_first[id(lay_cuda_first.position_ids)] is lay_cuda_first | |
| origin_cuda_first = ours_cuda_first._target_origin(lay_cuda_first) | |
| assert np.allclose( | |
| _cond_ts(lay_cuda_first), | |
| [origin_cuda_first + FRAME_RESCALE * i for i in range(4)]) | |
| assert mm_cuda_first.PackedLayout.__init__ is not cuda_wrapper_first | |
| mm_cuda_second = make_mm() | |
| first_cuda_second = load_patch(mm_cuda_second) | |
| assert first_cuda_second.apply_patch() | |
| patched_cuda_second = mm_cuda_second.PackedLayout.__init__ | |
| observed_cuda_second = {} | |
| def install_cuda_second(): | |
| original_init = patched_cuda_second | |
| def __init__(self, *args, **kwargs): | |
| original_init(self, *args, **kwargs) | |
| observed_cuda_second[id(self.position_ids)] = self | |
| __init__.__module__ = cuda_module | |
| mm_cuda_second.PackedLayout.__init__ = __init__ | |
| install_cuda_second() | |
| cuda_wrapper_second = mm_cuda_second.PackedLayout.__init__ | |
| sys.modules.pop("patch_layout", None) | |
| second_cuda_second = importlib.import_module("patch_layout") | |
| assert second_cuda_second._already_patched() == "same", \ | |
| second_cuda_second._already_patched() | |
| assert second_cuda_second.apply_patch() | |
| assert mm_cuda_second.PackedLayout.__init__ is cuda_wrapper_second | |
| lay_cuda_second = mm_cuda_second.PackedLayout( | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| assert observed_cuda_second[id(lay_cuda_second.position_ids)] is lay_cuda_second | |
| mm_cuda_bad = make_mm() | |
| bad_original = mm_cuda_bad.PackedLayout.__init__ | |
| def install_cuda_without_audited_closure(): | |
| captured_init = bad_original | |
| def __init__(self, *args, **kwargs): | |
| captured_init(self, *args, **kwargs) | |
| __init__.__module__ = cuda_module | |
| mm_cuda_bad.PackedLayout.__init__ = __init__ | |
| install_cuda_without_audited_closure() | |
| ours_cuda_bad = load_patch(mm_cuda_bad) | |
| assert ours_cuda_bad._already_patched() == "foreign" | |
| assert not ours_cuda_bad.apply_patch(), \ | |
| "trusted CUDA PR117 module identity without original_init closure" | |
| print("12. CUDA PR117's exact path-valued SolAttn observer composes in " | |
| "both load orders while malformed lookalikes remain refused") | |
| # 13. The same audited observer is often copied to a helper folder whose | |
| # name is chosen by the user. Recognise its behavior, not that name, and | |
| # look through multiple read-only copies rather than rejecting the stack. | |
| mm_struct = make_mm() | |
| observer_a = install_structural_solattn( | |
| mm_struct, | |
| r"C:\Comfy\custom_nodes\sol_attn_minimax_v2", | |
| ) | |
| observer_b = install_structural_solattn( | |
| mm_struct, | |
| "/workspace/custom_nodes/my-renamed-sol-attn-helper", | |
| ) | |
| ours_struct = load_patch(mm_struct) | |
| assert ours_struct._already_patched() == "solattn", \ | |
| ours_struct._already_patched() | |
| assert ours_struct.apply_patch() | |
| lay_struct = mm_struct.PackedLayout( | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| assert id(lay_struct.position_ids) in observer_a["_SPANS"] | |
| assert id(lay_struct.position_ids) in observer_b["_SPANS"] | |
| mm_partial = make_mm() | |
| install_partial_solattn_lookalike( | |
| mm_partial, | |
| "/workspace/custom_nodes/partial-sol-attn-lookalike", | |
| ) | |
| ours_partial = load_patch(mm_partial) | |
| assert ours_partial._already_patched() == "foreign" | |
| assert not ours_partial.apply_patch(), \ | |
| "a partial observer fingerprint was accepted" | |
| print("13. defensive-getattr and nested renamed SolAttn observers compose " | |
| "without a folder allowlist; a partial fingerprint remains refused") | |
| # 14. The explicit workflow node may promote the current pack over an | |
| # older compatible vendor. Unlike the same-name reload cases above, real | |
| # custom-node packages have distinct module names, so the active wrapper's | |
| # module remains addressable and exposes the stock constructor it captured. | |
| mm13 = make_mm() | |
| older13 = load_patch_named(mm13, "h3_layout_vendor_older") | |
| assert older13.apply_patch() | |
| newer13 = load_patch_named(mm13, "h3_layout_vendor_newer") | |
| assert newer13.apply_patch() | |
| assert mm13.PackedLayout.__init__ is older13._patched_init | |
| claimed, detail = newer13.claim_patch_ownership() | |
| assert claimed, detail | |
| assert mm13.PackedLayout.__init__ is newer13._patched_init | |
| lay13 = mm13.PackedLayout( | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| origin13 = newer13._target_origin(lay13) | |
| assert np.allclose( | |
| _cond_ts(lay13), | |
| [origin13 + FRAME_RESCALE * i for i in range(4)]) | |
| # Preserve SolAttn even when it loaded around the older owner: claim the | |
| # captured inner constructor instead of replacing the observer itself. | |
| observed14 = {} | |
| mm14 = make_mm() | |
| older14 = load_patch_named(mm14, "h3_layout_solattn_older") | |
| assert older14.apply_patch() | |
| captured14 = mm14.PackedLayout.__init__ | |
| def install_solattn_around_older(): | |
| original_init = captured14 | |
| def __init__(self, *args, **kwargs): | |
| original_init(self, *args, **kwargs) | |
| observed14[id(self.position_ids)] = self | |
| __init__.__module__ = "ComfyUI-SolAttn_triton._morton_h3" | |
| mm14.PackedLayout.__init__ = __init__ | |
| install_solattn_around_older() | |
| solattn14 = mm14.PackedLayout.__init__ | |
| newer14 = load_patch_named(mm14, "h3_layout_solattn_newer") | |
| assert newer14.apply_patch() | |
| claimed, detail = newer14.claim_patch_ownership() | |
| assert claimed, detail | |
| assert mm14.PackedLayout.__init__ is solattn14 | |
| assert newer14._solattn_wrapped_init(solattn14) is newer14._patched_init | |
| lay14 = mm14.PackedLayout( | |
| TEXT_LEN, LATENT_T, LH, LW, AUDIO_T, | |
| keyframes=run, refs=[dict(marked)], frame_count=FC) | |
| assert observed14[id(lay14.position_ids)] is lay14 | |
| print("14. explicit priority claims layout ownership from an older copy " | |
| "and preserves SolAttn when it wraps that owner") | |
| print("all checks passed") | |
| if __name__ == "__main__": | |
| main() | |