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