EthanfelandClaude Opus 4.8
ea099ccac4
Per clarification: the 721 group lines are clip markers, NOT segments. Segments come ONLY from user split points, numbered globally from 1 — so the first split is seg1, not seg23. 0 splits = 0 segments; N splits = N segments (seg k = point k-1 .. split k, tail after the last split excluded). group lines still drawn (bold) + used for the group#. JS matches (segPoints/nSegs). Empty summary now prompts to add splits. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
234 lines
10 KiB
Python
234 lines
10 KiB
Python
"""
|
||
Audio Wave Segments node for ComfyUI.
|
||
|
||
An audio node with its own upload that (via the JS widget in web/audio_wave.js)
|
||
displays the waveform, plays the clip, and lets you click segment boundaries on
|
||
the waveform and give each a note. It outputs the same waveform IMAGE + timing
|
||
`audio_summary` as Audio Prompt Guide (feed those to the judge in chat mode),
|
||
plus the AUDIO for downstream use.
|
||
|
||
Works without the JS too: leave `segments_json` as "[]" and it auto-splits, or type
|
||
per-segment notes in `notes` using the seg1:/10s:/global syntax.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
|
||
import numpy as np
|
||
import torch
|
||
|
||
from .audio_guide import (
|
||
_rms_envelope, _tempo_beats, _snap8, _render, _summary, _attach_notes,
|
||
)
|
||
|
||
|
||
def _load_audio_file(path):
|
||
"""Load an audio file -> (waveform [C, N] float32 tensor, sample_rate).
|
||
Tries torchaudio, then soundfile, then librosa."""
|
||
try:
|
||
import torchaudio
|
||
wav, sr = torchaudio.load(path)
|
||
return wav.to(torch.float32), int(sr)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
import soundfile as sf
|
||
data, sr = sf.read(path, dtype="float32", always_2d=True) # [N, C]
|
||
return torch.from_numpy(data.T.copy()), int(sr)
|
||
except Exception:
|
||
pass
|
||
import librosa
|
||
y, sr = librosa.load(path, sr=None, mono=False)
|
||
y = np.atleast_2d(y)
|
||
return torch.from_numpy(np.ascontiguousarray(y, dtype=np.float32)), int(sr)
|
||
|
||
|
||
def _segments_from_boundaries(rms_n, times, duration, fps, starts, beats):
|
||
"""Build segments from user-placed boundary start times (JS click points)."""
|
||
starts = sorted({0.0} | {round(float(s), 3) for s in starts if 0.0 < float(s) < duration})
|
||
stages = ["establish", "build", "peak", "settle"]
|
||
bounds = np.array(starts + [duration])
|
||
segs = []
|
||
for i, t0 in enumerate(starts):
|
||
t1 = starts[i + 1] if i + 1 < len(starts) else duration
|
||
mask = (times >= t0) & (times < t1)
|
||
e = float(rms_n[mask].mean()) if mask.any() else 0.0
|
||
peak_t = float(times[mask][np.argmax(rms_n[mask])]) if mask.any() else t0
|
||
dur = round(t1 - t0, 2)
|
||
label = "high" if e > 0.66 else ("medium" if e > 0.33 else "low")
|
||
stage = stages[i] if i < len(stages) else ("peak" if e > 0.6 else "settle")
|
||
segs.append({
|
||
"segment": i + 1, "start_s": round(t0, 2), "duration_s": dur,
|
||
"frames": _snap8(round(dur * fps)), "energy": label,
|
||
"energy_val": round(e, 3), "peak_s": round(peak_t, 2),
|
||
"stage_hint": stage, "note": "",
|
||
"beats_in": [round(b, 2) for b in beats if t0 <= b < t1],
|
||
})
|
||
return segs, bounds
|
||
|
||
|
||
_NO_AUDIO = "(put an audio file in ComfyUI/input)"
|
||
|
||
|
||
def _audio_files():
|
||
files = []
|
||
try:
|
||
import folder_paths
|
||
d = folder_paths.get_input_directory()
|
||
try:
|
||
files = sorted(folder_paths.filter_files_content_types(os.listdir(d), ["audio", "video"]))
|
||
except Exception:
|
||
exts = (".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aac")
|
||
files = sorted(f for f in os.listdir(d) if f.lower().endswith(exts))
|
||
except Exception:
|
||
files = []
|
||
# A combo must never be empty, or ComfyUI can't build the node.
|
||
return files or [_NO_AUDIO]
|
||
|
||
|
||
def _parse_range(sel, n):
|
||
"""'' / '0' / 'all' -> (0,0) = whole clip; 'N' -> (N,N); 'A-B' -> (A,B). 1-based, clamped."""
|
||
sel = (sel or "").strip()
|
||
if not sel or sel in ("0", "all"):
|
||
return (0, 0)
|
||
m = re.match(r"^(\d+)\s*[-–:]\s*(\d+)$", sel)
|
||
if m:
|
||
a, b = int(m.group(1)), int(m.group(2))
|
||
else:
|
||
try:
|
||
a = b = int(sel)
|
||
except ValueError:
|
||
return (0, 0)
|
||
if a > b:
|
||
a, b = b, a
|
||
return (max(1, min(a, n)), max(1, min(b, n)))
|
||
|
||
|
||
def _segments_grouped(rms_n, times, duration, fps, group_frames, user_starts, beats):
|
||
"""Segments come ONLY from the user's split points (the 721 group lines are just clip
|
||
markers, NOT segments). Points = [0] + splits; a segment runs between consecutive
|
||
points, so N splits -> N segments (the tail after the last split is not a segment).
|
||
Global numbering. `group` = which 721 group each segment starts in (for reference)."""
|
||
gstep = max(group_frames, 1) / max(fps, 1)
|
||
gbounds, k = [], 1
|
||
while k * gstep < duration - 1e-6: # for drawing + the group number only
|
||
gbounds.append(round(k * gstep, 3)); k += 1
|
||
pts = [0.0] + sorted({round(s, 3) for s in user_starts if 0 < s < duration - 1e-6})
|
||
stages = ["establish", "build", "peak", "settle"]
|
||
segs, bounds = [], [0.0]
|
||
for i in range(len(pts) - 1): # exclude the tail (last point -> end)
|
||
t0, t1 = pts[i], pts[i + 1]
|
||
bounds.append(round(t1, 3))
|
||
mask = (times >= t0) & (times < t1)
|
||
e = float(rms_n[mask].mean()) if mask.any() else 0.0
|
||
peak_t = float(times[mask][np.argmax(rms_n[mask])]) if mask.any() else t0
|
||
dur = round(t1 - t0, 2)
|
||
segs.append({
|
||
"segment": i + 1, "group": 1 + sum(1 for gb in gbounds if t0 >= gb - 0.02),
|
||
"start_s": round(t0, 2), "duration_s": dur, "frames": _snap8(round(dur * fps)),
|
||
"energy": "high" if e > 0.66 else ("medium" if e > 0.33 else "low"),
|
||
"energy_val": round(e, 3), "peak_s": round(peak_t, 2),
|
||
"stage_hint": stages[i] if i < len(stages) else ("peak" if e > 0.6 else "settle"),
|
||
"note": "", "beats_in": [round(b, 2) for b in beats if t0 <= b < t1],
|
||
})
|
||
return segs, np.array(bounds), gbounds
|
||
|
||
|
||
class AudioWaveSegments:
|
||
CATEGORY = "prompt_calibrator"
|
||
FUNCTION = "run"
|
||
RETURN_TYPES = ("IMAGE", "STRING", "AUDIO")
|
||
RETURN_NAMES = ("waveform_image", "audio_summary", "audio")
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
# Pick a file from ComfyUI/input, or use the widget's "upload" button (JS).
|
||
"audio": (_audio_files(),),
|
||
"fps": ("INT", {"default": 24, "min": 1, "max": 120}),
|
||
# Hard split every this many frames (721@24fps ~= 30.04s = one LTX clip).
|
||
# Segments live inside a group and never cross a group line.
|
||
"group_frames": ("INT", {"default": 721, "min": 1, "max": 100000}),
|
||
# "" / "0" = whole clip. "N" = only segment N. "A-B" = segments A..B —
|
||
# crops image + audio + summary to that range (for generating/skipping beats).
|
||
"segment_select": ("STRING", {"default": ""}),
|
||
"notes": ("STRING", {"default": "", "multiline": True}),
|
||
# Machine field: the JS widget writes only USER split points here (hidden).
|
||
"segments_json": ("STRING", {"default": "[]"}),
|
||
},
|
||
}
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, audio, fps, group_frames, segment_select, notes, segments_json):
|
||
try:
|
||
import folder_paths
|
||
p = folder_paths.get_annotated_filepath(audio)
|
||
mt = os.path.getmtime(p) if os.path.isfile(p) else ""
|
||
except Exception:
|
||
mt = ""
|
||
return f"{audio}|{fps}|{group_frames}|{segment_select}|{notes}|{segments_json}|{mt}"
|
||
|
||
def run(self, audio, fps, group_frames, segment_select, notes, segments_json):
|
||
try:
|
||
import folder_paths
|
||
path = folder_paths.get_annotated_filepath(audio)
|
||
except Exception:
|
||
path = audio
|
||
if not path or not os.path.isfile(path):
|
||
blank = torch.zeros((1, 64, 512, 3))
|
||
return (blank, f"[AudioWaveSegments] audio not found: {audio}", None)
|
||
|
||
wav, sr = _load_audio_file(path) # wav: [C, N]
|
||
y = wav.mean(dim=0).cpu().numpy().astype(np.float32)
|
||
duration = len(y) / sr if sr else 0.0
|
||
rms_n, times = _rms_envelope(y, sr)
|
||
bpm, beats = _tempo_beats(y, sr)
|
||
|
||
user_starts, seg_notes = [], {} # user split points from the JS widget
|
||
try:
|
||
for i, seg in enumerate(json.loads(segments_json) if segments_json.strip() else []):
|
||
user_starts.append(float(seg.get("start_s", 0.0)))
|
||
if seg.get("note"):
|
||
seg_notes[round(float(seg.get("start_s", 0.0)), 2)] = str(seg["note"])
|
||
except Exception as e:
|
||
print(f"[AudioWaveSegments] bad segments_json ({e})")
|
||
|
||
if group_frames < 1:
|
||
group_frames = 721
|
||
segs, bounds, gbounds = _segments_grouped(rms_n, times, duration, fps, group_frames, user_starts, beats)
|
||
for s in segs: # notes attached by start time
|
||
if s["start_s"] in seg_notes:
|
||
s["note"] = seg_notes[s["start_s"]]
|
||
global_notes = _attach_notes(segs, notes) # notes-box (segN:) merges on top
|
||
|
||
a, b = _parse_range(segment_select, len(segs))
|
||
window, sel = None, None
|
||
audio_out = {"waveform": wav.unsqueeze(0), "sample_rate": sr}
|
||
if a >= 1 and segs: # crop to segment range A..B
|
||
s0 = segs[a - 1]["start_s"]
|
||
s1 = min(duration, segs[b - 1]["start_s"] + segs[b - 1]["duration_s"])
|
||
i0, i1 = int(s0 * sr), int(s1 * sr)
|
||
if i1 > i0:
|
||
audio_out = {"waveform": wav[:, i0:i1].unsqueeze(0), "sample_rate": sr}
|
||
window, sel = (s0, s1), (a, b)
|
||
total = sum(s["frames"] for s in segs[a - 1:b])
|
||
g0, g1 = segs[a - 1]["group"], segs[b - 1]["group"]
|
||
gtxt = f"group {g0}" + ("" if g0 == g1 else f"–{g1}")
|
||
summary = (f"SELECTED SEGMENTS {a}-{b} of {len(segs)} ({gtxt}), {total} frames total "
|
||
f"— generate these beats.\n"
|
||
+ _summary(duration, sr, bpm, beats, segs[a - 1:b], global_notes))
|
||
else:
|
||
summary = _summary(duration, sr, bpm, beats, segs, global_notes)
|
||
|
||
image = _render(rms_n, times, duration, beats, bounds, segs,
|
||
fps=fps, group_frames=group_frames, window=window, sel=sel)
|
||
return (image, summary, audio_out)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {"AudioWaveSegments": AudioWaveSegments}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {"AudioWaveSegments": "Audio Wave + Segments"}
|