diff --git a/README.md b/README.md index c41bfb9..1d5c401 100644 --- a/README.md +++ b/README.md @@ -18,9 +18,10 @@ ## Filter Banks -| Name | Description | -|-----------------------|---------------------------------| -| Mel-scale Filter Bank | Create a mel-scale filter bank. | +| Name | Description | +|-----------------------|-----------------------------------------------------| +| Mel-scale Filter Bank | Create a mel-scale filter bank. | +| Apply Filter Bank | Apply the specified filter bank to the spectrogram. | ## Plotting diff --git a/__init__.py b/__init__.py index 9a7da26..a6ec169 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,6 @@ # Copyright (C) 2024 Reece H. Dunn. SPDX-License-Identifier: GPL-3 -from .fb_nodes import MelScaleFilterBank +from .fb_nodes import (MelScaleFilterBank, ApplyFilterBank) from .io_nodes import LoadAudio from .plot_nodes import (PlotWaveform, PlotSpectrogram, PlotFilterBank) from .spect_nodes import ( @@ -12,6 +12,7 @@ from .spect_nodes import ( NODE_CLASS_MAPPINGS = { # Filter Banks "ComfyAudio.MelScaleFilterBank": MelScaleFilterBank, + "ComfyAudio.ApplyFilterBank": ApplyFilterBank, # I/O "ComfyAudio.LoadAudio": LoadAudio, # Plotting @@ -28,6 +29,7 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { # Filter Banks "ComfyAudio.MelScaleFilterBank": "Mel-scale Filter Bank", + "ComfyAudio.ApplyFilterBank": "Apply Filter Bank", # I/O "ComfyAudio.LoadAudio": "Load Audio From Path", # Plotting diff --git a/fb_nodes.py b/fb_nodes.py index 3926885..cab7419 100644 --- a/fb_nodes.py +++ b/fb_nodes.py @@ -1,5 +1,6 @@ # Copyright (C) 2024 Reece H. Dunn. SPDX-License-Identifier: GPL-3 +import torch import torchaudio.functional as F @@ -34,3 +35,26 @@ class MelScaleFilterBank: "f_min": f_min, "f_max": f_max, "n_mels": n_mels}, ) + + +class ApplyFilterBank: + @classmethod + def INPUT_TYPES(s): + return {"required": {"spectrogram": ("SPECT",), + "fbank": ("FILTER_BANK",)}} + + CATEGORY = "AudioProcessing/filterbank" + + RETURN_TYPES = ("SPECT", ) + FUNCTION = "filter" + + def filter(self, spectrogram, fbank): + spect = spectrogram["spectrogram"] + fb = fbank["fbank"] + filtered = torch.matmul(spect.transpose(-1, -2), fb).transpose(-1, -2) + return ({"spectrogram": filtered, + "sample_rate": spectrogram["sample_rate"], + "stype": spectrogram["stype"], + "n_fft": spectrogram["n_fft"], + "win_length": spectrogram["win_length"], + "hop_length": spectrogram["hop_length"]}, )