""" 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"}