File size: 37,310 Bytes
f2b53cf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import { chainCallback, addMiddleClickPan, addWheelPassthrough } from './utility.js';
const { app } = window.comfyAPI.app;

// --- begin context window math (port of comfy/context_windows.py) ---
const MAX_WINDOWS = 512;
const TRUNC_MSG = `too many windows, truncated at ${MAX_WINDOWS}`;
function capReached(windows, errors) {
    if (windows.length < MAX_WINDOWS) return false;
    errors.push(TRUNC_MSG);
    return true;
}

function pythonRound(x) {
    // python round() is half-even
    const floor = Math.floor(x);
    const diff = x - floor;
    if (diff > 0.5) return floor + 1;
    if (diff < 0.5) return floor;
    return floor % 2 === 0 ? floor : floor + 1;
}

function orderedHalving(val) {
    const bin = BigInt(val).toString(2).padStart(64, '0').split('').reverse().join('');
    return Number(BigInt('0b' + bin)) / 18446744073709551616;
}

function range(start, stop) {
    const out = [];
    for (let i = start; i < stop; i++) out.push(i);
    return out;
}

// loop min/max; avoids Math.min(...arr) spread-limit crash on long windows
function arrMin(a) {
    let m = Infinity;
    for (let i = 0; i < a.length; i++) if (a[i] < m) m = a[i];
    return m;
}
function arrMax(a) {
    let m = -Infinity;
    for (let i = 0; i < a.length; i++) if (a[i] > m) m = a[i];
    return m;
}

function arraysEqual(a, b) {
    if (a.length !== b.length) return false;
    for (let i = 0; i < a.length; i++) if (a[i] !== b[i]) return false;
    return true;
}

function doesWindowRollOver(window, numFrames) {
    let prev = -1;
    for (let i = 0; i < window.length; i++) {
        const val = ((window[i] % numFrames) + numFrames) % numFrames;
        if (val < prev) return [true, i];
        prev = val;
    }
    return [false, -1];
}

function shiftWindowToStart(window, numFrames) {
    const startVal = window[0];
    for (let i = 0; i < window.length; i++) {
        window[i] = ((window[i] - startVal) + numFrames) % numFrames;
    }
}

function shiftWindowToEnd(window, numFrames) {
    shiftWindowToStart(window, numFrames);
    const endDelta = numFrames - window[window.length - 1] - 1;
    for (let i = 0; i < window.length; i++) window[i] += endDelta;
}

