Ethanfel
c4d7224118
Labels crowded when splits are near each other. Now labels are shortened to S{n} and
skipped if within ~26px (canvas) / 30px (image) of the previous drawn label; notes gated
wider. Lines still drawn for every split.
246 lines
11 KiB
Python
246 lines
11 KiB
Python
"""
|
||
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"}
|