837 lines
36 KiB
JavaScript
837 lines
36 KiB
JavaScript
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();
|
||
});
|
||
},
|
||
});
|