// p: {length, overlap, stride, closedLoop, step}
function createWindows(numFrames, schedule, p, errors) {
    const L = p.length;
    if (schedule === 'standard_static') {
        if (numFrames <= L) return [range(0, numFrames)];
        const delta = L - p.overlap;
        if (delta <= 0) {
            errors.push(`context_overlap (${p.overlap}) must be smaller than context_length (${L}); the sampler would error`);
            return [];
        }
        const windows = [];
        for (let startIdx = 0; startIdx < numFrames; startIdx += delta) {
            const ending = startIdx + L;
            if (ending >= numFrames) {
                const finalStart = startIdx - (ending - numFrames);
                windows.push(range(finalStart, finalStart + L));
                break;
            }
            windows.push(range(startIdx, startIdx + L));
            if (capReached(windows, errors)) break;
        }
        return windows;
    }
    if (schedule === 'standard_static_balanced') {
        // standard_static count, placed evenly so seams are uniform (±1 frame)
        if (numFrames <= L) return [range(0, numFrames)];
        if (L - p.overlap <= 0) errors.push(`context_overlap (${p.overlap}) >= context_length (${L}); windows collapse to stride 1`);
        const strideTarget = Math.max(1, L - p.overlap);
        let N = Math.ceil((numFrames - L) / strideTarget) + 1; // stride <= L => no gaps
        if (N > MAX_WINDOWS) {
            errors.push(TRUNC_MSG);
            N = MAX_WINDOWS;
        }
        const windows = [];
        for (let i = 0; i < N; i++) {
            const start = Math.floor(i * (numFrames - L) / (N - 1) + 0.5);
            windows.push(range(start, start + L));
        }
        return windows;
    }
    if (schedule === 'batched') {
        if (numFrames <= L) return [range(0, numFrames)];
        const windows = [];
        for (let startIdx = 0; startIdx < numFrames; startIdx += L) {
            windows.push(range(startIdx, Math.min(startIdx + L, numFrames)));
            if (capReached(windows, errors)) break;
        }
        return windows;
    }
    if (schedule === 'batched_shifted') {
        // zero-overlap windows whose seams roll each step (~1.0x compute)
        if (numFrames <= L) return [range(0, numFrames)];
        const offset = pythonRound(numFrames * orderedHalving(p.step)) % L; // 0 at step 0 (plain batched)
        const windows = [];
        let start = 0;
        if (offset > 0) { windows.push(range(0, offset)); start = offset; }
        while (start < numFrames) {
            const end = Math.min(start + L, numFrames);
            windows.push(range(start, end));
            start = end;
            if (capReached(windows, errors)) break;
        }
        return windows;
    }
    if (schedule === 'looped_uniform' || schedule === 'standard_uniform') {
        const looped = schedule === 'looped_uniform';
        if (looped ? numFrames < L : numFrames <= L) return [range(0, numFrames)];
        const windows = [];
        const strideCount = Math.min(p.stride, Math.ceil(Math.log2(numFrames / L)) + 1);
        for (let sPow = 0; sPow < strideCount; sPow++) {
            const contextStep = 1 << sPow;
            const pad = pythonRound(numFrames * orderedHalving(p.step));
            const start = Math.floor(orderedHalving(p.step) * contextStep) + pad;
            const stop = numFrames + pad + (looped && p.closedLoop ? 0 : -p.overlap);
            const stepSize = L * contextStep - p.overlap;
            if (stepSize <= 0) {
                errors.push(`window advance is ${stepSize} (≤ 0) at stride ${contextStep}; the sampler would hang/error`);
                continue;
            }
            for (let j = start; j < stop; j += stepSize) {
                const w = [];
                for (let e = j; e < j + L * contextStep; e += contextStep) {
                    w.push(((e % numFrames) + numFrames) % numFrames);
                }
                windows.push(w);
                if (capReached(windows, errors)) return windows;
            }
        }
        if (!looped) {
            // shift roll-over windows to the end, fill gaps, drop duplicates
            const deleteIdxs = [];
            let winI = 0;
            while (winI < windows.length) {
                const [isRoll, rollIdx] = doesWindowRollOver(windows[winI], numFrames);
                if (isRoll) {
                    const rollVal = windows[winI][rollIdx];
                    shiftWindowToEnd(windows[winI], numFrames);
                    if (!windows[(winI + 1) % windows.length].includes(rollVal)) {
                        windows.splice(winI + 1, 0, range(rollVal, rollVal + L));
                    }
                }
                for (let preI = 0; preI < winI; preI++) {
                    if (arraysEqual(windows[winI], windows[preI])) {
                        deleteIdxs.push(winI);
                        break;
                    }
                }
                winI++;
            }
            deleteIdxs.reverse();
            for (const i of deleteIdxs) windows.splice(i, 1);
        }
        return windows;
    }
    return [range(0, numFrames)];
}

// extra overlap on the last static window from its shove-back; 0 = even tiling
function staticFinalDelta(F, L, overlap) {
    if (F <= L) return 0;
    const d = L - overlap;
    if (d <= 0) return 0; // invalid config, flagged elsewhere
    const mod = (F - L) % d;
    return mod === 0 ? 0 : d - mod;
}

function linspace(a, b, n) {
    if (n === 1) return [a];
    const out = [];
    for (let i = 0; i < n; i++) out.push(a + (b - a) * i / (n - 1));
    return out;
}

function weightsPyramid(length) {
    const seq = [];
    if (length % 2 === 0) {
        const m = length / 2;
        for (let i = 1; i <= m; i++) seq.push(i);
        for (let i = m; i >= 1; i--) seq.push(i);
    } else {
        const m = (length + 1) / 2;
        for (let i = 1; i < m; i++) seq.push(i);
        seq.push(m);
        for (let i = m - 1; i >= 1; i--) seq.push(i);
    }
    return seq;
}

function weightsOverlapLinear(length, fullLength, idxs, overlap, errors) {
    const w = new Array(length).fill(1);
    if (overlap > length && errors) {
        errors.push(`overlap-linear with context_overlap > context_length would error in the sampler`);
    }
    const ov = Math.min(overlap, length);
    if (ov > 0) {
        if (arrMin(idxs) > 0) {
            const ramp = linspace(1e-37, 1, ov);
            for (let i = 0; i < ov; i++) w[i] = ramp[i];
        }
        if (arrMax(idxs) < fullLength - 1) {
            const ramp = linspace(1, 1e-37, ov);
            for (let i = 0; i < ov; i++) w[length - ov + i] = ramp[i];
        }
    }
    return w;
}

