97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
# Copyright (C) 2024 Reece H. Dunn. SPDX-License-Identifier: GPL-3
|
|
|
|
import torchaudio.transforms as T
|
|
|
|
|
|
SPECTROGRAM_TYPES = ("complex", "magnitude", "power")
|
|
|
|
SPECTROGRAM_TYPE_TO_POWER = {
|
|
"complex": None,
|
|
"magnitude": 1,
|
|
"power": 2,
|
|
}
|
|
|
|
|
|
class Spectrogram:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"audio": ("AUDIO",)},
|
|
"optional": {"stype": (SPECTROGRAM_TYPES, {"default": "power"}),
|
|
"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 = "audio processing/spectrogram"
|
|
|
|
RETURN_TYPES = ("SPECT", )
|
|
FUNCTION = "spectrogram"
|
|
|
|
def spectrogram(self,
|
|
audio,
|
|
stype="power",
|
|
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(0, -1).squeeze()
|
|
spectrogram = T.Spectrogram(n_fft=n_fft,
|
|
win_length=win_length,
|
|
hop_length=hop_length,
|
|
power=SPECTROGRAM_TYPE_TO_POWER[stype])(waveform)
|
|
return ({"spectrogram": spectrogram,
|
|
"sample_rate": audio["sample_rate"],
|
|
"stype": stype,
|
|
"n_fft": n_fft,
|
|
"win_length": win_length,
|
|
"hop_length": hop_length}, )
|
|
|
|
|
|
class InverseSpectrogram:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"spectrogram": ("SPECT",)}}
|
|
|
|
CATEGORY = "audio processing/spectrogram"
|
|
|
|
RETURN_TYPES = ("AUDIO", )
|
|
FUNCTION = "inverse"
|
|
|
|
def inverse(self, spectrogram):
|
|
if spectrogram["stype"] != "complex":
|
|
raise Exception(f"The spectrogram is not a complex-valued spectrogram: {spectrogram['stype']}")
|
|
waveform = T.InverseSpectrogram(n_fft=spectrogram["n_fft"],
|
|
win_length=spectrogram["win_length"],
|
|
hop_length=spectrogram["hop_length"])(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 GriffinLim:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"spectrogram": ("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 = "audio processing/spectrogram"
|
|
|
|
RETURN_TYPES = ("AUDIO", )
|
|
FUNCTION = "griffin_lim"
|
|
|
|
def griffin_lim(self, spectrogram, n_iter=32, momentum=0.99):
|
|
if spectrogram["stype"] == "complex":
|
|
raise Exception(f"The spectrogram is not a real-valued spectrogram: {spectrogram['stype']}")
|
|
waveform = T.GriffinLim(n_fft=spectrogram["n_fft"],
|
|
n_iter=n_iter,
|
|
win_length=spectrogram["win_length"],
|
|
hop_length=spectrogram["hop_length"],
|
|
power=SPECTROGRAM_TYPE_TO_POWER[spectrogram["stype"]],
|
|
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"]}, )
|