ComfyUI-Prompt-Calibrator/nodes/audio_wave_segments.py

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


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}),
                "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": "Audio Wave + Segments"}