""" 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, fps=None, group_frames=0, frames_marker=0, window=None, sel=None): """Render a mirrored waveform + the 721-frame group grid (bold) + segment lines (thin), labels along the top and per-segment notes along the bottom. window=(t0,t1) crops to a range; sel=(a,b) shades segments a..b.""" W, H = 1024, 256 img = Image.new("RGB", (W, H), (18, 18, 22)) d = ImageDraw.Draw(img) t0, t1 = window if window else (0.0, duration) span = max(t1 - t0, 1e-6) mid = H // 2 def X(t): return int(max(0, min(W - 1, (t - t0) / span * (W - 1)))) if sel and not window: # shade the selected segment range chosen = [s for s in segs if sel[0] <= s["segment"] <= sel[1]] if chosen: xa = X(chosen[0]["start_s"]) xb = X(chosen[-1]["start_s"] + chosen[-1]["duration_s"]) d.rectangle([xa, 0, xb, H], fill=(38, 54, 82)) amp = H * 0.44 # mirrored waveform around the centre line for i in range(len(rms_n)): if t0 <= times[i] <= t1: x = X(times[i]); h = int(rms_n[i] * amp) d.line([(x, mid - h), (x, mid + h)], fill=(60, 140, 220)) if fps and group_frames: # bold 721-frame group grid gstep = group_frames / fps k = 1 while k * gstep < duration + 1e-6: gt = k * gstep if t0 <= gt <= t1: d.line([(X(gt), 0), (X(gt), H)], fill=(235, 235, 242), width=2) k += 1 for b in beats: # beat ticks (faint, centre band) if t0 <= b <= t1: d.line([(X(b), mid - 4), (X(b), mid + 4)], fill=(150, 90, 60), width=1) last_lx, last_nx = -999, -999 # skip labels/notes that would overlap for s in segs: # segment line + label top + note bottom if not (t0 <= s["start_s"] <= t1): continue x = X(s["start_s"]) d.line([(x, 0), (x, H)], fill=(120, 190, 150), width=1) if x - last_lx > 30: d.text((x + 2, 3), f"S{s['segment']}", fill=(240, 240, 240)); last_lx = x if s.get("note") and x - last_nx > 70: d.text((x + 2, H - 13), s["note"][:22], fill=(255, 210, 110)); last_nx = x 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): if not segs: return (f"AUDIO GUIDE β€” {duration:.2f}s clip. No segments defined yet: double-click " f"the waveform to add split points (each split starts a new beat).") 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": "Audio Prompt Guide"}