ComfyUI-Prompt-Calibrator/web/audio_wave.js

// Audio Wave + Segments — waveform display + playback + click-to-segment.
// Writes the segment boundaries/notes into the node's hidden `segments_json` widget.
// First cut: open the browser console for [audiowave] logs if something misbehaves.
import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

function getW(node, name) { return node.widgets?.find((w) => w.name === name); }

function computePeaks(buf, n) {
  const data = buf.getChannelData(0);
  const block = Math.max(1, Math.floor(data.length / n));
  const peaks = new Float32Array(n);
  let max = 1e-6;
  for (let i = 0; i < n; i++) {
    let m = 0;
    const start = i * block;
    for (let j = 0; j < block && start + j < data.length; j++) {
      const v = Math.abs(data[start + j]);
      if (v > m) m = v;
    }
    peaks[i] = m;
    if (m > max) max = m;
  }
  for (let i = 0; i < n; i++) peaks[i] /= max;
  return peaks;
}

function setupWave(node) {
  const st = { duration: 0, peaks: null, boundaries: [], notes: {}, audio: new Audio(), playing: false };
  node._wave = st;

  const wrap = document.createElement("div");
  wrap.style.cssText = "display:flex;flex-direction:column;gap:4px;width:100%;";
  const bar = document.createElement("div");
  bar.style.cssText = "display:flex;gap:6px;align-items:center;font-size:10px;color:#bbb;flex-wrap:wrap;";
  const mk = (t) => { const b = document.createElement("button"); b.textContent = t; b.style.cssText = "font-size:10px;padding:1px 6px;"; return b; };
  const playBtn = mk("▶ play"), clearBtn = mk("clear"), upBtn = mk("upload");
  const hint = document.createElement("span");
  hint.textContent = "click=add split · drag=move · dblclick=note · shift/right-click=delete";
  bar.append(playBtn, clearBtn, upBtn, hint);
  const canvas = document.createElement("canvas");
  canvas.width = 640; canvas.height = 150;
  canvas.style.cssText = "width:100%;height:150px;background:#141418;border-radius:4px;cursor:crosshair;";
  wrap.append(bar, canvas);
  node.addDOMWidget("wave", "wave", wrap, { serialize: false });
  const ctx = canvas.getContext("2d");

  const t2x = (t) => (st.duration ? (t / st.duration) * canvas.width : 0);
  const x2t = (x) => (st.duration ? (x / canvas.width) * st.duration : 0);
  const starts = () => [0, ...st.boundaries.slice().sort((a, b) => a - b)];

  function serialize() {
    const segs = starts().map((s, i) => ({ start_s: Math.round(s * 100) / 100, note: st.notes[i] || "" }));
    const w = getW(node, "segments_json");
    if (w) { w.value = JSON.stringify(segs); w.callback?.(w.value); }
  }

  function draw() {
    ctx.fillStyle = "#141418"; ctx.fillRect(0, 0, canvas.width, canvas.height);
    if (st.peaks) {
      ctx.fillStyle = "#3c8cdc";
      const n = st.peaks.length, bw = canvas.width / n;
      for (let i = 0; i < n; i++) { const h = st.peaks[i] * (canvas.height - 26); ctx.fillRect(i * bw, canvas.height - h, Math.max(1, bw), h); }
    }
    ctx.font = "10px monospace";
    starts().forEach((s, i) => {
      const x = t2x(s);
      ctx.strokeStyle = "#eee"; ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, canvas.height); ctx.stroke();
      ctx.fillStyle = "#fff"; ctx.fillText("S" + (i + 1), x + 3, 11);
      if (st.notes[i]) { ctx.fillStyle = "#ffd27a"; ctx.fillText(st.notes[i].slice(0, 24), x + 3, 23); }
    });
    if (st.playing) { const x = t2x(st.audio.currentTime); ctx.strokeStyle = "#ff5a3c"; ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, canvas.height); ctx.stroke(); }
  }

  async function loadFile(name) {
    if (!name) return;
    const url = `/view?filename=${encodeURIComponent(name)}&type=input&subfolder=`;
    try {
      const ab = await (await fetch(url)).arrayBuffer();
      const ac = new (window.AudioContext || window.webkitAudioContext)();
      const buf = await ac.decodeAudioData(ab.slice(0));
      st.duration = buf.duration; st.peaks = computePeaks(buf, canvas.width);
      st.audio.src = url; draw();
    } catch (e) { console.error("[audiowave] could not load/decode", name, e); }
  }

  const tol = () => x2t(6);
  const nearBoundary = (t) => { let bi = -1, bd = 1e9; st.boundaries.forEach((b, i) => { const d = Math.abs(b - t); if (d < tol() && d < bd) { bd = d; bi = i; } }); return bi; };
  let dragIdx = -1;

  canvas.addEventListener("mousedown", (e) => {
    const r = canvas.getBoundingClientRect();
    const t = x2t(((e.clientX - r.left) / r.width) * canvas.width);
    if (e.button === 2 || e.shiftKey) { const bi = nearBoundary(t); if (bi >= 0) { st.boundaries.splice(bi, 1); serialize(); draw(); } e.preventDefault(); return; }
    const bi = nearBoundary(t);
    if (bi >= 0) { dragIdx = bi; }
    else { st.boundaries.push(t); serialize(); draw(); }
  });
  canvas.addEventListener("mousemove", (e) => {
    if (dragIdx < 0) return;
    const r = canvas.getBoundingClientRect();
    st.boundaries[dragIdx] = Math.max(0.01, Math.min(st.duration - 0.01, x2t(((e.clientX - r.left) / r.width) * canvas.width)));
    draw();
  });
  window.addEventListener("mouseup", () => { if (dragIdx >= 0) { serialize(); dragIdx = -1; draw(); } });
  canvas.addEventListener("dblclick", (e) => {
    const r = canvas.getBoundingClientRect();
    const t = x2t(((e.clientX - r.left) / r.width) * canvas.width);
    const ss = starts(); let si = 0; for (let i = 0; i < ss.length; i++) if (t >= ss[i]) si = i;
    const val = window.prompt(`Note for segment ${si + 1}:`, st.notes[si] || "");
    if (val !== null) { st.notes[si] = val; serialize(); draw(); }
  });
  canvas.addEventListener("contextmenu", (e) => e.preventDefault());

  const loop = () => { if (!st.playing) return; draw(); requestAnimationFrame(loop); };
  playBtn.onclick = () => {
    if (st.audio.paused) { st.audio.play().catch((e) => console.error("[audiowave] play", e)); st.playing = true; playBtn.textContent = "⏸ pause"; loop(); }
    else { st.audio.pause(); st.playing = false; playBtn.textContent = "▶ play"; }
  };
  st.audio.onended = () => { st.playing = false; playBtn.textContent = "▶ play"; draw(); };
  clearBtn.onclick = () => { st.boundaries = []; st.notes = {}; serialize(); draw(); };

  upBtn.onclick = () => {
    const inp = document.createElement("input"); inp.type = "file"; inp.accept = "audio/*";
    inp.onchange = async () => {
      const f = inp.files[0]; if (!f) return;
      const fd = new FormData(); fd.append("image", f, f.name); fd.append("type", "input");
      try {
        const res = await api.fetchApi("/upload/image", { method: "POST", body: fd });
        const j = await res.json(); const name = j.name || f.name;
        const w = getW(node, "audio");
        if (w) { if (!w.options.values.includes(name)) w.options.values.push(name); w.value = name; }
        await loadFile(name);
      } catch (e) { console.error("[audiowave] upload failed (drop the file in ComfyUI/input instead)", e); }
    };
    inp.click();
  };

  // react to the audio combo changing + initial load
  const aw = getW(node, "audio");
  if (aw) { const cb = aw.callback; aw.callback = function () { const r = cb ? cb.apply(this, arguments) : undefined; loadFile(aw.value); return r; }; if (aw.value) loadFile(aw.value); }
  // restore boundaries/notes from a reloaded workflow
  try { const sj = getW(node, "segments_json"); if (sj && sj.value) { const arr = JSON.parse(sj.value); st.boundaries = arr.slice(1).map((s) => s.start_s); arr.forEach((s, i) => { if (s.note) st.notes[i] = s.note; }); } } catch (e) { /* ignore */ }
  draw();
}

app.registerExtension({
  name: "promptcalib.audiowave",
  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData?.name !== "AudioWaveSegments") return;
    const orig = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      const r = orig?.apply(this, arguments);
      try { setupWave(this); } catch (e) { console.error("[audiowave] setup failed", e); }
      return r;
    };
  },
});