function weightsRelative(idxList) {
    // per-index triangular bias (sequential weighted avg == bias/Σbias share)
    const first = idxList[0];
    const last = idxList[idxList.length - 1];
    return idxList.map((idx) => {
        const bias = 1 - Math.abs(idx - (first + last) / 2) / ((last - first + 1e-2) / 2);
        return Math.max(1e-2, bias);
    });
}

// smooth overlap-linear: flat interior, raised-cosine ramps that sum to 1
function weightsHann(length, fullLength, idxs, overlap, errors) {
    const w = new Array(length).fill(1);
    if (overlap > length && errors) {
        errors.push(`hann with context_overlap (${overlap}) > context_length (${length}); ramp clamped`);
    }
    const ov = Math.min(Math.max(overlap, 0), length);
    if (ov >= 1) {
        const denom = Math.max(ov - 1, 1);
        if (arrMin(idxs) > 0) {
            for (let j = 0; j < ov; j++) w[j] = Math.max(0.5 * (1 - Math.cos(Math.PI * j / denom)), 1e-37);
        }
        if (arrMax(idxs) < fullLength - 1) {
            for (let j = 0; j < ov; j++) w[length - ov + j] = Math.max(0.5 * (1 - Math.cos(Math.PI * (denom - j) / denom)), 1e-37);
        }
    }
    return w;
}

// smooth bell over the window (smooth twin of pyramid)
function weightsGaussian(length) {
    if (length <= 1) return new Array(length).fill(1);
    const c = (length - 1) / 2;
    const std = Math.max(length / 4, 1e-6);
    const out = [];
    for (let i = 0; i < length; i++) out.push(Math.exp(-0.5 * ((i - c) / std) ** 2));
    return out;
}

function computeViz(p) {
    const errors = [];
    const numFrames = Math.max(1, p.numFrames);
    const rawWindows = createWindows(numFrames, p.schedule, {
        length: Math.max(1, p.length),
        overlap: Math.max(0, p.overlap),
        stride: Math.max(1, p.stride),
        closedLoop: p.closedLoop,
        step: Math.max(0, p.step),
    }, errors);

    const windows = rawWindows.map((idxList) => {
        let anchor = -1;
        if (p.causalFix) {
            const a = idxList[0] - 1;
            if (a >= 0 && a < numFrames) anchor = a;
        }
        let weights;
        if (p.fuse === 'relative') weights = weightsRelative(idxList);
        else if (p.fuse === 'flat') weights = new Array(idxList.length).fill(1);
        else if (p.fuse === 'overlap-linear') weights = weightsOverlapLinear(idxList.length, numFrames, idxList, Math.max(0, p.overlap), errors);
        else if (p.fuse === 'hann') weights = weightsHann(idxList.length, numFrames, idxList, Math.max(0, p.overlap), errors);
        else if (p.fuse === 'gaussian') weights = weightsGaussian(idxList.length);
        else weights = weightsPyramid(idxList.length);
        const [rolls] = doesWindowRollOver(idxList, numFrames);
        return { idxList, anchor, weights, rolls };
    });

    // per-frame total weight (counts_final in the handler)
    const totals = new Float64Array(numFrames);
    for (const w of windows) {
        for (let pos = 0; pos < w.idxList.length; pos++) totals[w.idxList[pos]] += w.weights[pos];
    }
    const missing = [];
    for (let f = 0; f < numFrames; f++) if (totals[f] === 0) missing.push(f);
    let evals = 0;
    for (const w of windows) evals += w.idxList.length + (w.anchor >= 0 ? 1 : 0);

    return { numFrames, windows, totals, missing, evals, errors: [...new Set(errors)] };
}
// --- end context window math ---

