commit 35e0209ac77c5c92cefede5a41cab1892908be3f Author: m1_4k Date: Fri Feb 13 01:49:07 2026 +0900 Initial release: Audio to Spectrogram ComfyUI custom node - AudioToSpectrogram node: AUDIO input → IMAGE output - Supports Mel and Linear (STFT) spectrograms with dB scaling - 7 colormaps, configurable FFT params and output resolution - torchaudio-based transforms with NumPy STFT fallback Co-Authored-By: Claude Sonnet 4.5 diff --git a/README.md b/README.md new file mode 100644 index 0000000..7747182 --- /dev/null +++ b/README.md @@ -0,0 +1,81 @@ +# ComfyUI Audio to Spectrogram + +A ComfyUI custom node that converts an audio input into a spectrogram image. + +## Features + +- Supports both **Mel** and **Linear** (STFT) spectrograms +- dB-scaled output (top 80 dB dynamic range) +- 7 built-in colormaps: `inferno`, `magma`, `viridis`, `plasma`, `hot`, `cool`, `gray` +- Configurable frequency range, FFT parameters, and output resolution +- Optional axes, colorbar, and title overlay +- Uses `torchaudio` when available; falls back to NumPy STFT automatically + +## Installation + +Copy this directory into ComfyUI's `custom_nodes` folder: + +``` +ComfyUI/ +└── custom_nodes/ + └── CustomNode_Bispectrum/ + ├── __init__.py + ├── nodes.py + └── pyproject.toml +``` + +Then restart ComfyUI. The node will appear under the **audio** category. + +## Dependencies + +| Package | Required | Notes | +|---|---|---| +| `torch` | Yes | Comes with ComfyUI | +| `matplotlib` | Yes | Used for rendering | +| `Pillow` | Yes | Image conversion | +| `numpy` | Yes | Array operations | +| `torchaudio` | Recommended | Faster STFT / Mel transforms | + +Install recommended dependencies: + +```bash +pip install torchaudio matplotlib Pillow +``` + +## Node: Audio to Spectrogram + +**Category:** `audio` + +### Inputs + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `audio` | AUDIO | — | ComfyUI audio input | +| `use_mel` | BOOLEAN | `True` | `True` = Mel spectrogram, `False` = Linear spectrogram | +| `n_fft` | INT | `2048` | FFT window size (256–8192) | +| `hop_length` | INT | `512` | Frame hop size in samples (64–2048) | +| `win_length` | INT | `2048` | Analysis window length (256–8192) | +| `n_mels` | INT | `128` | Number of Mel filter banks (Mel mode only) | +| `fmin` | FLOAT | `0.0` | Minimum frequency in Hz | +| `fmax` | FLOAT | `8000.0` | Maximum frequency in Hz (set to `0` for Nyquist) | +| `colormap` | COMBO | `inferno` | Matplotlib colormap | +| `width` | INT | `1024` | Output image width in pixels | +| `height` | INT | `512` | Output image height in pixels | +| `show_axes` | BOOLEAN | `True` | Show time/frequency axes, colorbar, and title | + +### Output + +| Name | Type | Shape | Description | +|---|---|---|---| +| `spectrogram` | IMAGE | `[1, H, W, 3]` | Float32 RGB image, values in `[0, 1]` | + +## Example Workflow + +1. Add a **Load Audio** node and load any audio file. +2. Connect its output to the `audio` input of **Audio to Spectrogram**. +3. Connect `spectrogram` to a **Preview Image** or **Save Image** node. +4. Adjust `use_mel`, `fmax`, and `colormap` to taste. + +## License + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..5a49e65 --- /dev/null +++ b/nodes.py @@ -0,0 +1,203 @@ +import torch +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.colors as mcolors +from io import BytesIO +from PIL import Image + + +COLORMAPS = ["inferno", "magma", "viridis", "plasma", "hot", "cool", "gray"] + + +def _to_mono(waveform: torch.Tensor) -> torch.Tensor: + """waveform: [1, C, T] → [T] mono float tensor""" + audio = waveform.squeeze(0) # [C, T] + if audio.ndim == 2: + audio = audio.mean(dim=0) # mono + return audio.float() + + +def _compute_spectrogram( + audio: torch.Tensor, + sample_rate: int, + n_fft: int, + hop_length: int, + win_length: int, + use_mel: bool, + n_mels: int, + fmin: float, + fmax: float, +) -> np.ndarray: + """STFT 또는 Mel-spectrogram 계산 후 dB 스케일 numpy 배열 반환 [freq, time]""" + try: + import torchaudio + import torchaudio.transforms as T + + if use_mel: + transform = T.MelSpectrogram( + sample_rate=sample_rate, + n_fft=n_fft, + hop_length=hop_length, + win_length=win_length, + n_mels=n_mels, + f_min=fmin, + f_max=fmax if fmax > 0 else sample_rate / 2, + power=2.0, + ) + spec = transform(audio) # [n_mels, T] + else: + transform = T.Spectrogram( + n_fft=n_fft, + hop_length=hop_length, + win_length=win_length, + power=2.0, + ) + spec = transform(audio) # [F, T] + + to_db = T.AmplitudeToDB(stype="power", top_db=80) + spec_db = to_db(spec).numpy() + + except ImportError: + # torchaudio 없을 경우 numpy STFT fallback + audio_np = audio.numpy() + window = np.hanning(win_length) + step = hop_length + frames = [] + for start in range(0, max(len(audio_np) - win_length + 1, 1), step): + frame = audio_np[start: start + win_length] + if len(frame) < win_length: + frame = np.pad(frame, (0, win_length - len(frame))) + frames.append(np.fft.rfft(frame * window)) + spec = np.abs(np.array(frames)).T ** 2 # [F, T] + spec_db = 10 * np.log10(np.maximum(spec, 1e-10)) + spec_db = np.clip(spec_db, spec_db.max() - 80, spec_db.max()) + + return spec_db + + +def _render_spectrogram( + spec_db: np.ndarray, + sample_rate: int, + hop_length: int, + use_mel: bool, + fmin: float, + fmax: float, + colormap: str, + width: int, + height: int, + show_axes: bool, +) -> torch.Tensor: + """스펙트로그램 배열을 matplotlib으로 렌더링 → [1, H, W, 3] float tensor""" + dpi = 100 + fig_w = width / dpi + fig_h = height / dpi + + fig, ax = plt.subplots(figsize=(fig_w, fig_h), dpi=dpi) + + n_frames = spec_db.shape[1] + duration = n_frames * hop_length / sample_rate + n_freqs = spec_db.shape[0] + freq_max = fmax if (use_mel and fmax > 0) else (sample_rate / 2) + + extent = [0, duration, fmin, freq_max] + im = ax.imshow( + spec_db, + origin="lower", + aspect="auto", + extent=extent, + cmap=colormap, + ) + + if show_axes: + ax.set_xlabel("Time (s)") + ax.set_ylabel("Mel Frequency (Hz)" if use_mel else "Frequency (Hz)") + plt.colorbar(im, ax=ax, format="%+2.0f dB", pad=0.02) + plt.title("Mel Spectrogram" if use_mel else "Spectrogram") + plt.tight_layout() + else: + ax.axis("off") + plt.subplots_adjust(left=0, right=1, top=1, bottom=0) + + buf = BytesIO() + fig.savefig(buf, format="png", dpi=dpi) + plt.close(fig) + buf.seek(0) + + img = Image.open(buf).convert("RGB").resize((width, height), Image.LANCZOS) + img_np = np.array(img).astype(np.float32) / 255.0 + return torch.from_numpy(img_np).unsqueeze(0) # [1, H, W, 3] + + +class AudioToSpectrogram: + """ + ComfyUI 커스텀 노드: 오디오 → 스펙트로그램 이미지 + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO",), + "use_mel": ("BOOLEAN", {"default": True, "label_on": "Mel", "label_off": "Linear"}), + "n_fft": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 256}), + "hop_length": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}), + "win_length": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 256}), + "n_mels": ("INT", {"default": 128, "min": 32, "max": 512, "step": 16}), + "fmin": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 8000.0, "step": 10.0}), + "fmax": ("FLOAT", {"default": 8000.0, "min": 0.0, "max": 24000.0, "step": 100.0, + "tooltip": "0 = Nyquist (sample_rate / 2)"}), + "colormap": (COLORMAPS, {"default": "inferno"}), + "width": ("INT", {"default": 1024, "min": 256, "max": 4096, "step": 64}), + "height": ("INT", {"default": 512, "min": 128, "max": 4096, "step": 64}), + "show_axes": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("spectrogram",) + FUNCTION = "process" + CATEGORY = "audio" + + def process( + self, + audio, + use_mel, + n_fft, + hop_length, + win_length, + n_mels, + fmin, + fmax, + colormap, + width, + height, + show_axes, + ): + waveform = audio["waveform"] # [1, C, T] + sample_rate = audio["sample_rate"] + + mono = _to_mono(waveform) + + spec_db = _compute_spectrogram( + mono, sample_rate, n_fft, hop_length, win_length, + use_mel, n_mels, fmin, fmax, + ) + + image = _render_spectrogram( + spec_db, sample_rate, hop_length, + use_mel, fmin, fmax, + colormap, width, height, show_axes, + ) + + return (image,) + + +NODE_CLASS_MAPPINGS = { + "AudioToSpectrogram": AudioToSpectrogram, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AudioToSpectrogram": "Audio to Spectrogram", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8d3aa9f --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "comfyui-audio-spectrogram" +version = "1.0.0" +description = "ComfyUI custom node: Audio → Spectrogram image" +license = { file = "LICENSE" } +dependencies = ["matplotlib", "Pillow", "numpy", "torch"] + +[project.optional-dependencies] +recommended = ["torchaudio"] + +[tool.comfy] +PublisherId = "local" +DisplayName = "Audio to Spectrogram" +Icon = ""