Add toJSON method to FX classes for JSON serialization of parameters

modules/fx.py

@@ -2,7 +2,7 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 from torch.nn.utils.parametrize import register_parametrization
-from torchcomp import ms2coef, coef2ms, db2amp
+from torchcomp import ms2coef, coef2ms, db2amp, amp2db
 from torchaudio.transforms import Spectrogram, InverseSpectrogram
 
 from typing import List, Tuple, Union, Any, Optional, Callable
@@ -72,6 +72,9 @@ class FX(nn.Module):
 
         self.params = nn.ParameterDict({k: float2param(v) for k, v in kwargs.items()})
 
+    def toJSON(self) -> dict[str, Any]:
+        return {k: v.item() for k, v in self.params.items() if v.numel() == 1}
+
 
 class SmoothingCoef(nn.Module):
     def forward(self, x):
@@ -196,6 +199,18 @@ class CompressorExpander(FX):
                 s += f"\nlookahead: {self.params.lookahead.item()} (ms)"
         return s
 
+    def toJSON(self) -> dict[str, Any]:
+        return {
+            "Attack (ms)": coef2ms(self.params.at, self.sr).item(),
+            "Release (ms)": coef2ms(self.params.rt, self.sr).item(),
+            "Average Coefficient": self.params.avg_coef.item(),
+            "Compressor Ratio": self.params.cmp_ratio.item(),
+            "Expander Ratio": self.params.exp_ratio.item(),
+            "Compressor Threshold (dB)": self.params.cmp_th.item(),
+            "Expander Threshold (dB)": self.params.exp_th.item(),
+            "Make Up (dB)": self.params.make_up.item(),
+        } | ({"Lookahead (ms)": self.params.lookahead.item()} if self.lookahead else {})
+
     def forward(self, x):
         if self.lookahead:
             lookahead_in_samples = self.params.lookahead * 0.001 * self.sr
@@ -230,6 +245,11 @@ class Panning(FX):
             s = f"pan: {self.params.pan.item() * 200 - 100}"
         return s
 
+    def toJSON(self) -> dict[str, Any]:
+        return {
+            "Pan": self.params.pan.item() * 200 - 100,
+        }
+
     def forward(self, x: torch.Tensor):
         angle = self.params.pan.view(1) * torch.pi * 0.5
         amp = torch.concat([angle.cos(), angle.sin()]).view(2, 1) * STEREO_NORM
@@ -310,6 +330,12 @@ class LowPass(FX):
             s = f"freq: {self.params.freq.item():.4f}, Q: {self.params.Q.item():.4f}"
         return s
 
+    def toJSON(self) -> dict[str, Any]:
+        return {
+            "Frequency (Hz)": self.params.freq.item(),
+            "Q": self.params.Q.item(),
+        }
+
 
 class HighPass(LowPass):
     def __init__(
@@ -363,6 +389,13 @@ class Peak(FX):
             s = f"freq: {self.params.freq.item():.4f}, gain: {self.params.gain.item():.4f}, Q: {self.params.Q.item():.4f}"
         return s
 
+    def toJSON(self) -> dict[str, Any]:
+        return {
+            "Frequency (Hz)": self.params.freq.item(),
+            "Gain (dB)": self.params.gain.item(),
+            "Q": self.params.Q.item(),
+        }
+
 
 class LowShelf(FX):
     def __init__(
@@ -394,6 +427,12 @@ class LowShelf(FX):
             s = f"freq: {self.params.freq.item():.4f}, gain: {self.params.gain.item():.4f}"
         return s
 
+    def toJSON(self) -> dict[str, Any]:
+        return {
+            "Frequency (Hz)": self.params.freq.item(),
+            "Gain (dB)": self.params.gain.item(),
+        }
+
 
 class HighShelf(LowShelf):
     def __init__(
@@ -611,6 +650,15 @@ class Delay(FX):
             )
         return s
 
+    def toJSON(self) -> dict[str, Any]:
+        return {
+            "Delay (ms)": self.params.delay.item(),
+            "Feedback (dB)": self.params.feedback.log10().mul(20).item(),
+            "Gain (dB)": self.params.gain.log10().mul(20).item(),
+            "Odd delays": self.odd_pan.toJSON(),
+            "Even delays": self.even_pan.toJSON(),
+        }
+
 
 class SurrogateDelay(Delay):
     def __init__(self, *args, dropout=0.5, straight_through=False, **kwargs):
@@ -992,3 +1040,17 @@ class FDN(FX):
             F.pad(x, (self.ir_length - 1, 0)),
             h.flip(-1),
         )
+
+    def toJSON(self) -> dict[str, Any]:
+        return {
+            "T60 (s)": {
+                f"{f:.2f} Hz": g.item()
+                for f, g in zip(
+                    torch.linspace(0, 22050, self.params.gamma.numel()),
+                    -60 * self.delays.min() / amp2db(self.params.gamma) / 44100,
+                )
+            },
+            "Gain (dB, approx)": amp2db(
+                torch.linalg.norm(self.params.b) * torch.linalg.norm(self.params.c)
+            ).item(),
+        }