Apply Filter Bank

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
 

__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

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"]}, )