""" Audio Prompt Guide node for ComfyUI. Turns an audio clip + your free-text motion notes into two things the (vision-only) VLM can use to write an audio-aligned LTX beat timeline: 1. waveform_image (IMAGE) - a rendered energy envelope with beat markers and suggested segment boundaries, so the model can *see* the audio's shape. 2. audio_summary (STRING) - duration, tempo/beats (if librosa is installed), a per-segment energy breakdown with snapped frame counts, and your notes, formatted as guidance the model can map onto beats. Wire waveform_image -> the judge node's `reference_image` and audio_summary -> its `user_prompt` (mode=chat, json_output=true), with your LTX system prompt. librosa is optional: without it you still get the energy envelope + segments; with it you also get tempo (BPM) and beat times. """ from __future__ import annotations import re import numpy as np import torch from PIL import Image, ImageDraw def _split_notes(notes: str): """Parse the notes box into per-segment / per-time / global notes. Lines like: seg2: fast motion (or 2: fast motion / S2: ...) -> segment 2 5s: build (or 0-5s: ... / 12.3s: ...) -> by time cinematic, moody (no prefix) -> global Returns (seg_map {num: [notes]}, time_notes [(t0,t1,text)], global_notes [str]).""" seg_map, time_notes, glob = {}, [], [] for raw in notes.splitlines(): line = raw.strip() if not line: continue mt = re.match(r"(?i)^(\d+(?:\.\d+)?)\s*s(?:\s*[-–to]+\s*(\d+(?:\.\d+)?)\s*s?)?\s*[:)\-]\s*(.+)$", line) ms = (re.match(r"(?i)^(?:seg(?:ment)?|s)\s*(\d+)\s*[:)\-]\s*(.+)$", line) or re.match(r"^(\d+)\s*[:)]\s*(.+)$", line)) if mt: t0 = float(mt.group(1)); t1 = float(mt.group(2)) if mt.group(2) else t0 time_notes.append((t0, t1, mt.group(3).strip())) elif ms: seg_map.setdefault(int(ms.group(1)), []).append(ms.group(2).strip()) else: glob.append(line) return seg_map, time_notes, glob def _attach_notes(segs, notes): """Attach the parsed notes to each segment (by number or overlapping time). Returns the leftover global notes.""" seg_map, time_notes, glob = _split_notes(notes) for s in segs: s0, s1 = s["start_s"], s["start_s"] + s["duration_s"] existing = (s.get("note") or "").strip() found = ([existing] if existing else []) + list(seg_map.get(s["segment"], [])) found += [txt for (t0, t1, txt) in time_notes if t0 < s1 and t1 > s0] s["note"] = "; ".join(found) return glob def _to_mono(audio) -> tuple[np.ndarray, int]: """ComfyUI AUDIO dict -> (mono float32 samples, sample_rate).""" wf = audio["waveform"] sr = int(audio["sample_rate"]) arr = wf.detach().cpu().numpy() if hasattr(wf, "detach") else np.asarray(wf) arr = np.asarray(arr, dtype=np.float32) while arr.ndim > 2: # [B, C, N] -> [C, N] arr = arr[0] if arr.ndim == 2: # [C, N] -> mono arr = arr.mean(axis=0) return arr, sr def _rms_envelope(y: np.ndarray, sr: int, fps_env: int = 100) -> tuple[np.ndarray, np.ndarray]: """RMS energy per ~1/fps_env-second frame, normalized 0..1, with frame times.""" hop = max(1, sr // fps_env) n = (len(y) // hop) * hop if n < hop: return np.array([0.0]), np.array([0.0]) frames = y[:n].reshape(-1, hop) rms = np.sqrt((frames ** 2).mean(axis=1) + 1e-9) rms_n = rms / (rms.max() + 1e-9) times = np.arange(len(rms_n)) * hop / sr return rms_n, times def _tempo_beats(y: np.ndarray, sr: int): """(bpm, beat_times) via librosa if available, else (None, []).""" try: import librosa tempo, beat_frames = librosa.beat.beat_track(y=y.astype(np.float32), sr=sr) beats = librosa.frames_to_time(beat_frames, sr=sr).tolist() return float(np.atleast_1d(tempo)[0]), beats except Exception: return None, [] def _snap8(frames: int) -> int: """Snap a frame count to LTX's 8n+1 grid (e.g. 120 -> 121).""" n = max(0, round((frames - 1) / 8)) return int(8 * n + 1) def _segments(rms_n, times, duration, fps, max_segments, beats): """Split into ~5s segments; label each by mean energy + snapped frame count.""" seg_count = int(min(max(3, round(duration / 5.0)), max(3, max_segments))) bounds = np.linspace(0.0, duration, seg_count + 1) stages = ["establish", "build", "peak", "settle"] segs = [] for i in range(seg_count): t0, t1 = float(bounds[i]), float(bounds[i + 1]) 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, "beats_in": [round(b, 2) for b in beats if t0 <= b < t1], }) return segs, bounds def _render(rms_n, times, duration, beats, bounds, segs): """Render the envelope + beats + segment boundaries to a ComfyUI IMAGE tensor.""" W, H = 1024, 256 img = Image.new("RGB", (W, H), (18, 18, 22)) d = ImageDraw.Draw(img) dur = max(duration, 1e-6) def X(t): return int(max(0, min(W - 1, t / dur * (W - 1)))) pts = [(0, H)] + [(X(times[i]), H - int(rms_n[i] * (H - 26))) for i in range(len(rms_n))] + [(W - 1, H)] d.polygon(pts, fill=(60, 140, 220)) for b in beats: # beat markers (orange) d.line([(X(b), 0), (X(b), H)], fill=(230, 110, 60), width=1) for bd in bounds: # segment boundaries (white) d.line([(X(bd), 0), (X(bd), H)], fill=(240, 240, 240), width=1) for s in segs: x = X(s["start_s"]) + 4 d.text((x, 4), f"S{s['segment']} {s['energy']}", fill=(255, 255, 255)) if s.get("note"): # per-segment note (amber) d.text((x, 18), s["note"][:30], fill=(255, 210, 110)) arr = np.asarray(img, dtype=np.float32) / 255.0 return torch.from_numpy(arr)[None, ...] # [1, H, W, 3] def _summary(duration, sr, bpm, beats, segs, global_notes): lines = ["AUDIO GUIDE β€” align the video beats to this audio:", f"- duration: {duration:.2f}s | sample_rate: {sr} Hz"] if bpm: lines.append(f"- tempo: ~{bpm:.0f} BPM") if beats: preview = ", ".join(f"{b:.2f}" for b in beats[:24]) lines.append(f"- beat times (s): {preview}{' ...' if len(beats) > 24 else ''}") lines.append("- suggested segments (use these durations/frames and energy; honor the NOTE):") for s in segs: peak = f", energy peak @ {s['peak_s']}s" if s["energy"] == "high" else "" note = f" <<< NOTE: {s['note']}" if s.get("note") else "" lines.append( f" seg{s['segment']}: {s['start_s']}–{s['start_s'] + s['duration_s']:.2f}s, " f"{s['duration_s']}s, {s['frames']} frames, energy {s['energy'].upper()} " f"({s['stage_hint']}){peak}{note}") lines.append("") lines.append("Guidance: higher energy -> faster motion and bigger camera moves; " "lower energy -> slower, settle. Put escalation on rising energy and the " "release on the final segment. Snap each beat's frames to 8n+1. Where a " "segment has a NOTE, that instruction OVERRIDES the energy default.") if global_notes: lines.append("") lines.append("GLOBAL NOTES (apply throughout):") lines.extend(f" - {n}" for n in global_notes) return "\n".join(lines) class AudioPromptGuide: CATEGORY = "prompt_calibrator" FUNCTION = "guide" RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("waveform_image", "audio_summary") @classmethod def INPUT_TYPES(cls): return { "required": { "audio": ("AUDIO",), "notes": ("STRING", {"default": "", "multiline": True}), "fps": ("INT", {"default": 24, "min": 1, "max": 120}), "max_segments": ("INT", {"default": 6, "min": 3, "max": 12}), }, } def guide(self, audio, notes, fps, max_segments): y, sr = _to_mono(audio) duration = len(y) / sr if sr else 0.0 rms_n, times = _rms_envelope(y, sr) bpm, beats = _tempo_beats(y, sr) segs, bounds = _segments(rms_n, times, duration, fps, max_segments, beats) global_notes = _attach_notes(segs, notes) # per-segment / per-time / global image = _render(rms_n, times, duration, beats, bounds, segs) summary = _summary(duration, sr, bpm, beats, segs, global_notes) return (image, summary) NODE_CLASS_MAPPINGS = {"AudioPromptGuide": AudioPromptGuide} NODE_DISPLAY_NAME_MAPPINGS = {"AudioPromptGuide": "SxCP Audio Prompt Guide"}