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 <noreply@anthropic.com>
2026-02-13 01:49:07 +09:00
commit 35e0209ac7
4 changed files with 301 additions and 0 deletions
@@ -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
@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
@@ -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",
+}
@@ -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 = ""