414 lines
14 KiB
Python
414 lines
14 KiB
Python
import math
|
|
import os
|
|
|
|
import torch
|
|
from comfy_api.latest import io
|
|
|
|
_PACKAGE_ROOT = os.path.dirname(os.path.abspath(__file__))
|
|
_AUDIO_DIR = os.path.join(_PACKAGE_ROOT, "audio")
|
|
|
|
|
|
def _list_audio_wavs() -> list[str]:
|
|
"""Scan audio/ directory for .wav files."""
|
|
if not os.path.isdir(_AUDIO_DIR):
|
|
return []
|
|
return sorted(
|
|
f for f in os.listdir(_AUDIO_DIR)
|
|
if f.lower().endswith(".wav") and os.path.isfile(os.path.join(_AUDIO_DIR, f))
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rendering constants
|
|
# ---------------------------------------------------------------------------
|
|
_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 (t ** 2 - 0.5) * 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
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _trail_color(ratio: float, device: torch.device) -> torch.Tensor:
|
|
"""Return an RGB color for an old-to-new trail ratio."""
|
|
red = ratio
|
|
green = max(0.0, 1.0 - abs(ratio - 0.5) * 2.0)
|
|
blue = 1.0 - ratio
|
|
return torch.tensor([red, green, blue], device=device, dtype=torch.float32)
|
|
|
|
|
|
def _draw_filled_circle(
|
|
frame: torch.Tensor,
|
|
x: float,
|
|
y: float,
|
|
radius: float,
|
|
color: torch.Tensor,
|
|
) -> None:
|
|
"""Draw a filled circle into an RGB frame tensor."""
|
|
height, width = frame.shape[:2]
|
|
cx = int(round(x))
|
|
cy = int(round(y))
|
|
r = max(1, int(round(radius)))
|
|
|
|
x0 = max(0, cx - r)
|
|
x1 = min(width, cx + r + 1)
|
|
y0 = max(0, cy - r)
|
|
y1 = min(height, cy + r + 1)
|
|
if x0 >= x1 or y0 >= y1:
|
|
return
|
|
|
|
yy = torch.arange(y0, y1, device=frame.device).view(-1, 1)
|
|
xx = torch.arange(x0, x1, device=frame.device).view(1, -1)
|
|
mask = (xx - cx).square() + (yy - cy).square() <= r * r
|
|
patch = frame[y0:y1, x0:x1]
|
|
patch[mask] = color
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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.5, 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.25, 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
|
|
|
|
coords = torch.stack([coords_x, coords_y], dim=-1)
|
|
render_frames = torch.zeros(num_frames, height, width, 3, device=device)
|
|
radius_scale = min(width, height) / _REF_SHORT_SIDE
|
|
|
|
for t in range(num_frames):
|
|
start = max(0, t - _MAX_TRAIL)
|
|
for source_frame in range(start, t + 1):
|
|
age = t - source_frame
|
|
ratio = 1.0 - age / _MAX_TRAIL
|
|
radius = (
|
|
_MIN_RADIUS + (_MAX_RADIUS - _MIN_RADIUS) * ratio
|
|
) * radius_scale
|
|
color = _trail_color(ratio, device)
|
|
x, y = coords[source_frame].tolist()
|
|
_draw_filled_circle(render_frames[t], x, y, radius, color)
|
|
|
|
# RGB -> BGR to match IC-LoRA training data format
|
|
out = render_frames[..., [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)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RhythmGenerator Node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _load_wav(path: str) -> tuple[torch.Tensor, int]:
|
|
"""Load a wav file into a [channels, samples] float32 tensor."""
|
|
import torchaudio
|
|
|
|
waveform, sample_rate = torchaudio.load(path)
|
|
if waveform.dtype != torch.float32:
|
|
waveform = waveform.to(torch.float32)
|
|
return waveform, int(sample_rate)
|
|
|
|
|
|
class RhythmGenerator(io.ComfyNode):
|
|
"""Generate BPM-synced audio by placing a wav sample on every beat."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
wav_options = _list_audio_wavs() or ["kick.wav"]
|
|
default_wav = "kick.wav" if "kick.wav" in wav_options else wav_options[0]
|
|
return io.Schema(
|
|
node_id="RhythmGenerator",
|
|
category="rhythm/audio",
|
|
display_name="Rhythm Generator",
|
|
description=(
|
|
"Generate BPM-synced audio by repeating a wav sample on each beat. "
|
|
"Monophonic playback: each beat cuts the previous sample."
|
|
),
|
|
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 sample trigger.",
|
|
),
|
|
io.Combo.Input(
|
|
"wav", options=wav_options, default=default_wav,
|
|
tooltip="Wav sample file in the package's audio/ directory.",
|
|
),
|
|
io.Float.Input(
|
|
"length", default=5.0, min=0.1, max=60.0, step=0.1,
|
|
tooltip="Output duration in seconds.",
|
|
),
|
|
io.Float.Input(
|
|
"delay", default=0.0, min=-10.0, max=10.0, step=0.01,
|
|
tooltip="Time offset in seconds for the first beat. "
|
|
"Negative = earlier, positive = later.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Audio.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
bpm: float,
|
|
wav: str,
|
|
length: float,
|
|
delay: float,
|
|
) -> io.NodeOutput:
|
|
wav_path = os.path.join(_AUDIO_DIR, wav)
|
|
if not os.path.isfile(wav_path):
|
|
raise FileNotFoundError(f"Wav file not found: {wav_path}")
|
|
|
|
source_waveform, source_sr = _load_wav(wav_path)
|
|
channels, source_samples = source_waveform.shape
|
|
|
|
output_samples = int(math.floor(length * source_sr))
|
|
if output_samples <= 0:
|
|
empty = torch.zeros(1, channels, 0, dtype=torch.float32)
|
|
return io.NodeOutput({"waveform": empty, "sample_rate": source_sr})
|
|
|
|
beat_period_samples = max(1, int(round(60.0 / bpm * source_sr)))
|
|
delay_samples = int(round(delay * source_sr))
|
|
|
|
output = torch.zeros(channels, output_samples, dtype=torch.float32)
|
|
|
|
beat_idx = 0
|
|
while True:
|
|
start = delay_samples + beat_idx * beat_period_samples
|
|
if start >= output_samples:
|
|
break
|
|
|
|
# Monophonic: cut at next beat. Truncate at source end and output end.
|
|
copy_len = min(source_samples, beat_period_samples, output_samples - start)
|
|
|
|
# For negative-delay beats: absorb the pre-zero portion as a source offset.
|
|
src_offset = max(0, -start)
|
|
dst_offset = max(0, start)
|
|
n = copy_len - src_offset
|
|
if n > 0:
|
|
output[:, dst_offset:dst_offset + n] = \
|
|
source_waveform[:, src_offset:src_offset + n]
|
|
|
|
beat_idx += 1
|
|
|
|
audio = {
|
|
"waveform": output.unsqueeze(0),
|
|
"sample_rate": source_sr,
|
|
}
|
|
return io.NodeOutput(audio)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Node Registration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DrawRhythmTracks": DrawRhythmTracks,
|
|
"BPMDetector": BPMDetector,
|
|
"RhythmGenerator": RhythmGenerator,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DrawRhythmTracks": "Draw Rhythm Tracks",
|
|
"BPMDetector": "BPM Detector",
|
|
"RhythmGenerator": "Rhythm Generator",
|
|
}
|