const STYLE_ID = 'kj-cwviz-style';
function ensureStyles() {
    if (document.getElementById(STYLE_ID)) return;
    const s = document.createElement('style');
    s.id = STYLE_ID;
    s.textContent = `
.kj-cw-root { width:100%; height:100%; display:flex; flex-direction:column; background:#181818; overflow:hidden; user-select:none; border-radius:4px; }
.kj-cw-bar { display:flex; align-items:center; gap:6px; padding:4px 6px; background:#222; color:#ddd; font:12px sans-serif; flex:0 0 auto; }
.kj-cw-btn { background:#333; color:#ddd; border:1px solid #444; padding:1px 10px; cursor:pointer; font:12px sans-serif; border-radius:3px; }
.kj-cw-btn:hover { background:#444; }
.kj-cw-scrubber { position:relative; flex:1 1 auto; height:20px; background:#2a2a2a; border:1px solid #444; cursor:ew-resize; user-select:none; touch-action:none; border-radius:3px; overflow:hidden; }
.kj-cw-scrubber-fill { position:absolute; top:0; bottom:0; left:0; width:0; background:#3a5a7a; pointer-events:none; }
.kj-cw-scrubber-text { position:absolute; inset:0; display:flex; align-items:center; justify-content:center; color:#ddd; font:11px sans-serif; pointer-events:none; text-shadow:0 1px 2px rgba(0,0,0,0.6); }
.kj-cw-steps { width:42px; background:#2a2a2a; color:#ddd; border:1px solid #444; font:11px sans-serif; padding:2px 4px; border-radius:3px; }
.kj-cw-info { padding:2px 8px; color:#aaa; font:11px sans-serif; flex:0 0 auto; white-space:nowrap; overflow:hidden; text-overflow:ellipsis; }
.kj-cw-info.error { color:#e06c5a; }
.kj-cw-canvas { flex:1 1 auto; min-height:0; width:100%; display:block; cursor:crosshair; }
`;
    document.head.appendChild(s);
}

const STEP_DEPENDENT = ['standard_uniform', 'looped_uniform', 'batched_shifted']; // schedules whose windows change per sampling step

const WINDOW_HUE_STEP = 137.508; // golden angle, distinct adjacent colors
function windowColor(i, sat = 65, light = 60, alpha = 1) {
    return `hsla(${(i * WINDOW_HUE_STEP) % 360}, ${sat}%, ${light}%, ${alpha})`;
}

function pickTickStep(numFrames, plotW) {
    const steps = [1, 2, 5, 10, 20, 25, 50, 100, 200, 500, 1000];
    for (const s of steps) {
        if ((plotW / numFrames) * s >= 34) return s;
    }
    return steps[steps.length - 1];
}

