"""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__ ''', "", "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__ ''', "", "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__ @functools.wraps(orig10) 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()