From 6f3f29f384ffafc8a8dbb9e4aa7cba07cc6e6b7d Mon Sep 17 00:00:00 2001
From: "Reece H. Dunn" <msclrhd@gmail.com>
Date: Fri, 26 Jul 2024 21:34:01 +0100
Subject: [PATCH] Spectrogram

---
 __init__.py    |  3 +++
 spect_nodes.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 spect_nodes.py

diff --git a/__init__.py b/__init__.py
index 963f924..fc26f45 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,13 +2,16 @@
 
 from .io_nodes import LoadAudio
 from .plot_nodes import PlotWaveform
+from .spect_nodes import Spectrogram
 
 NODE_CLASS_MAPPINGS = {
     "ComfyAudio.LoadAudio": LoadAudio,
     "ComfyAudio.PlotWaveform": PlotWaveform,
+    "ComfyAudio.Spectrogram": Spectrogram,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ComfyAudio.LoadAudio": "Load Audio From Path",
     "ComfyAudio.PlotWaveform": "Plot Waveform",
+    "ComfyAudio.Spectrogram": "Spectrogram",
 }
diff --git a/spect_nodes.py b/spect_nodes.py
new file mode 100644
index 0000000..4d8a133
--- /dev/null
+++ b/spect_nodes.py
@@ -0,0 +1,35 @@
+# Copyright (C) 2024 Reece H. Dunn. SPDX-License-Identifier: GPL-3
+
+import torchaudio.transforms as T
+
+
+class Spectrogram:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"audio": ("AUDIO",)},
+                "optional": {"n_fft": ("INT", {"default": 400, "min": 1, "max": 2**32}),
+                             "win_length": ("INT", {"default": -1, "min": -1, "max": 2**32}),
+                             "hop_length": ("INT", {"default": -1, "min": -1, "max": 2**32})}}
+
+    CATEGORY = "AudioProcessing/features"
+
+    RETURN_TYPES = ("SPECT", )
+    FUNCTION = "spectrogram"
+
+    def spectrogram(self,
+                    audio,
+                    n_fft=400,
+                    win_length=-1,
+                    hop_length=-1):
+        if win_length == -1: win_length = None
+        if hop_length == -1: hop_length = None
+        waveform = audio["waveform"].movedim(1, -1).squeeze()
+        spectrogram = T.Spectrogram(n_fft=n_fft,
+                                    win_length=win_length,
+                                    hop_length=hop_length,
+                                    power=None)(waveform)
+        return ({"spectrogram": spectrogram,
+                 "sample_rate": audio["sample_rate"],
+                 "n_fft": n_fft,
+                 "win_length": win_length,
+                 "hop_length": hop_length}, )