app.registerExtension({
    name: 'KJNodes.ContextWindowsVisualizer',
    async beforeRegisterNodeDef(nodeType, nodeData) {
        if (nodeData?.name !== 'ContextWindowsVisualizerKJ') return;

        chainCallback(nodeType.prototype, 'onNodeCreated', function () {
            ensureStyles();
            const node = this;
            if (node.size[0] < 560 || node.size[1] < 480) node.setSize([Math.max(node.size[0], 560), Math.max(node.size[1], 480)]);

            const root = document.createElement('div');
            root.className = 'kj-cw-root';

            const bar = document.createElement('div');
            bar.className = 'kj-cw-bar';
            const playBtn = document.createElement('button');
            playBtn.className = 'kj-cw-btn';
            playBtn.textContent = '▶';
            playBtn.title = 'Animate sampling steps (uniform schedules shift their windows each step)';
            const scrubber = document.createElement('div');
            scrubber.className = 'kj-cw-scrubber';
            const scrubberFill = document.createElement('div');
            scrubberFill.className = 'kj-cw-scrubber-fill';
            const scrubberText = document.createElement('span');
            scrubberText.className = 'kj-cw-scrubber-text';
            scrubber.appendChild(scrubberFill);
            scrubber.appendChild(scrubberText);
            const stepsInput = document.createElement('input');
            stepsInput.className = 'kj-cw-steps';
            stepsInput.type = 'number';
            stepsInput.min = '1';
            stepsInput.max = '999';
            stepsInput.title = 'Total sampling steps to scrub through';
            bar.appendChild(playBtn);
            bar.appendChild(scrubber);
            bar.appendChild(stepsInput);

            const info = document.createElement('div');
            info.className = 'kj-cw-info';
            const canvas = document.createElement('canvas');
            canvas.className = 'kj-cw-canvas';

            root.appendChild(bar);
            root.appendChild(info);
            root.appendChild(canvas);
            node.addDOMWidget('context_windows_viz', 'div', root, { serialize: false });

            const removeMiddleClickPan = addMiddleClickPan(root);
            addWheelPassthrough(root);

            const state = {
                step: 0,
                steps: 20,
                viz: null,
                latent: null,     // converted params used for the viz
                hoverWin: -1,
                hoverFrame: -1,
                mouseX: 0,
                mouseY: 0,
                layout: null,     // hit-test geometry from last draw
            };

            const getWidgetValue = (name, fallback) => {
                const w = node.widgets?.find((w) => w.name === name);
                return w ? w.value : fallback;
            };

            const persist = () => {
                node.properties = node.properties || {};
                node.properties.cw_viz = { step: state.step, steps: state.steps };
            };

            const readParams = () => {
                const units = getWidgetValue('frame_units', 'pixel');
                const factor = Math.max(getWidgetValue('temporal_downscale', 4), 1);
                let numFrames = getWidgetValue('num_frames', 161);
                let length = getWidgetValue('context_length', 81);
                let overlap = getWidgetValue('context_overlap', 30);
                if (units === 'pixel') {
                    numFrames = Math.max(Math.floor((numFrames - 1) / factor) + 1, 1);
                    length = Math.max(Math.floor((length - 1) / factor) + 1, 1);
                    overlap = Math.max(Math.floor((overlap - 1) / factor) + 1, 0);
                }
                const schedule = getWidgetValue('context_schedule', 'standard_static');
                // flag the oversized final seam plain standard_static can leave
                const finalDelta = schedule === 'standard_static' ? staticFinalDelta(numFrames, length, overlap) : 0;
                return {
                    units,
                    factor,
                    numFrames,
                    length,
                    overlap,
                    finalDelta,
                    schedule,
                    stride: getWidgetValue('context_stride', 1),
                    closedLoop: !!getWidgetValue('closed_loop', false),
                    fuse: getWidgetValue('fuse_method', 'pyramid'),
                    causalFix: !!getWidgetValue('causal_window_fix', false),
                    step: state.step,
                };
            };

            const updateScrubberUI = () => {
                const t = state.steps > 1 ? state.step / (state.steps - 1) : 0;
                scrubberFill.style.width = (t * 100) + '%';
                const stepIndependent = state.latent && !STEP_DEPENDENT.includes(state.latent.schedule);
                scrubberText.textContent = `step ${state.step + 1} / ${state.steps}` + (stepIndependent ? ' (static schedule, no effect)' : '');
                stepsInput.value = String(state.steps);
            };

            const updateInfo = () => {
                if (!state.viz) return;
                const v = state.viz;
                if (v.errors.length) {
                    info.classList.add('error');
                    info.textContent = '⚠ ' + v.errors.join(' · ');
                    info.title = v.errors.join('\n');
                    return;
                }
                info.classList.remove('error');
                const lat = state.latent;
                const overhead = v.numFrames > 0 ? (v.evals / v.numFrames) : 0;
                const px = lat.units === 'pixel' ? ` (${(v.numFrames - 1) * lat.factor + 1}px frames)` : '';
                let txt = `${v.windows.length} windows over ${v.numFrames} latent frames${px} · ${v.evals} frame evals/step (${overhead.toFixed(2)}×)`;
                if (v.missing.length) txt += ` · ⚠ ${v.missing.length} uncovered frames`;
                if (lat.schedule === 'standard_static' && lat.finalDelta > 0) {
                    txt += ` · ⚠ last window overlaps +${lat.finalDelta} frames; standard_static_balanced spreads seams evenly`;
                } else if (lat.schedule === 'standard_static_balanced') {
                    txt += ` · balanced placement: even seams (±1 frame)`;
                } else if (lat.schedule === 'batched_shifted') {
                    const off = pythonRound(v.numFrames * orderedHalving(lat.step)) % lat.length;
                    txt += v.windows.length > 1 ? ` · zero-overlap rolling seams (step offset ${off})` : ` · zero-overlap`;
                }
                info.textContent = txt;
                info.title = txt;
            };

            let rafPending = false;
            const scheduleRender = () => {
                if (rafPending) return;
                rafPending = true;
                requestAnimationFrame(() => {
                    rafPending = false;
                    draw();
                });
            };

            const recompute = () => {
                state.latent = readParams();
                state.viz = computeViz(state.latent);
                state.hoverWin = -1;
                state.hoverFrame = -1;
                updateScrubberUI();
                updateInfo();
                scheduleRender();
            };

            const draw = () => {
                const viz = state.viz;
                if (!viz) return;
                const dpr = window.devicePixelRatio || 1;
                const cw = canvas.clientWidth;
                const ch = canvas.clientHeight;
                if (cw <= 0 || ch <= 0) return;
                const W = Math.floor(cw * dpr);
                const H = Math.floor(ch * dpr);
                if (canvas.width !== W) canvas.width = W;
                if (canvas.height !== H) canvas.height = H;
                const ctx = canvas.getContext('2d');
                ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
                ctx.clearRect(0, 0, cw, ch);
                ctx.fillStyle = '#181818';
                ctx.fillRect(0, 0, cw, ch);

                const n = viz.numFrames;
                const pad = 8;
                const gutter = 30;
                const rulerH = 16;
                const chartH = Math.max(50, Math.min(110, ch * 0.28));
                const chartGap = 18;
                const plotX = pad + gutter;
                const plotW = cw - plotX - pad;
                const rowsY = pad + rulerH + 4;
                const chartY = ch - pad - chartH;
                const rowsH = chartY - chartGap - rowsY;
                if (plotW <= 0 || rowsH <= 10) return;
                const cellW = plotW / n;
                const barW = Math.max(cellW - (cellW > 3 ? 0.5 : 0), 0.5); // cell width minus a hairline gap
                const wins = viz.windows;
                const rowH = wins.length ? Math.min(20, Math.max(3, rowsH / wins.length)) : 0;
                state.layout = { plotX, plotW, rowsY, rowH, chartY, chartH, n };

                // ruler
                ctx.font = '10px sans-serif';
                ctx.textBaseline = 'middle';
                ctx.fillStyle = '#777';
                ctx.strokeStyle = '#3a3a3a';
                const tick = pickTickStep(n, plotW);
                for (let f = 0; f < n; f += tick) {
                    const x = plotX + (f + 0.5) * cellW;
                    ctx.textAlign = 'center';
                    ctx.fillText(String(f), x, pad + rulerH / 2);
                    ctx.beginPath();
                    ctx.moveTo(x, rowsY);
                    ctx.lineTo(x, chartY + chartH);
                    ctx.stroke();
                }
                // uncovered frames
                if (viz.missing.length) {
                    ctx.fillStyle = 'rgba(220,70,60,0.25)';
                    for (const f of viz.missing) ctx.fillRect(plotX + f * cellW, rowsY, cellW, (chartY + chartH) - rowsY);
                }

                // window rows
                const hov = state.hoverWin;
                let drawnRows = wins.length;
                for (let i = 0; i < wins.length; i++) {
                    const w = wins[i];
                    const y = rowsY + i * rowH;
                    if (y + rowH > rowsY + rowsH + 1) { drawnRows = i; break; }
                    const dim = hov >= 0 && hov !== i;
                    const maxW = arrMax(w.weights);
                    // row label
                    if (rowH >= 9) {
                        ctx.fillStyle = dim ? '#555' : '#999';
                        ctx.textAlign = 'right';
                        ctx.fillText(String(i), plotX - 4, y + rowH / 2);
                    }
                    for (let pos = 0; pos < w.idxList.length; pos++) {
                        const f = w.idxList[pos];
                        const wt = maxW > 0 ? w.weights[pos] / maxW : 1;
                        const alpha = (0.25 + 0.7 * wt) * (dim ? 0.25 : 1);
                        // also encode weight as cell height when rows are tall
                        const hFrac = rowH >= 8 ? 0.35 + 0.65 * wt : 1;
                        const cellH = Math.max(1.5, (rowH - 1) * hFrac);
                        ctx.fillStyle = windowColor(i, 65, 60, alpha);
                        ctx.fillRect(plotX + f * cellW, y + (rowH - 1 - cellH), barW, cellH);
                    }
                    // causal anchor marker
                    if (w.anchor >= 0 && !dim) {
                        ctx.strokeStyle = windowColor(i, 65, 75, 0.9);
                        ctx.setLineDash([2, 2]);
                        ctx.strokeRect(plotX + w.anchor * cellW + 0.5, y + 0.5, Math.max(cellW - 1, 1), Math.max(rowH - 2, 1));
                        ctx.setLineDash([]);
                    }
                    if (hov === i) {
                        const lo = arrMin(w.idxList);
                        const hi = arrMax(w.idxList);
                        ctx.strokeStyle = '#fff';
                        ctx.lineWidth = 1;
                        ctx.strokeRect(plotX + lo * cellW, y, (hi - lo + 1) * cellW, rowH - 1);
                    }
                }

                state.layout.drawnRows = drawnRows;
                // blend chart: stacked normalized contribution per frame
                ctx.fillStyle = '#888';
                ctx.textAlign = 'left';
                ctx.fillText('blend (normalized contribution per frame)', plotX, chartY - 8);
                if (drawnRows < wins.length) {
                    ctx.textAlign = 'right';
                    ctx.fillText(`…+${wins.length - drawnRows} more windows (enlarge node)`, plotX + plotW, chartY - 8);
                }
                ctx.strokeStyle = '#3a3a3a';
                ctx.strokeRect(plotX, chartY, plotW, chartH);
                const stackOffsets = new Float64Array(n);
                for (let i = 0; i < wins.length; i++) {
                    const w = wins[i];
                    const dim = hov >= 0 && hov !== i;
                    ctx.fillStyle = windowColor(i, dim ? 25 : 65, dim ? 38 : 58, dim ? 0.6 : 0.92);
                    for (let pos = 0; pos < w.idxList.length; pos++) {
                        const f = w.idxList[pos];
                        const total = viz.totals[f];
                        if (total <= 0) continue;
                        const h = (w.weights[pos] / total) * chartH;
                        const yOff = (stackOffsets[f] / total) * chartH;
                        stackOffsets[f] += w.weights[pos];
                        ctx.fillRect(plotX + f * cellW, chartY + chartH - yOff - h, barW, h);
                    }
                }

                // hover guide + tooltip
                if (state.hoverFrame >= 0) {
                    const x = plotX + (state.hoverFrame + 0.5) * cellW;
                    ctx.strokeStyle = 'rgba(255,255,255,0.35)';
                    ctx.beginPath();
                    ctx.moveTo(x, rowsY);
                    ctx.lineTo(x, chartY + chartH);
                    ctx.stroke();
                }
                const lines = [];
                if (hov >= 0 && hov < wins.length) {
                    const w = wins[hov];
                    const strided = w.idxList.length > 1 ? w.idxList[1] - w.idxList[0] : 1;
                    let l = `window ${hov + 1}/${wins.length}: ${w.idxList.length} frames [${w.idxList[0]}${w.idxList[w.idxList.length - 1]}]`;
                    if (Math.abs(strided) > 1) l += ` stride ${strided}`;
                    if (w.rolls) l += ' (wraps)';
                    if (w.anchor >= 0) l += ` · anchor ${w.anchor}`;
                    lines.push(l);
                }
                if (state.hoverFrame >= 0) {
                    const f = state.hoverFrame;
                    const cover = [];
                    for (let i = 0; i < wins.length; i++) {
                        const pos = wins[i].idxList.indexOf(f);
                        if (pos >= 0) cover.push([i, wins[i].weights[pos]]);
                    }
                    const px = state.latent.units === 'pixel' ? ` (≈px ${f * state.latent.factor})` : '';
                    if (cover.length === 0) {
                        lines.push(`frame ${f}${px}: not covered by any window!`);
                    } else {
                        const total = viz.totals[f];
                        const shares = cover.map(([i, wt]) => `w${i + 1} ${(100 * wt / total).toFixed(0)}%`).join(', ');
                        lines.push(`frame ${f}${px}: ${shares}`);
                    }
                }
                if (lines.length) {
                    ctx.font = '11px sans-serif';
                    const tw = Math.max(...lines.map((l) => ctx.measureText(l).width)) + 12;
                    const th = lines.length * 15 + 8;
                    let tx = state.mouseX + 14;
                    let ty = state.mouseY + 14;
                    if (tx + tw > cw - 4) tx = state.mouseX - tw - 8;
                    if (ty + th > ch - 4) ty = state.mouseY - th - 8;
                    ctx.fillStyle = 'rgba(20,20,20,0.92)';
                    ctx.strokeStyle = '#555';
                    ctx.fillRect(tx, ty, tw, th);
                    ctx.strokeRect(tx, ty, tw, th);
                    ctx.fillStyle = '#ddd';
                    ctx.textAlign = 'left';
                    for (let i = 0; i < lines.length; i++) ctx.fillText(lines[i], tx + 6, ty + 12 + i * 15);
                }
            };

            // hover handling
            canvas.addEventListener('pointermove', (e) => {
                if (!state.viz || !state.layout) return;
                const rect = canvas.getBoundingClientRect();
                const x = e.clientX - rect.left;
                const y = e.clientY - rect.top;
                state.mouseX = x;
                state.mouseY = y;
                const lay = state.layout;
                let frame = Math.floor((x - lay.plotX) / (lay.plotW / lay.n));
                if (x < lay.plotX || frame < 0 || frame >= lay.n) frame = -1;
                let win = -1;
                const visibleRows = Math.min(lay.drawnRows ?? state.viz.windows.length, state.viz.windows.length);
                if (y >= lay.rowsY && y < lay.rowsY + lay.rowH * visibleRows && lay.rowH > 0) {
                    win = Math.floor((y - lay.rowsY) / lay.rowH);
                    if (win >= visibleRows) win = -1;
                } else if (y >= lay.chartY && y <= lay.chartY + lay.chartH && frame >= 0) {
                    // pick the stacked band under the cursor
                    const total = state.viz.totals[frame];
                    if (total > 0) {
                        const fromBottom = (lay.chartY + lay.chartH - y) / lay.chartH * total;
                        let acc = 0;
                        for (let i = 0; i < state.viz.windows.length; i++) {
                            const w = state.viz.windows[i];
                            const pos = w.idxList.indexOf(frame);
                            if (pos < 0) continue;
                            acc += w.weights[pos];
                            if (fromBottom <= acc) { win = i; break; }
                        }
                    }
                }
                if (win !== state.hoverWin || frame !== state.hoverFrame) {
                    state.hoverWin = win;
                    state.hoverFrame = frame;
                }
                scheduleRender();
            });
            canvas.addEventListener('pointerleave', () => {
                state.hoverWin = -1;
                state.hoverFrame = -1;
                scheduleRender();
            });

            // step scrubber
            let scrubbing = false;
            const scrubFromEvent = (e) => {
                const rect = scrubber.getBoundingClientRect();
                const t = Math.max(0, Math.min(1, (e.clientX - rect.left) / rect.width));
                const step = Math.round(t * (state.steps - 1));
                if (step !== state.step) {
                    state.step = step;
                    persist();
                    recompute();
                }
            };
            scrubber.addEventListener('pointerdown', (e) => {
                if (e.button !== 0) return;
                stopPlayback();
                scrubbing = true;
                scrubber.setPointerCapture(e.pointerId);
                scrubFromEvent(e);
                e.preventDefault();
            });
            scrubber.addEventListener('pointermove', (e) => {
                if (scrubbing) scrubFromEvent(e);
            });
            const endScrub = (e) => {
                if (!scrubbing) return;
                scrubbing = false;
                try { scrubber.releasePointerCapture(e.pointerId); } catch (_) { }
            };
            scrubber.addEventListener('pointerup', endScrub);
            scrubber.addEventListener('pointercancel', endScrub);

            stepsInput.addEventListener('change', () => {
                const v = parseInt(stepsInput.value, 10);
                if (Number.isFinite(v) && v >= 1 && v <= 999) {
                    state.steps = v;
                    if (state.step >= v) state.step = v - 1;
                    persist();
                    recompute();
                } else {
                    stepsInput.value = String(state.steps);
                }
            });
            stepsInput.addEventListener('pointerdown', (e) => e.stopPropagation());

            // step playback
            let playHandle = 0;
            const stopPlayback = () => {
                if (playHandle) {
                    clearInterval(playHandle);
                    playHandle = 0;
                    playBtn.textContent = '▶';
                }
            };
            playBtn.addEventListener('click', () => {
                if (playHandle) {
                    stopPlayback();
                    return;
                }
                playBtn.textContent = '⏸';
                playHandle = setInterval(() => {
                    state.step = (state.step + 1) % state.steps;
                    persist();
                    recompute();
                }, 300);
            });

            // live updates on widget changes
            const WATCHED = ['frame_units', 'temporal_downscale', 'num_frames', 'context_length', 'context_overlap',
                'context_schedule', 'context_stride', 'closed_loop', 'fuse_method', 'causal_window_fix'];
            for (const name of WATCHED) {
                const w = node.widgets?.find((w) => w.name === name);
                if (!w) continue;
                const orig = w.callback;
                w.callback = function (...args) {
                    const r = orig ? orig.apply(this, args) : undefined;
                    recompute();
                    return r;
                };
            }

            const ro = new ResizeObserver(() => scheduleRender());
            ro.observe(canvas);

            chainCallback(node, 'onConfigure', function () {
                const saved = this.properties?.cw_viz;
                if (saved) {
                    if (Number.isFinite(saved.steps)) state.steps = Math.min(999, Math.max(1, saved.steps));
                    if (Number.isFinite(saved.step)) state.step = Math.min(state.steps - 1, Math.max(0, saved.step));
                }
                recompute();
            });

            chainCallback(node, 'onRemoved', function () {
                stopPlayback();
                ro.disconnect();
                removeMiddleClickPan();
            });

            recompute();
        });
    },
});