grmchnandClaude Opus 4.6
a65caecfaa
DrawRhythmTracks generates BPM-synced sparse track point images with independent X/Y swing patterns (pendulum, bounce, linear). Rendering is GPU-accelerated and compatible with LTXVDrawTracks output format. BPMDetector extracts tempo from audio input using librosa beat tracking. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
392 lines
13 KiB
Python
392 lines
13 KiB
Python
import math
|
|
import json
|
|
|
|
import numpy as np
|
|
import torch
|
|
from comfy_api.latest import io
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Constants (matching LTXVDrawTracks rendering)
|
|
# ---------------------------------------------------------------------------
|
|
_MIN_RADIUS = 2
|
|
_MAX_RADIUS = 8
|
|
_MAX_TRAIL = 50
|
|
_REF_SHORT_SIDE = 1080
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Swing pattern functions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _pendulum(phase: torch.Tensor, size: int, rate: float) -> torch.Tensor:
|
|
"""Sinusoidal swing: offset = sin(phase) * size * rate / 2"""
|
|
return torch.sin(phase) * size * rate / 2
|
|
|
|
|
|
def _bounce(phase: torch.Tensor, size: int, rate: float) -> torch.Tensor:
|
|
"""Parabolic bounce arches with phase shift so offset=0 at phase=0.
|
|
|
|
offset = (0.5 - t^2) * size * rate
|
|
where t = (((phase + phi) / pi) % 2) - 1
|
|
phi = pi * (1 + 1/sqrt(2))
|
|
"""
|
|
phi = math.pi * (1 + 1 / math.sqrt(2))
|
|
t = ((phase + phi) / math.pi % 2) - 1
|
|
return (0.5 - t ** 2) * size * rate
|
|
|
|
|
|
def _linear(phase: torch.Tensor, size: int, rate: float) -> torch.Tensor:
|
|
"""Triangle wave: offset = trianglewave(phase) * size * rate / 2
|
|
|
|
trianglewave(phase) = (2/pi) * arcsin(sin(phase))
|
|
"""
|
|
return (2 / math.pi) * torch.arcsin(torch.sin(phase)) * size * rate / 2
|
|
|
|
|
|
_PATTERNS = {
|
|
"pendulum": _pendulum,
|
|
"bounce": _bounce,
|
|
"linear": _linear,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rendering helpers (ported from LTXVDrawTracks)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _render_resolution(width: int, height: int, reference_short_side: int):
|
|
"""Compute the higher render resolution that preserves aspect ratio."""
|
|
if height <= width:
|
|
rw = int(width * reference_short_side / height)
|
|
rh = reference_short_side
|
|
else:
|
|
rw = reference_short_side
|
|
rh = int(height * reference_short_side / width)
|
|
scale_x = rw / width
|
|
scale_y = rh / height
|
|
return rw, rh, scale_x, scale_y
|
|
|
|
|
|
def _age_color_batch(ratios: torch.Tensor, device: torch.device) -> torch.Tensor:
|
|
"""Vectorised age-ratio -> RGB [0..1] mapping on GPU.
|
|
|
|
Gradient: blue -> green -> yellow -> red.
|
|
"""
|
|
colors = torch.zeros(ratios.shape[0], 3, device=device)
|
|
|
|
m1 = ratios <= 1 / 3
|
|
tr1 = ratios[m1] * 3
|
|
colors[m1, 1] = tr1
|
|
colors[m1, 2] = 1 - tr1
|
|
|
|
m2 = (ratios > 1 / 3) & (ratios <= 2 / 3)
|
|
tr2 = (ratios[m2] - 1 / 3) * 3
|
|
colors[m2, 0] = tr2
|
|
colors[m2, 1] = 1
|
|
|
|
m3 = ratios > 2 / 3
|
|
tr3 = (ratios[m3] - 2 / 3) * 3
|
|
colors[m3, 0] = 1
|
|
colors[m3, 1] = 1 - tr3
|
|
|
|
return colors
|
|
|
|
|
|
def _rasterise_circles(
|
|
frame: torch.Tensor,
|
|
pts: torch.Tensor,
|
|
radii: torch.Tensor,
|
|
colors: torch.Tensor,
|
|
template_dist_sq: torch.Tensor,
|
|
half_d: int,
|
|
max_d: int,
|
|
H: int,
|
|
W: int,
|
|
) -> None:
|
|
"""Stamp filled circles onto *frame* fully on-device."""
|
|
M = pts.shape[0]
|
|
if M == 0:
|
|
return
|
|
device = pts.device
|
|
|
|
radii_sq = (radii * radii).view(M, 1, 1)
|
|
circle_masks = template_dist_sq.unsqueeze(0) <= radii_sq
|
|
|
|
cx = pts[:, 0].round().long().view(M, 1, 1)
|
|
cy = pts[:, 1].round().long().view(M, 1, 1)
|
|
offsets_y = torch.arange(max_d, device=device).sub(half_d).view(1, max_d, 1)
|
|
offsets_x = torch.arange(max_d, device=device).sub(half_d).view(1, 1, max_d)
|
|
fy = (cy + offsets_y).expand(M, max_d, max_d)
|
|
fx = (cx + offsets_x).expand(M, max_d, max_d)
|
|
|
|
valid = circle_masks & (fy >= 0) & (fy < H) & (fx >= 0) & (fx < W)
|
|
|
|
flat_fy = fy[valid]
|
|
flat_fx = fx[valid]
|
|
flat_lin = (flat_fy * W + flat_fx).long()
|
|
|
|
j_map = torch.arange(M, device=device, dtype=torch.float32).view(M, 1, 1)
|
|
j_map = j_map.expand_as(valid)
|
|
flat_j = j_map[valid]
|
|
|
|
priority = torch.full((H * W,), -1.0, device=device)
|
|
priority.scatter_reduce_(0, flat_lin, flat_j, reduce="amax", include_self=False)
|
|
priority = priority.view(H, W).long()
|
|
|
|
has_circle = priority >= 0
|
|
frame[has_circle] = colors[priority[has_circle]]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DrawRhythmTracks Node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class DrawRhythmTracks(io.ComfyNode):
|
|
"""Generate BPM-synced sparse track point images."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="DrawRhythmTracks",
|
|
category="rhythm/motion_tracking",
|
|
display_name="Draw Rhythm Tracks",
|
|
description=(
|
|
"Generate BPM-synced sparse track point images. "
|
|
"A single point swings to the rhythm with independent "
|
|
"X/Y patterns and timing controls."
|
|
),
|
|
inputs=[
|
|
io.Float.Input(
|
|
"bpm", default=140.0, min=30.0, max=300.0, step=0.1,
|
|
tooltip="Tempo in beats per minute. One beat = one swing cycle.",
|
|
),
|
|
io.Float.Input(
|
|
"fps", default=24.0, min=1.0, max=120.0, step=0.1,
|
|
tooltip="Output frame rate.",
|
|
),
|
|
io.Float.Input(
|
|
"length", default=5.0, min=0.1, max=60.0, step=0.1,
|
|
tooltip="Output duration in seconds.",
|
|
),
|
|
io.Int.Input(
|
|
"width", default=512, min=8, max=8192, step=8,
|
|
tooltip="Output image width in pixels.",
|
|
),
|
|
io.Int.Input(
|
|
"height", default=512, min=8, max=8192, step=8,
|
|
tooltip="Output image height in pixels.",
|
|
),
|
|
io.Int.Input(
|
|
"start_x", default=256, min=0, max=8192,
|
|
tooltip="Swing center point X coordinate.",
|
|
),
|
|
io.Int.Input(
|
|
"start_y", default=128, min=0, max=8192,
|
|
tooltip="Swing center point Y coordinate.",
|
|
),
|
|
io.Combo.Input(
|
|
"pattern_x", options=["pendulum", "bounce", "linear"],
|
|
default="pendulum",
|
|
tooltip="X-axis swing pattern.",
|
|
),
|
|
io.Combo.Input(
|
|
"pattern_y", options=["pendulum", "bounce", "linear"],
|
|
default="bounce",
|
|
tooltip="Y-axis swing pattern.",
|
|
),
|
|
io.Float.Input(
|
|
"rate_x", default=0.0, min=0.0, max=1.0, step=0.01,
|
|
tooltip="Horizontal swing width as ratio of image width.",
|
|
),
|
|
io.Float.Input(
|
|
"rate_y", default=0.25, min=0.0, max=1.0, step=0.01,
|
|
tooltip="Vertical swing width as ratio of image height.",
|
|
),
|
|
io.Float.Input(
|
|
"rhythm_x", default=0.5, min=0.01, max=10.0, step=0.01,
|
|
tooltip="X-axis BPM multiplier. Effective BPM = bpm * rhythm_x.",
|
|
),
|
|
io.Float.Input(
|
|
"rhythm_y", default=1.0, min=0.01, max=10.0, step=0.01,
|
|
tooltip="Y-axis BPM multiplier. Effective BPM = bpm * rhythm_y.",
|
|
),
|
|
io.Float.Input(
|
|
"delay", default=0.0, min=-10.0, max=10.0, step=0.01,
|
|
tooltip="Time offset in seconds. Negative = earlier, positive = later.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
bpm: float,
|
|
fps: float,
|
|
length: float,
|
|
width: int,
|
|
height: int,
|
|
start_x: int,
|
|
start_y: int,
|
|
pattern_x: str,
|
|
pattern_y: str,
|
|
rate_x: float,
|
|
rate_y: float,
|
|
rhythm_x: float,
|
|
rhythm_y: float,
|
|
delay: float,
|
|
) -> io.NodeOutput:
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
num_frames = math.ceil(length * fps)
|
|
beat_period = 60.0 / bpm
|
|
|
|
# Compute per-frame coordinates
|
|
frame_indices = torch.arange(num_frames, device=device, dtype=torch.float32)
|
|
frame_times = frame_indices / fps + delay
|
|
|
|
phase_x = (frame_times / beat_period * rhythm_x) * 2 * math.pi
|
|
phase_y = (frame_times / beat_period * rhythm_y) * 2 * math.pi
|
|
|
|
pattern_fn_x = _PATTERNS[pattern_x]
|
|
pattern_fn_y = _PATTERNS[pattern_y]
|
|
|
|
offset_x = pattern_fn_x(phase_x, width, rate_x)
|
|
offset_y = pattern_fn_y(phase_y, height, rate_y)
|
|
|
|
coords_x = start_x + offset_x
|
|
coords_y = start_y + offset_y
|
|
|
|
# Build track data: [1 track, num_frames, 2]
|
|
point_xy = torch.stack([coords_x, coords_y], dim=-1).unsqueeze(0)
|
|
vis = torch.ones(1, num_frames, dtype=torch.bool, device=device)
|
|
|
|
# Rendering (ported from LTXVDrawTracks)
|
|
rw, rh, sx, sy = _render_resolution(width, height, _REF_SHORT_SIDE)
|
|
scaled_xy = point_xy.clone()
|
|
scaled_xy[..., 0] *= sx
|
|
scaled_xy[..., 1] *= sy
|
|
|
|
max_d = 2 * _MAX_RADIUS + 3
|
|
half_d = max_d // 2
|
|
offsets = torch.arange(max_d, device=device) - half_d
|
|
oy, ox = torch.meshgrid(offsets, offsets, indexing="ij")
|
|
template_dist_sq = oy.float().square() + ox.float().square()
|
|
|
|
render_frames = torch.zeros(num_frames, rh, rw, 3, device=device)
|
|
|
|
for t in range(num_frames):
|
|
tau_min = max(0, t - _MAX_TRAIL)
|
|
window = t - tau_min + 1
|
|
|
|
active_xy = scaled_xy[:, tau_min:t + 1]
|
|
active_vis = vis[:, tau_min:t + 1]
|
|
|
|
ages = torch.arange(
|
|
window - 1, -1, -1, device=device, dtype=torch.float32
|
|
)
|
|
ratios = 1.0 - ages / _MAX_TRAIL
|
|
radii = _MIN_RADIUS + (_MAX_RADIUS - _MIN_RADIUS) * ratios
|
|
colors = _age_color_batch(ratios, device)
|
|
|
|
flat_xy = active_xy.reshape(-1, 2)
|
|
flat_vis = active_vis.reshape(-1)
|
|
flat_radii = radii.unsqueeze(0).expand(1, -1).reshape(-1)
|
|
flat_colors = colors.unsqueeze(0).expand(1, -1, -1).reshape(-1, 3)
|
|
|
|
idx = flat_vis.nonzero(as_tuple=True)[0]
|
|
if idx.shape[0] == 0:
|
|
continue
|
|
|
|
pts = flat_xy[idx]
|
|
r = flat_radii[idx]
|
|
c = flat_colors[idx]
|
|
|
|
flat_ages = ages.unsqueeze(0).expand(1, -1).reshape(-1)
|
|
sort_order = flat_ages[idx].argsort(descending=True)
|
|
pts = pts[sort_order]
|
|
r = r[sort_order]
|
|
c = c[sort_order]
|
|
|
|
_rasterise_circles(
|
|
render_frames[t], pts, r, c,
|
|
template_dist_sq, half_d, max_d, rh, rw,
|
|
)
|
|
|
|
# Downscale to target resolution
|
|
out = torch.nn.functional.interpolate(
|
|
render_frames.permute(0, 3, 1, 2),
|
|
size=(height, width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
).permute(0, 2, 3, 1)
|
|
|
|
# RGB -> BGR to match IC-LoRA training data format
|
|
out = out[..., [2, 1, 0]]
|
|
return io.NodeOutput(out.half())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# BPMDetector Node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class BPMDetector(io.ComfyNode):
|
|
"""Detect BPM from audio input."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="BPMDetector",
|
|
category="rhythm/audio",
|
|
display_name="BPM Detector",
|
|
description="Detect BPM (tempo) from audio input using beat tracking.",
|
|
inputs=[
|
|
io.Audio.Input(
|
|
"audio",
|
|
tooltip="Audio input from Load Audio node.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Float.Output("bpm"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, audio) -> io.NodeOutput:
|
|
import librosa
|
|
|
|
waveform = audio["waveform"]
|
|
sample_rate = audio["sample_rate"]
|
|
|
|
# Convert to mono numpy array
|
|
# waveform shape: [1, channels, samples]
|
|
audio_np = waveform.squeeze(0).mean(dim=0).cpu().numpy()
|
|
|
|
# Detect BPM
|
|
tempo, _ = librosa.beat.beat_track(y=audio_np, sr=sample_rate)
|
|
|
|
# librosa may return an array; extract scalar
|
|
if hasattr(tempo, "__len__"):
|
|
bpm_value = float(tempo[0]) if len(tempo) > 0 else 120.0
|
|
else:
|
|
bpm_value = float(tempo)
|
|
|
|
return io.NodeOutput(bpm_value)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Node Registration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DrawRhythmTracks": DrawRhythmTracks,
|
|
"BPMDetector": BPMDetector,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DrawRhythmTracks": "Draw Rhythm Tracks",
|
|
"BPMDetector": "BPM Detector",
|
|
}
|