diff --git a/__init__.py b/__init__.py index e509c4e..054207b 100644 --- a/__init__.py +++ b/__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", diff --git a/spect_nodes.py b/spect_nodes.py index 6209194..f1bd5c4 100644 --- a/spect_nodes.py +++ b/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):