nodes/mcDandy-more_math
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

run ruff

Browse Source
This commit is contained in:
mcDandy
2026-01-08 14:14:44 +01:00
parent 01cbdf2e26
commit f273d74ae6
28 changed files with 10849 additions and 2576 deletions
Download Patch File Download Diff File Expand all files Collapse all files
more_math/AudioToSpectrogramNode.py
+21 -14
View File
@@ -1,7 +1,8 @@
from comfy_api.latest import io
import torch
windows = {'bartlet':torch.bartlett_window, 'blackman':torch.blackman_window, 'hamming':torch.hamming_window,'hann':torch.hann_window}
windows = {"bartlet": torch.bartlett_window, "blackman": torch.blackman_window, "hamming": torch.hamming_window, "hann": torch.hann_window}
class AudioToSpectrogram(io.ComfyNode):
"""
@@ -9,6 +10,7 @@ class AudioToSpectrogram(io.ComfyNode):
Red = Real, Green = logarithm of value (just so it looks good), Blue = Imaginary.
Each audio channel is stacked vertically.
"""
@classmethod
def define_schema(cls) -> io.Schema:
return io.Schema(
@@ -19,7 +21,9 @@ class AudioToSpectrogram(io.ComfyNode):
io.Audio.Input(id="audio", tooltip="Input audio"),
io.Int.Input(id="window_length", default=1024, min=16, max=4096, tooltip="Window length in samples"),
io.Int.Input(id="hop_length", default=256, min=1, max=4096, tooltip="Stride of the window (hop length) in samples"),
io.Int.Input(id="bucket_count", default=513, min=2, max=4096, tooltip="Number of frequency buckets (determines resolution)"),
io.Int.Input(
id="bucket_count", default=513, min=2, max=4096, tooltip="Number of frequency buckets (determines resolution)"
),
io.Combo.Input(id="window_type", default="hann", options=list(windows.keys()), tooltip="Type of window function to apply"),
],
outputs=[
@@ -30,26 +34,29 @@ class AudioToSpectrogram(io.ComfyNode):
)
@classmethod
def execute(cls, audio, window_length, hop_length, bucket_count,window_type):
waveform = audio['waveform']
sample_rate = audio['sample_rate']
def execute(cls, audio, window_length, hop_length, bucket_count, window_type):
waveform = audio["waveform"]
sample_rate = audio["sample_rate"]
B, C, S = waveform.shape
n_fft = (bucket_count - 1) * 2
flat_waveform = waveform.reshape(B * C, S)
window = windows[window_type](window_length, device=waveform.device)
stft_out = torch.stft(
flat_waveform, n_fft=n_fft, hop_length=hop_length, win_length=window_length,
window=window, center=True, pad_mode='reflect', normalized=False,
onesided=True, return_complex=True
flat_waveform,
n_fft=n_fft,
hop_length=hop_length,
win_length=window_length,
window=window,
center=True,
pad_mode="reflect",
normalized=False,
onesided=True,
return_complex=True,
)
F, T = stft_out.shape[1], stft_out.shape[2]
stft_out = stft_out.reshape(B, C, F, T).reshape(B, C * F, T)
image = torch.stack([
stft_out.real,
torch.log1p(torch.abs(stft_out)),
stft_out.imag
], dim=-1)
image = torch.stack([stft_out.real, torch.log1p(torch.abs(stft_out)), stft_out.imag], dim=-1)
print(image.type())
return (image, C, sample_rate)
return (image, C, sample_rate)