forked from NikoDemon80/ComfyUI-H3-Motion-Context
-
Notifications
You must be signed in to change notification settings - Fork 29
Expand file tree
/
Copy pathpatch_layout.py
More file actions
987 lines (878 loc) · 44.1 KB
/
Copy pathpatch_layout.py
File metadata and controls
987 lines (878 loc) · 44.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
"""Place MiniMax H3 continuation guides on the target timeline.
Stock ComfyUI builds keyframe conditioning rows at one of two time
coordinates and rejects everything else:
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(_video_t_spans(latent_t)) - FRAME_RESCALE
else:
raise ValueError("only first/last keyframe anchors are supported")
Both branches are the same expression. Each video token spans
FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] and covers FRAME_PER_TOKEN[k % 5]
pixel frames, so the cumulative time at pixel frame p is exactly
FRAME_RESCALE * p, for every p. Substituting p = frame_count - 1
reproduces the second branch identically:
text_len + FRAME_RESCALE * (frame_count - 1)
== text_len + FRAME_RESCALE * frame_count - FRAME_RESCALE
== text_len + sum(_video_t_spans(latent_t)) - FRAME_RESCALE
So the general position is:
cond_t = text_len + FRAME_RESCALE * pixel_index
We do NOT rewrite the source of PackedLayout.__init__. Instead every
keyframe is handed to stock code with resolved_frame_index = 0, which is
always legal, and the real index rides along under MC_KEY. After the
stock constructor returns we rewrite the time column of each cond
segment's rows in position_ids. RoPE is built at forward time from
position_ids, so this lands before anything reads it.
That keeps the patch surface to one attribute we can verify rather than a
copy of a 90-line constructor that would rot on the next ComfyUI change.
ComfyUI PR #15439 adds native arbitrary-position video/audio guides, merges
guide and Ref2VA payloads, aligns guides to the target origin after reference
blocks, and drops the ``frame_count`` constructor argument. On that API this
module detects the native implementation and leaves ``PackedLayout`` entirely
core-owned. Everything below is retained only as a guarded fallback for older
ComfyUI builds.
"""
import inspect
import logging
import sys
import torch
import comfy.ldm.minimax.model as mm
MC_KEY = "motion_context_index"
MC_AUDIO_KEY = "motion_context_audio_end_frame"
# Marker set on our wrapper so a second copy of this file, vendored into
# another pack, can recognise it and stand down instead of wrapping it.
# Shared ABI across every pack that vendors this patch, exactly like
# MC_KEY and MC_AUDIO_KEY: rename it in all of them or in none.
PATCH_MARKER = "_h3_motion_context_layout_patch"
NATIVE_GUIDES_MARKER = "_h3_motion_context_native_guides"
_LOG = logging.getLogger("h3_motion_context")
_orig_init = None
_applied = False
# SolAttn's H3 Morton hook wraps PackedLayout only to record the final video
# span. It calls the constructor it captured, does not mutate the layout, and
# keeps the layout object alive so later in-place position fixes remain visible.
# That makes it safe to compose with this patch in either load order. Keep this
# recognition deliberately narrow: unknown wrappers still fail closed.
SOLATTN_LAYOUT_MODULE_SUFFIX = "._morton_h3"
# The CUDA PR117 development checkout currently installs its H3 observer from
# the pack root rather than ``_morton_h3`` and assigns the checkout path to the
# nested wrapper's ``__module__``. ComfyUI therefore reports the literal path
# below as PackedLayout.__init__'s owner. Keep this identity exact: it is safe
# to compose only after the same ``original_init`` closure check used for the
# upstream Kijai wrapper succeeds.
SOLATTN_CUDA_PR117_LAYOUT_MODULE = "ComfyUI-SolAttn-CUDA-PR117"
# Folder names are not an ABI. ComfyUI loads package roots under their full
# install path, and PR test helpers are commonly copied to names such as
# ``sol_attn_minimax_v2``. Recognise those copies by the upstream observer's
# read-only span-registration fingerprint instead of growing a folder-name
# allowlist forever. A future upstream marker is accepted as the clean ABI.
SOLATTN_LAYOUT_OBSERVER_MARKER = "_sol_attn_h3_layout_observer"
# The Kijai observer's fingerprint is a nested ``__init__`` in a module
# that owns two very specific globals: a callable ``_video_span`` that reads
# the target video span off the layout, and a dict ``_SPANS`` that records it
# by ``position_ids`` identity. The wrapper must reference BOTH those globals
# from its body -- that is what proves this specific module's observer is what
# is installed, not some unrelated wrapper that happens to live alongside them.
SOLATTN_LAYOUT_SIGNATURE_GLOBALS = frozenset({"_video_span", "_SPANS"})
# The wrapper must also touch both PackedLayout attributes used by the
# observer. Two access patterns appear in the wild and both are equally
# observe-only:
# * ``self.segments`` / ``self.position_ids`` -- attribute lookup, name lands
# in ``code.co_names`` (Kijai's upstream ``_morton_h3`` does this).
# * ``getattr(self, "segments", [])`` -- defensive lookup, name lands in
# ``code.co_consts`` as a string literal (comfy-kitchen PR #117's
# ``sol_attn_minimax_v2`` does this).
# Requiring both attributes via either access route accepts both styles while
# ensuring that an arbitrary ``__init__`` in the same module does not match.
SOLATTN_LAYOUT_ATTRIBUTE_NAMES = frozenset({"position_ids", "segments"})
REF_SEGMENT_KINDS = ("ref_img", "ref_audio")
def _target_origin(layout):
"""The coordinate the target clip starts at, read off the built layout.
Stock lays reference blocks out from a cursor that starts at text_len,
and the target rows take the cursor's final value as their origin.
Keyframe coordinates are computed from text_len directly and never
compensated, so without this term any reference slides the anchors
backwards relative to the clip they are anchoring.
Earlier versions recomputed that cursor with a local copy of stock's
per-kind advance arithmetic. Reading it back out of the layout instead
means there is nothing to keep in sync: if upstream changes how a
reference kind advances the cursor, the number here changes with it.
The target video segment is always last and always has at least one
latent step, and _video_grid puts its first row exactly on the cursor.
"""
a, b, kind = layout.segments[-1]
if kind != "video" or b <= a:
raise RuntimeError(
"h3_motion_context: expected the target video rows to be the "
"last layout segment, found %r spanning %d rows. Upstream "
"layout change; refusing to rewrite positions." % (kind, b - a))
return float(layout.position_ids[a, 0])
def _expected_ref_segments(blk):
"""The segment kinds one reference block emits, in emission order.
Mirrors the branches of the stock constructor:
image one ref_img
audio one ref_audio, or nothing at all when the window is
empty (stock skips the segment but still advances)
video the block's audio rows pack immediately before its
video_audio video rows, so ref_audio then ref_img
"""
kind = blk.get("kind")
if kind == "image":
return ("ref_img",)
if kind == "audio":
return ("ref_audio",) if int(blk.get("ref_audio_t", 0)) > 0 else ()
if kind in ("video", "video_audio"):
if int(blk.get("ref_audio_t", 0)) > 0:
return ("ref_audio", "ref_img")
return ("ref_img",)
raise RuntimeError(
"h3_motion_context: unknown reference kind %r; cannot tell which "
"layout rows belong to it." % (kind,))
def _ref_segment_map(layout, refs):
"""Which rows each reference block actually produced.
Returns {block_index: {segment_kind: (start, stop)}}.
This is the whole point of the multi-reference support. The rows of
one reference block could be found by working out the coordinate span
it ought to occupy and selecting everything inside it, but that means
duplicating stock's cursor arithmetic and then hoping nothing else
shares the range. Stock keyframe rows genuinely do land inside it,
which is why the coordinate approach needed an explicit exclusion for
them, and a second reference of the same kind would need another.
The layout already publishes a segment table, and reference blocks
emit their segments in list order, so the mapping is exact. Nothing
is inferred from coordinates and nothing has to be excluded.
"""
ref_segs = [(a, b, k) for a, b, k in layout.segments
if k in REF_SEGMENT_KINDS]
want = [(i, k) for i, blk in enumerate(refs or [])
for k in _expected_ref_segments(blk)]
if len(want) != len(ref_segs):
raise RuntimeError(
"h3_motion_context: %d reference blocks should have produced %d "
"layout segments, the layout has %d. Upstream layout change; "
"refusing to move rows." % (len(refs or []), len(want),
len(ref_segs)))
out = {}
for (i, kind), (a, b, got) in zip(want, ref_segs):
if got != kind:
raise RuntimeError(
"h3_motion_context: reference block %d (%r) should have "
"emitted a %s segment, the layout has %s. Upstream layout "
"change; refusing to move rows."
% (i, refs[i].get("kind"), kind, got))
out.setdefault(i, {})[kind] = (a, b)
return out
def _cond_t(text_len, latent_t, frame_count, p):
"""Time coordinate for a keyframe anchored at pixel frame p.
The endpoints reuse stock's exact expressions rather than the general
formula. They are mathematically identical, but stock accumulates
latent_t float additions where the general form does one multiply, and
those differ in the last bits (about 7e-15). Matching stock bit for bit
means an existing first/last graph builds byte-identical positions
after this patch is applied, and lets the self-test stay strict.
"""
if p == 0:
return float(text_len)
if frame_count is not None and p == frame_count - 1:
return float(text_len) + sum(mm._video_t_spans(latent_t)) - mm.FRAME_RESCALE
return float(text_len) + mm.FRAME_RESCALE * float(p)
def _fixup(layout, text_len, latent_t, frame_count, keyframes, refs=None):
"""Rewrite cond-row time coordinates to the general position formula.
`refs` is accepted but no longer read for arithmetic: the compensation
a reference block owes the anchors is now taken from where the target
actually landed, not recomputed from the block list.
"""
offset = _target_origin(layout) - float(text_len)
if offset and any(kf.get(MC_KEY) is None for kf in keyframes):
# keyframes without MC_KEY are left exactly as stock built them,
# which means they do NOT get the reference compensation. Mixing
# them with MC keyframes under a reference would slide the stock
# anchors relative to ours and to the target. Nothing produces
# this today; refuse loudly in case something ever does.
raise RuntimeError(
"h3_motion_context: stock and motion-context keyframes mixed in "
"one graph alongside a ref; their coordinates would disagree. "
"Give every keyframe a %s entry or remove the refs." % MC_KEY)
cond_spans = [(a, b) for a, b, kind in layout.segments if kind == "cond"]
if len(cond_spans) != len(keyframes):
raise RuntimeError(
"h3_motion_context: expected %d cond segments, layout has %d. "
"Refusing to rewrite positions."
% (len(keyframes), len(cond_spans)))
for (a, b), kf in zip(cond_spans, keyframes):
p = kf.get(MC_KEY)
if p is None:
continue
layout.position_ids[a:b, 0] = _cond_t(text_len, latent_t, frame_count, p) + offset
def _fixup_audio(layout, text_len, refs):
"""Move the marked audio ref's rows onto the target timeline.
References and keyframes carry identical row machinery; what makes the
model read a reference as "a separate clip to imitate" rather than
"this clip, continued" is that its coordinates sit in a span before
the target. That distinction decided continuation vs reproduction for
video, and seam analysis showed the audio reference producing
phase-unlocked imitation. So: keep the audio on the reference path for
construction and payload (rows built, latents filled, all stock code
untouched) and TRANSLATE its time coordinates so the window END lands
at target frame MC_AUDIO_KEY, the same instant the pinned video ends.
Translation, not per-row assignment: new = old + shift preserves
whatever intra-block structure stock built. Stock lays an audio
reference out channel-major, the same rt coordinates once per stereo
channel, and a uniform shift keeps that intact without this code
having to know about it.
The block keeps its place in the cursor, so the coordinates it vacates
are left empty. An audio window longer than the video window therefore
spills backwards into empty space rather than onto the text rows, so
the collision that made `before` mode fail for video does not arise.
Other reference blocks are untouched. A Ref2VA graph can carry its own
image, video and audio references and this moves only the one block
the node marked, wherever in the list it sits.
"""
marked = [i for i, r in enumerate(refs or [])
if r.get(MC_AUDIO_KEY) is not None]
if len(marked) != 1:
raise RuntimeError(
"h3_motion_context: audio timeline placement needs exactly one "
"reference marked with %s; the layout has %d references and %d "
"marked. If this appeared during startup, check for more than "
"one H3 Motion Context folder in custom_nodes."
% (MC_AUDIO_KEY, len(refs or []), len(marked)))
idx = marked[0]
blk = refs[idx]
if blk.get("kind") != "audio":
raise RuntimeError(
"h3_motion_context: %s set on a %r ref; only audio refs can be "
"moved onto the timeline." % (MC_AUDIO_KEY, blk.get("kind")))
rt = int(blk.get("ref_audio_t", 0))
if rt <= 0:
return
seg = _ref_segment_map(layout, refs).get(idx, {}).get("ref_audio")
if seg is None:
raise RuntimeError(
"h3_motion_context: the marked audio reference produced no "
"ref_audio segment. Upstream layout change; refusing to move "
"rows.")
a, b = seg
if b - a != 2 * rt:
# stock emits exactly rt rows per stereo channel. An exact count
# rather than a tolerance: if this ever changes, the intra-block
# structure a translation is preserving has changed with it.
raise RuntimeError(
"h3_motion_context: the marked audio reference has %d rows for "
"%d latent steps, expected %d (stereo, channel-major). Upstream "
"layout change; refusing to move rows." % (b - a, rt, 2 * rt))
target_origin = _target_origin(layout)
slot_start = float(layout.position_ids[a, 0])
end_frame = float(blk[MC_AUDIO_KEY])
# window end at target time FRAME_RESCALE * end_frame, width rt steps
desired_start = target_origin + mm.FRAME_RESCALE * end_frame - float(rt)
layout.position_ids[a:b, 0] = (layout.position_ids[a:b, 0]
+ (desired_start - slot_start))
def _patched_init(self, text_len, latent_t, latent_h, latent_w, audio_t,
keyframes=None, refs=None, frame_count=None):
_orig_init(self, text_len, latent_t, latent_h, latent_w, audio_t,
keyframes=keyframes, refs=refs, frame_count=frame_count)
has_mc_kf = bool(keyframes) and any(
kf.get(MC_KEY) is not None for kf in keyframes)
has_mc_audio = bool(refs) and any(
r.get(MC_AUDIO_KEY) is not None for r in refs)
if has_mc_kf:
_fixup(self, text_len, latent_t, frame_count, keyframes, refs)
if has_mc_audio:
_fixup_audio(self, text_len, refs)
# neither marked: stock graph, leave it exactly as built
def _self_test():
"""Prove the rewrite reproduces stock positions before committing.
Builds the two anchors stock code already supports, once the stock way
and once through our mechanism, and requires the position tensors to
match exactly. Then exercises the parts stock has no equivalent of:
interior anchors, reference compensation, the audio move, and the
same audio move inside a multi-reference Ref2VA layout. If ComfyUI
changes the position maths or the segment table underneath us this
fails and the patch is not applied.
"""
text_len, latent_t, lh, lw, audio_t = 7, 7, 22, 38, 16
frame_count = sum(mm.FRAME_PER_TOKEN[k % 5] for k in range(latent_t))
def build(keyframes=None, refs=None, fix=False, move=False):
lay = mm.PackedLayout.__new__(mm.PackedLayout)
_orig_init(lay, text_len, latent_t, lh, lw, audio_t,
keyframes=keyframes, refs=refs, frame_count=frame_count)
if fix:
_fixup(lay, text_len, latent_t, frame_count, keyframes, refs)
if move:
_fixup_audio(lay, text_len, refs)
return lay
def cond_ts(lay):
return [float(lay.position_ids[a, 0])
for a, _, k in lay.segments if k == "cond"]
# 1. the two anchors stock supports must come out bit-identical
stock_kf = [{"resolved_frame_index": 0},
{"resolved_frame_index": frame_count - 1}]
ours_kf = [{"resolved_frame_index": 0, MC_KEY: 0},
{"resolved_frame_index": 0, MC_KEY: frame_count - 1}]
a = build(keyframes=stock_kf)
b = build(keyframes=ours_kf, fix=True)
if a.position_ids.shape != b.position_ids.shape:
raise RuntimeError("position_ids shape mismatch in self-test")
if not torch.equal(a.position_ids, b.position_ids):
bad = (a.position_ids != b.position_ids).any(dim=1).nonzero().flatten()
raise RuntimeError("position mismatch at rows %s" % bad[:8].tolist())
# 2. a consecutive run lands on strictly increasing coordinates inside
# the span the two endpoints define
run = [{"resolved_frame_index": 0, MC_KEY: i} for i in range(4)]
c = build(keyframes=run, fix=True)
ts = cond_ts(c)
if len(ts) != len(run):
raise RuntimeError("expected %d cond segments, got %d" % (len(run), len(ts)))
if any(ts[i] >= ts[i + 1] for i in range(len(ts) - 1)):
raise RuntimeError("consecutive anchors not strictly increasing: %s" % ts)
t_last = float(text_len) + mm.FRAME_RESCALE * (frame_count - 1)
if not (ts[0] == float(text_len) and ts[-1] < t_last):
raise RuntimeError("run %s escapes the [%.4f, %.4f] span"
% (ts, float(text_len), t_last))
# 3. adding a reference must not move the anchors relative to the
# target. Stock cond rows cannot be the reference here: stock computes
# them from text_len and never compensates, which is the very bug the
# compensation exists to fix. The ground truth is the target rows
# themselves, so the anchor-to-end gap must be identical with and
# without the reference.
ref = [{"kind": "audio", "ref_audio_t": 8}]
d = build(keyframes=run, refs=ref, fix=True)
ts_ref = cond_ts(d)
if len(ts_ref) != len(ts):
raise RuntimeError("cond segment count changed when a ref was added")
tol = 1e-3
gap = float(c.position_ids[:, 0].max()) - ts[0]
gap_ref = float(d.position_ids[:, 0].max()) - ts_ref[0]
if abs(gap - gap_ref) > tol:
raise RuntimeError(
"ref compensation off by %.6f: anchor-to-target gap %.6f without "
"ref, %.6f with. The target origin read back from the layout no "
"longer matches its cursor arithmetic." % (gap_ref - gap, gap, gap_ref))
shifts = [y - x for x, y in zip(ts, ts_ref)]
if any(abs(sh - shifts[0]) > tol for sh in shifts):
raise RuntimeError("ref shifted anchors unevenly: %s" % shifts)
# 4. the audio move: exactly the marked block's rows shift, all by one
# uniform amount, every other row bit-identical
end_frame, rt = 4, 8
ref_mc = [{"kind": "audio", "ref_audio_t": rt, MC_AUDIO_KEY: end_frame}]
e = build(keyframes=run, refs=ref_mc, fix=True, move=True)
_check_move(d, e, ref_mc, 0, "single-ref")
# 5. the same move inside a Ref2VA layout: image, video and audio
# references of the graph's own must come through untouched. The
# marked block sits in the MIDDLE of the list, not at the end, because
# nothing about locating it by segment depends on its position.
r_lh, r_lw, r_vt = 8, 12, 3
others = [
{"kind": "image", "latent_h": r_lh, "latent_w": r_lw},
{"kind": "video_audio", "latent_h": r_lh, "latent_w": r_lw,
"latent_t": r_vt, "ref_audio_t": 5},
{"kind": "audio", "ref_audio_t": 3},
]
marked = {"kind": "audio", "ref_audio_t": rt, MC_AUDIO_KEY: end_frame}
plain = {"kind": "audio", "ref_audio_t": rt}
multi_plain = others[:2] + [plain] + others[2:]
multi_marked = others[:2] + [marked] + others[2:]
f = build(keyframes=run, refs=multi_plain, fix=True)
g = build(keyframes=run, refs=multi_marked, fix=True, move=True)
_check_move(f, g, multi_marked, 2, "multi-ref")
# 6. the segment map must agree with how the layout is actually laid
# out. Not by recomputing the cursor, which is the duplication this
# rewrite exists to remove, but by checking the structural properties
# the move depends on: reference blocks appear in list order, their
# rows sit before the target, and no block's rows overlap another's.
smap = _ref_segment_map(f, multi_plain)
prev_hi = float(text_len) - 1e-9
origin = _target_origin(f)
for i in range(len(multi_plain)):
spans = smap.get(i)
if not spans:
continue
rows = [r for a0, b0 in spans.values() for r in range(a0, b0)]
lo = min(float(f.position_ids[r, 0]) for r in rows)
hi = max(float(f.position_ids[r, 0]) for r in rows)
if lo < prev_hi - 1e-9:
raise RuntimeError(
"reference block %d starts at %.6f, before block %d ended "
"at %.6f. Reference blocks are not laid out in list order."
% (i, lo, i - 1, prev_hi))
if hi >= origin - 1e-9:
raise RuntimeError(
"reference block %d reaches %.6f, at or past the target "
"origin %.6f. Reference rows should sit before the target."
% (i, hi, origin))
prev_hi = hi
def _check_move(before, after, refs, idx, label):
"""Only the marked block's rows moved, uniformly, on the time axis."""
if after.position_ids.shape != before.position_ids.shape:
raise RuntimeError("%s: audio move changed the layout shape" % label)
if not torch.equal(before.position_ids[:, 1:], after.position_ids[:, 1:]):
raise RuntimeError(
"%s: audio move touched a non-time coordinate column" % label)
a, b = _ref_segment_map(before, refs)[idx]["ref_audio"]
expect_moved = set(range(a, b))
tb, ta = before.position_ids[:, 0], after.position_ids[:, 0]
moved = set(i for i in range(len(tb)) if float(tb[i]) != float(ta[i]))
if not moved:
raise RuntimeError("%s: audio move moved no rows" % label)
if moved != expect_moved:
raise RuntimeError(
"%s: audio move touched the wrong rows: %d moved, %d expected, "
"e.g. %s" % (label, len(moved), len(expect_moved),
sorted(moved ^ expect_moved)[:8]))
deltas = [float(ta[i]) - float(tb[i]) for i in sorted(moved)]
if any(abs(dd - deltas[0]) > 1e-9 for dd in deltas):
raise RuntimeError("%s: audio rows shifted non-uniformly: %s"
% (label, deltas[:4]))
# The size of the shift is deliberately NOT asserted. It depends on
# how far the reference cursor advanced, which is stock's business,
# and pinning it here would put a copy of that arithmetic back in.
# What must hold is where the window ENDS: on the target timeline,
# FRAME_RESCALE * end_frame past the target origin.
blk = refs[idx]
rt = int(blk["ref_audio_t"])
want_end = (_target_origin(after)
+ mm.FRAME_RESCALE * float(blk[MC_AUDIO_KEY]))
got_end = float(after.position_ids[a, 0]) + float(rt)
if abs(got_end - want_end) > 1e-9:
raise RuntimeError(
"%s: audio window ends at %.6f, should end at %.6f"
% (label, got_end, want_end))
setattr(_patched_init, PATCH_MARKER, True)
def _is_supported_solattn_layout_module(init):
"""Whether ``init`` comes from one of the audited SolAttn observers."""
module = str(getattr(init, "__module__", "") or "")
if module.endswith(SOLATTN_LAYOUT_MODULE_SUFFIX):
return True
normalized = module.replace("\\", "/").rstrip("/")
tail = normalized.rsplit("/", 1)[-1]
return (tail == SOLATTN_CUDA_PR117_LAYOUT_MODULE
or tail.endswith("." + SOLATTN_CUDA_PR117_LAYOUT_MODULE))
def _is_structural_solattn_layout_observer(init):
"""Recognise a Kijai H3 layout observer independently of its folder.
The observer's fingerprint is a nested ``__init__`` in a module that owns
two very specific globals -- a callable ``_video_span`` that reads the
target span off the layout, and a dict ``_SPANS`` that records it by
``position_ids`` identity -- and the wrapper's body references both. It
must also touch both PackedLayout attributes. Each touch can be a direct
attribute lookup (``self.segments``, in which case ``segments`` is in
``co_names``) or a defensive getattr with a string literal
(``getattr(self, "segments", [])``, in which case ``"segments"`` is in
``co_consts``). Both patterns appear in the wild -- Kijai's upstream
``_morton_h3`` uses direct access, comfy-kitchen PR #117's
``sol_attn_minimax_v2`` uses getattr -- and both are equally observe-only.
Requiring the complete attribute pair keeps the match fail-closed without
imposing a particular access style.
"""
if bool(getattr(init, SOLATTN_LAYOUT_OBSERVER_MARKER, False)):
return True
code = getattr(init, "__code__", None)
namespace = getattr(init, "__globals__", None)
if code is None or code.co_name != "__init__" or not isinstance(
namespace, dict):
return False
# Signature globals: two module-level objects with matching types. These
# names are unique to the Kijai observer design; any module that
# ships both a callable ``_video_span`` and a dict ``_SPANS`` is one of
# the audited observer variants.
if not (callable(namespace.get("_video_span"))
and isinstance(namespace.get("_SPANS"), dict)):
return False
names = set(code.co_names)
# The wrapper must actually reference the signature globals from its body,
# not merely inherit them from the module. This rules out unrelated
# ``__init__`` methods that happen to live in the same module.
if not SOLATTN_LAYOUT_SIGNATURE_GLOBALS.issubset(names):
return False
# And it must touch both PackedLayout attributes we care about.
# Attribute access puts the name in ``co_names``; ``getattr(self, "name",
# default)`` puts the name in ``co_consts`` as a string literal. Accept
# either route -- a real observer references both under one of the two
# patterns.
consts = {c for c in getattr(code, "co_consts", ()) if isinstance(c, str)}
if not SOLATTN_LAYOUT_ATTRIBUTE_NAMES.issubset(names | consts):
return False
return True
def _is_supported_solattn_layout_observer(init):
return bool(
_is_supported_solattn_layout_module(init)
or _is_structural_solattn_layout_observer(init)
)
def _solattn_wrapped_init(init):
"""Return the constructor captured by Kijai's H3 Morton wrapper.
The upstream wrapper is a nested ``__init__`` in ``_patch_packed_layout``
and captures the previous constructor in a closure named
``original_init``. Checking that exact closure together with either an
audited module identity or the read-only span-registration fingerprint is
intentionally stricter than accepting an arbitrary renamed wrapper.
"""
if not _is_supported_solattn_layout_observer(init):
return None
closure = getattr(init, "__closure__", None) or ()
freevars = getattr(getattr(init, "__code__", None), "co_freevars", ())
for name, cell in zip(freevars, closure):
if name != "original_init":
continue
try:
original = cell.cell_contents
except ValueError:
return None
return original if callable(original) and original is not init else None
wrapped = getattr(init, "__wrapped__", None)
if bool(getattr(init, SOLATTN_LAYOUT_OBSERVER_MARKER, False)):
return wrapped if callable(wrapped) and wrapped is not init else None
return None
def _unwrap_solattn_layout_observers(init):
"""Return the constructor beneath every audited observer in ``init``.
Multiple renamed copies can be installed in one process. They are all
read-only, so unwrapping the full chain is safe; cycles and excessive
wrapper depth are treated as unknown rather than guessed through.
"""
current = init
wrappers = []
seen = set()
for _depth in range(16):
if id(current) in seen:
return None, tuple(wrappers)
seen.add(id(current))
original = _solattn_wrapped_init(current)
if original is None:
return current, tuple(wrappers)
wrappers.append(current)
current = original
return None, tuple(wrappers)
def _replace_solattn_wrapped_init(init, replacement):
"""Replace the constructor captured by the innermost audited observer."""
current = init
seen = set()
for _depth in range(16):
if id(current) in seen or not _is_supported_solattn_layout_observer(
current):
return False
seen.add(id(current))
closure = getattr(current, "__closure__", None) or ()
freevars = getattr(
getattr(current, "__code__", None), "co_freevars", ())
for name, cell in zip(freevars, closure):
if name != "original_init":
continue
try:
original = cell.cell_contents
except ValueError:
return False
if _solattn_wrapped_init(original) is not None:
current = original
break
cell.cell_contents = replacement
return True
else:
return False
return False
def _uses_native_guide_api(init):
"""Whether a constructor exposes PR #15439's frame_count-free API."""
original, _wrappers = _unwrap_solattn_layout_observers(init)
if original is None:
return False
if getattr(original, NATIVE_GUIDES_MARKER, False):
return True
try:
parameters = inspect.signature(original).parameters
except (TypeError, ValueError):
return False
return "frame_count" not in parameters
def native_guides_available():
"""Whether ComfyUI provides the merged, core-owned H3 guide API.
SolAttn may wrap the core constructor as a read-only observer, and an older
installed copy of this pack may already have wrapped the native API. The
narrow unwrapping/marker checks in ``_uses_native_guide_api`` recognise
both cases without changing either implementation.
"""
cls = getattr(mm, "PackedLayout", None)
init = getattr(cls, "__init__", None) if cls is not None else None
if init is None:
return False
original, _wrappers = _unwrap_solattn_layout_observers(init)
if original is None:
return False
if (getattr(original, NATIVE_GUIDES_MARKER, False)
or getattr(original, "__name__", "") == "_patched_init_native"):
return True
# A frame_count-free signature is meaningful only on core's own
# constructor. Arbitrary third-party wrappers often accept *args/**kwargs;
# treating those as the native API would silently bypass collision checks.
if hasattr(original, "__wrapped__"):
return False
home = str(getattr(cls, "__module__", "") or "")
where = str(getattr(original, "__module__", "") or "")
if home and where != home:
return False
return _uses_native_guide_api(original)
def compatibility_report():
"""Inspect guide-layout compatibility without installing any patch.
Preflight and authoring UIs call this before model execution. It must
remain read-only: no constructor replacement, self-test, or global owner
claim is performed here.
"""
cls = getattr(mm, "PackedLayout", None)
init = getattr(cls, "__init__", None) if cls is not None else None
if init is None:
return {
"ok": False,
"mode": "unavailable",
"owner": "",
"message": "MiniMax H3 PackedLayout is unavailable.",
"action": "Install or update a ComfyUI build with MiniMax H3 support.",
}
owner = str(getattr(init, "__module__", "") or "")
if native_guides_available():
return {
"ok": True,
"mode": "native",
"owner": owner,
"message": "ComfyUI's native H3 guide API is available.",
"action": "",
}
who = _already_patched()
if who == "foreign":
return {
"ok": False,
"mode": "conflict",
"owner": owner,
"message": (
"Another extension owns MiniMax H3 PackedLayout with an "
"unsupported wrapper."),
"action": (
"Disable the conflicting H3 layout extension or update "
"ComfyUI to the native H3 guide API, then restart ComfyUI."),
}
if who == "solattn":
return {
"ok": True,
"mode": "legacy_solattn_compatible",
"owner": owner,
"message": (
"The supported CUDA SolAttn layout observer can compose with "
"the legacy H3 guide bridge."),
"action": "",
}
if who in ("same", "other"):
return {
"ok": True,
"mode": "legacy_active",
"owner": owner,
"message": "A compatible legacy H3 guide bridge is already active.",
"action": "",
}
return {
"ok": True,
"mode": "legacy_available",
"owner": owner,
"message": "The legacy H3 guide bridge can be installed at runtime.",
"action": "Update ComfyUI to use the native H3 guide API when practical.",
}
def _already_patched():
"""Has another copy of this file already wrapped the constructor?
Returns None, "same", "other", "solattn", or "foreign".
Two copies of this patch in one ComfyUI is normal enough: several
packs vendor it, and forks of this repo carry their own. Whichever
loads second would otherwise capture the first's wrapper as its
original and wrap a wrapper. That is worse than it sounds, because
each copy self-tests through whatever is already installed, so a copy
with newer tests gets checked against older behaviour and refuses
over a limitation that no longer exists.
Three checks, in decreasing confidence.
The marker is set by copies new enough to set it. That is a matching
version and we stand down quietly.
A wrapper merely NAMED like ours is an older copy of this code, or a
fork. We stand down and say so, because whichever one loaded first is
the one deciding what the patch supports.
SolAttn's narrowly identified Morton wrapper is safe to compose: it only
observes the layout and records its video span. If it captured a copy of
this patch, we stand down for that copy; if it captured stock, we install
outside it so both features remain active.
Anything else sitting where the stock constructor should be is a
DIFFERENT pack patching the same thing. Several H3 packs lift the
same first/last restriction independently, and they cannot both own
the constructor. Detected by comparing where the function was defined
against where the class was: stock's __init__ comes from the same
module as PackedLayout itself, a wrapper comes from somewhere else.
functools.wraps copies __module__ across, so __wrapped__ is checked
too. A wrapper that hides both is indistinguishable from stock and
nothing can be done about that.
"""
cls = getattr(mm, "PackedLayout", None)
init = getattr(cls, "__init__", None)
if init is None:
return None
if getattr(init, PATCH_MARKER, False):
return "same"
if getattr(init, "__name__", "") in (
"_patched_init", "_patched_init_native"):
return "other"
solattn_original, solattn_wrappers = _unwrap_solattn_layout_observers(init)
if solattn_wrappers:
if solattn_original is None:
return "foreign"
if getattr(solattn_original, PATCH_MARKER, False):
return "same"
if getattr(solattn_original, "__name__", "") in (
"_patched_init", "_patched_init_native"):
return "other"
home = getattr(cls, "__module__", None)
where = getattr(solattn_original, "__module__", None)
if home and where and where == home and not hasattr(
solattn_original, "__wrapped__"):
return "solattn"
return "foreign"
if hasattr(init, "__wrapped__"):
return "foreign"
home = getattr(cls, "__module__", None)
where = getattr(init, "__module__", None)
if home and where and where != home:
return "foreign"
return None
def apply_patch():
global _orig_init, _applied
if _applied:
return True
current_init = getattr(getattr(mm, "PackedLayout", None), "__init__", None)
if native_guides_available():
_LOG.info(
"h3_motion_context: native H3 guides detected; PackedLayout "
"remains core-owned")
return True
who = _already_patched()
if who == "foreign":
_LOG.warning(
"h3_motion_context: another pack has already patched H3's "
"layout (PackedLayout.__init__ now comes from %r). Several "
"ComfyUI packs lift the same first/last keyframe restriction "
"independently and they cannot both own it, so this one is "
"refusing rather than wrapping a wrapper and producing joins "
"neither pack intended. Disable one of them and restart.",
getattr(getattr(getattr(mm, "PackedLayout", None), "__init__",
None), "__module__", "?"))
return False
if who and who != "solattn":
# report success: the patch IS active, just not ours, and the
# calling pack's nodes check is_applied() before they will run
_applied = True
if who == "same":
_LOG.info(
"h3_motion_context: interior keyframe anchors already "
"enabled by another pack, standing down")
else:
_LOG.warning(
"h3_motion_context: the H3 layout patch is already installed "
"by a DIFFERENT copy of this code (another version, or a "
"fork). Standing down; that copy decides what the patch "
"supports, so features added since it may be unavailable. "
"If you have more than one H3 Motion Context folder in "
"custom_nodes, keep one and remove the rest. Renaming a "
"folder does not stop ComfyUI loading it.")
return True
if who == "solattn":
_LOG.info(
"h3_motion_context: composing with SolAttn's H3 Morton layout "
"observer; interior anchors and Morton ordering remain enabled")
if not hasattr(mm, "PackedLayout") or not hasattr(mm, "FRAME_RESCALE"):
_LOG.warning("h3_motion_context: MiniMax H3 model module missing expected "
"attributes, patch not applied")
return False
_orig_init = mm.PackedLayout.__init__
try:
_self_test()
except Exception as exc:
_orig_init = None
_LOG.warning("h3_motion_context: self-test failed (%s), patch not "
"applied. Continuation guide placement unavailable.", exc)
_LOG.warning(
"h3_motion_context: if you have more than one H3 Motion Context "
"folder in custom_nodes (a fork, a backup, a manual clone "
"alongside a Manager install), that is the usual cause: each "
"copy self-tests against whichever one loaded first. Keep one "
"and remove the rest. Renaming a folder does not stop ComfyUI "
"loading it. Otherwise this is an upstream ComfyUI change and "
"the message above says what moved.")
return False
mm.PackedLayout.__init__ = _patched_init
_applied = True
_LOG.info("h3_motion_context: legacy interior keyframe anchors enabled")
return True
def claim_patch_ownership():
"""Prefer this copy over an older compatible copy of the same patch.
The operation is explicit because it changes process-global ownership.
Only wrappers carrying this patch family's marker, or its exact historical
function names, are replaceable. SolAttn's narrowly recognised observer is
retained around the new owner; unknown wrappers fail closed.
Returns ``(ok, detail)`` for the visible Patch Priority pass-through node.
"""
global _orig_init, _applied
cls = getattr(mm, "PackedLayout", None)
current = getattr(cls, "__init__", None) if cls is not None else None
if current is None or not hasattr(mm, "FRAME_RESCALE"):
return False, "MiniMax H3 PackedLayout is unavailable"
if native_guides_available():
return True, "native guides are core-owned; no patch required"
if current is _patched_init:
_applied = True
return True, "layout owned by this pack"
who = _already_patched()
if who in (None, "solattn"):
# No sibling owns the patch yet. Normal activation already has all
# compatibility checks needed for stock or SolAttn-first load order.
_applied = False
if not apply_patch():
return False, "layout activation failed its compatibility test"
return True, ("layout activated by this pack"
if who is None else
"layout activated while retaining SolAttn")
if who == "foreign":
return False, (
"layout owner is unknown; only another H3 Motion Context copy "
"can be safely replaced")
family_wrapper, solattn_wrappers = _unwrap_solattn_layout_observers(
current)
if family_wrapper is None:
return False, "SolAttn observer chain is cyclic or unexpectedly deep"
owner_module = sys.modules.get(str(getattr(
family_wrapper, "__module__", "")))
original = getattr(owner_module, "_orig_init", None)
if not callable(original) or original is family_wrapper:
return False, (
"the existing H3 Motion Context layout wrapper does not expose "
"its captured constructor; restart with only one copy enabled")
stock, _original_observers = _unwrap_solattn_layout_observers(original)
if stock is None:
return False, "captured SolAttn observer chain is invalid"
home = str(getattr(cls, "__module__", "") or "")
where = str(getattr(stock, "__module__", "") or "")
if (hasattr(stock, "__wrapped__") or (home and where != home)
or getattr(stock, PATCH_MARKER, False)):
return False, (
"the existing layout wrapper captured another unknown wrapper; "
"refusing to discard it")
previous_original = _orig_init
_orig_init = original
try:
_self_test()
except Exception as exc:
_orig_init = previous_original
return False, "replacement layout self-test failed: %s" % exc
replacement = _patched_init
if solattn_wrappers:
if not _replace_solattn_wrapped_init(current, replacement):
_orig_init = previous_original
return False, "could not preserve SolAttn's layout observer"
else:
cls.__init__ = replacement
_applied = True
_LOG.info(
"h3_motion_context: this pack claimed H3 layout ownership from "
"compatible module %s", getattr(family_wrapper, "__module__", "?"))
return True, "layout ownership claimed from a compatible older copy"
def is_applied():
return _applied
def native_guides_active():
"""Backward-compatible alias for native guide capability detection."""
return native_guides_available()