Add SxCP Audio Wave + Segments: interactive waveform node (upload/play/click-segments)
New node with a JS widget (web/audio_wave.js): upload an audio clip, play it, and click the waveform to place segment boundaries (click add / drag move / dblclick note / shift|right-click delete). Boundaries+notes serialize to a hidden segments_json that drives Python segmentation (falls back to auto-split / notes syntax). Python node (nodes/audio_wave_segments.py) loads the file (torchaudio/soundfile/librosa), builds segments from the boundaries, and outputs waveform_image + audio_summary + AUDIO — same contract as Audio Prompt Guide, so it feeds chat mode the same way. _attach_notes now merges (keeps clicked notes). WEB_DIRECTORY re-enabled. JS is a first cut — needs testing in ComfyUI (console logs on error); the Python node works standalone via segments_json/notes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,157 @@
|
||||
"""
|
||||
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 SxCP 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 numpy as np
|
||||
import torch
|
||||
|
||||
from .audio_guide import (
|
||||
_rms_envelope, _tempo_beats, _snap8, _segments, _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
|
||||
|
||||
|
||||
def _audio_files():
|
||||
try:
|
||||
import folder_paths
|
||||
d = folder_paths.get_input_directory()
|
||||
try:
|
||||
return sorted(folder_paths.filter_files_content_types(os.listdir(d), ["audio", "video"]))
|
||||
except Exception:
|
||||
exts = (".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aac")
|
||||
return sorted(f for f in os.listdir(d) if f.lower().endswith(exts))
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
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": {
|
||||
"audio": (_audio_files(), {"audio_upload": True}),
|
||||
"fps": ("INT", {"default": 24, "min": 1, "max": 120}),
|
||||
"max_segments": ("INT", {"default": 6, "min": 3, "max": 12}),
|
||||
"notes": ("STRING", {"default": "", "multiline": True}),
|
||||
# Written by the JS waveform widget: [{"start_s": 0.0, "note": "..."}, ...].
|
||||
# Leave as "[]" to auto-split. Also editable by hand.
|
||||
"segments_json": ("STRING", {"default": "[]"}),
|
||||
},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, audio, fps, max_segments, 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}|{max_segments}|{notes}|{segments_json}|{mt}"
|
||||
|
||||
def run(self, audio, fps, max_segments, 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 boundaries from the JS widget, else auto-split.
|
||||
starts, seg_notes = [], {}
|
||||
try:
|
||||
data = json.loads(segments_json) if segments_json.strip() else []
|
||||
for i, seg in enumerate(data):
|
||||
starts.append(float(seg.get("start_s", 0.0)))
|
||||
seg_notes[i + 1] = str(seg.get("note", "") or "")
|
||||
except Exception as e:
|
||||
print(f"[AudioWaveSegments] bad segments_json ({e}); auto-splitting.")
|
||||
data = []
|
||||
|
||||
if data:
|
||||
segs, bounds = _segments_from_boundaries(rms_n, times, duration, fps, starts, beats)
|
||||
for s in segs: # notes placed on the waveform
|
||||
s["note"] = seg_notes.get(s["segment"], "")
|
||||
else:
|
||||
segs, bounds = _segments(rms_n, times, duration, fps, max_segments, beats)
|
||||
|
||||
global_notes = _attach_notes(segs, notes) # merge the text-box notes on top
|
||||
image = _render(rms_n, times, duration, beats, bounds, segs)
|
||||
summary = _summary(duration, sr, bpm, beats, segs, global_notes)
|
||||
audio_out = {"waveform": wav.unsqueeze(0), "sample_rate": sr} # [1, C, N]
|
||||
return (image, summary, audio_out)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"AudioWaveSegments": AudioWaveSegments}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"AudioWaveSegments": "SxCP Audio Wave + Segments"}
|
||||
Reference in New Issue
Block a user