EthanfelandClaude Opus 4.8
75b20f9656
Pivot to the user's model: segments are a fixed grid of subsegment_frames frames (default 721 @ 24fps = one LTX clip), not arbitrary clicks. New inputs: subsegment_frames (grid size) and segment_select (0=all, N=output ONLY chunk N — crops waveform_image, AUDIO, and summary so you can generate/skip one beat at a time). _render gains window-crop + frame markers + per-segment time/frame labels. JS rewritten: draws the fixed grid, auto-fills the notes box with one segN: line per chunk (type or dblclick to note), click-to-seek, playhead + time readout. Workflows updated for the new widgets. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
212 lines
9.0 KiB
Python
212 lines
9.0 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 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
|
|
|
|
|
|
_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 _segments_fixed(rms_n, times, duration, fps, sub_frames, beats):
|
|
"""Split into fixed chunks of `sub_frames` frames (the LTX clip length). The last
|
|
chunk gets whatever's left (snapped to 8n+1)."""
|
|
step = max(sub_frames, 1) / max(fps, 1) # seconds per chunk
|
|
n = max(1, int(np.ceil(duration / step - 1e-6)))
|
|
stages = ["establish", "build", "peak", "settle"]
|
|
segs, bounds = [], [0.0]
|
|
for i in range(n):
|
|
t0 = i * step
|
|
t1 = min(duration, (i + 1) * step)
|
|
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)
|
|
frames = sub_frames if i < n - 1 else _snap8(round(dur * fps))
|
|
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": frames,
|
|
"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, np.array(bounds)
|
|
|
|
|
|
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}),
|
|
# Fixed LTX clip length. Segments = chunks of this many frames (721@24fps
|
|
# ~= 30.04s). This is also the waveform grid. Set 0 to use segments_json.
|
|
"subsegment_frames": ("INT", {"default": 721, "min": 0, "max": 100000}),
|
|
# 0 = whole clip. N = output ONLY subsegment N (crops image + audio + summary),
|
|
# so you can generate/skip one beat at a time.
|
|
"segment_select": ("INT", {"default": 0, "min": 0, "max": 999}),
|
|
"notes": ("STRING", {"default": "", "multiline": True}),
|
|
# Machine field: per-chunk notes/boundaries from the JS widget (hidden).
|
|
"segments_json": ("STRING", {"default": "[]"}),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, audio, fps, subsegment_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}|{subsegment_frames}|{segment_select}|{notes}|{segments_json}|{mt}"
|
|
|
|
def run(self, audio, fps, subsegment_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)
|
|
|
|
if subsegment_frames > 0: # fixed LTX-clip grid
|
|
segs, bounds = _segments_fixed(rms_n, times, duration, fps, subsegment_frames, beats)
|
|
else: # manual boundaries from the JS widget
|
|
starts = []
|
|
try:
|
|
for seg in (json.loads(segments_json) if segments_json.strip() else []):
|
|
starts.append(float(seg.get("start_s", 0.0)))
|
|
except Exception as e:
|
|
print(f"[AudioWaveSegments] bad segments_json ({e})")
|
|
segs, bounds = _segments_from_boundaries(rms_n, times, duration, fps, starts, beats)
|
|
|
|
# per-chunk notes from segments_json (JS) then the notes-box (segN: syntax) on top.
|
|
try:
|
|
for i, seg in enumerate(json.loads(segments_json) if segments_json.strip() else []):
|
|
if i < len(segs) and seg.get("note"):
|
|
segs[i]["note"] = str(seg["note"])
|
|
except Exception:
|
|
pass
|
|
global_notes = _attach_notes(segs, notes)
|
|
|
|
window, render_segs = None, segs
|
|
audio_out = {"waveform": wav.unsqueeze(0), "sample_rate": sr}
|
|
sel = int(segment_select)
|
|
if 1 <= sel <= len(segs): # crop to one subsegment
|
|
seg = segs[sel - 1]
|
|
s0 = seg["start_s"]
|
|
s1 = min(duration, s0 + seg["duration_s"])
|
|
a, b = int(s0 * sr), int(s1 * sr)
|
|
if b > a:
|
|
audio_out = {"waveform": wav[:, a:b].unsqueeze(0), "sample_rate": sr}
|
|
window, render_segs = (s0, s1), [seg]
|
|
summary = (f"SELECTED SUBSEGMENT {sel} of {len(segs)} — generate ONLY this beat.\n"
|
|
+ _summary(duration, sr, bpm, beats, [seg], global_notes))
|
|
else:
|
|
summary = _summary(duration, sr, bpm, beats, segs, global_notes)
|
|
|
|
image = _render(rms_n, times, duration, beats, bounds, render_segs,
|
|
fps=fps, frames_marker=subsegment_frames, window=window)
|
|
return (image, summary, audio_out)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"AudioWaveSegments": AudioWaveSegments}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"AudioWaveSegments": "Audio Wave + Segments"}
|