Griffin Lim

__init__.py

@@ -4,11 +4,12 @@ from .io_nodes import LoadAudio
 from .plot_nodes import PlotWaveform
 from .spect_nodes import (
     Spectrogram, InverseSpectrogram,
-    PowerSpectrogram,
+    PowerSpectrogram, GriffinLim,
     MelSpectrogram
 )
 
 NODE_CLASS_MAPPINGS = {
+    "ComfyAudio.GriffinLim": GriffinLim,
     "ComfyAudio.InverseSpectrogram": InverseSpectrogram,
     "ComfyAudio.LoadAudio": LoadAudio,
     "ComfyAudio.MelSpectrogram": MelSpectrogram,
@@ -18,6 +19,7 @@ NODE_CLASS_MAPPINGS = {
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
+    "ComfyAudio.GriffinLim": "Griffin Lim",
     "ComfyAudio.InverseSpectrogram": "Inverse Spectrogram",
     "ComfyAudio.LoadAudio": "Load Audio From Path",
     "ComfyAudio.MelSpectrogram": "Mel-scale Spectrogram",

spect_nodes.py

@@ -87,6 +87,31 @@ class PowerSpectrogram:
                  "hop_length": hop_length}, )
 
 
+class GriffinLim:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"spectrogram": ("POWER_SPECT",)},
+                "optional": {"n_iter": ("INT", {"default": 32, "min": 1, "max": 2**32}),
+                             "momentum": ("FLOAT", {"default": 0.99, "min": 0.0, "max": 1.0, "step": 0.01})}}
+
+    CATEGORY = "AudioProcessing/features"
+
+    RETURN_TYPES = ("AUDIO", )
+    FUNCTION = "griffin_lim"
+
+    def griffin_lim(self, spectrogram, n_iter=32, momentum=0.99):
+        waveform = T.GriffinLim(n_fft=spectrogram["n_fft"],
+                                n_iter=n_iter,
+                                win_length=spectrogram["win_length"],
+                                hop_length=spectrogram["hop_length"],
+                                power=2,
+                                momentum=momentum)(spectrogram["spectrogram"])
+        if len(waveform.shape) == 1:
+            waveform = waveform.unsqueeze(0)  # 1 channel (mono) audio
+        return ({"waveform": waveform.unsqueeze(0),
+                 "sample_rate": spectrogram["sample_rate"]}, )
+
+
 class MelSpectrogram:
     @classmethod
     def INPUT_TYPES(s):