feat: add plotting functionality for PEQ frequency response

app.py

@@ -1,16 +1,17 @@
 import gradio as gr
 import numpy as np
+import matplotlib.pyplot as plt
 import torch
 import yaml
 import json
 import pyloudnorm as pyln
 from hydra.utils import instantiate
-from random import normalvariate
 from soxr import resample
 from functools import partial
 
 from modules.utils import chain_functions, vec2statedict, get_chunks
 from modules.fx import clip_delay_eq_Q
+from plot_utils import get_log_mags_from_eq
 
 
 title_md = "# Vocal Effects Generator"
@@ -41,11 +42,13 @@ TEMPERATURE = 0.7
 CONFIG_PATH = "presets/rt_config.yaml"
 PCA_PARAM_FILE = "presets/internal/gaussian.npz"
 INFO_PATH = "presets/internal/info.json"
+MASK_PATH = "presets/internal/feature_mask.npy"
 
 
 with open(CONFIG_PATH) as fp:
     fx_config = yaml.safe_load(fp)["model"]
 
+# Global effect
 fx = instantiate(fx_config)
 fx.eval()
 
@@ -58,6 +61,8 @@ eigvecs = np.flip(eigvecs, axis=1)[:, :75]
 U = eigvecs * np.sqrt(eigvals)
 U = torch.from_numpy(U).float()
 mean = torch.from_numpy(mean).float()
+feature_mask = torch.from_numpy(np.load(MASK_PATH))
+# Global latent variable
 z = torch.zeros(75)
 
 with open(INFO_PATH) as f:
@@ -85,11 +90,35 @@ vec2dict = partial(
         )
     ),
 )
+fx.load_state_dict(vec2dict(mean), strict=False)
 
 
 meter = pyln.Meter(44100)
 
 
+@torch.no_grad()
+def z2fx():
+    # close all figures to avoid too many open figures
+    plt.close("all")
+    x = U @ z + mean
+    # print(z)
+    fx.load_state_dict(vec2dict(x), strict=False)
+    return
+
+
+def fx2z(func):
+    @torch.no_grad()
+    def wrapper(*args, **kwargs):
+        ret = func(*args, **kwargs)
+        state_dict = fx.state_dict()
+        flattened = torch.cat([state_dict[k].flatten() for k in param_keys])
+        x = flattened[feature_mask]
+        z.copy_(U.T @ (x - mean))
+        return ret
+
+    return wrapper
+
+
 @torch.no_grad()
 def inference(audio):
     sr, y = audio
@@ -107,21 +136,6 @@ def inference(audio):
     if y.shape[1] != 1:
         y = y.mean(dim=1, keepdim=True)
 
-    # M = eigvals.shape[0]
-    # z = torch.cat(
-    #     [
-    #         torch.tensor([float(x) for x in pcs]),
-    #         (
-    #             torch.randn(M - len(pcs)) * TEMPERATURE
-    #             if randomise_rest
-    #             else torch.zeros(M - len(pcs))
-    #         ),
-    #     ]
-    # )
-    x = U @ z + mean
-    # print(z)
-
-    fx.load_state_dict(vec2dict(x), strict=False)
     fx.apply(partial(clip_delay_eq_Q, Q=0.707))
 
     rendered = fx(y).squeeze(0).T.numpy()
@@ -161,6 +175,23 @@ def model2json():
     )
 
 
+@torch.no_grad()
+def plot_eq():
+    fig, ax = plt.subplots(figsize=(8, 4))
+    w, eq_log_mags = get_log_mags_from_eq(fx[:6])
+    ax.plot(w, sum(eq_log_mags), color="black", linestyle="-")
+    for i, eq_log_mag in enumerate(eq_log_mags):
+        ax.plot(w, eq_log_mag, "k-", alpha=0.3)
+        ax.fill_between(w, eq_log_mag, 0, facecolor="gray", edgecolor="none", alpha=0.1)
+    ax.set_xlabel("Frequency (Hz)")
+    ax.set_ylabel("Magnitude (dB)")
+    ax.set_xlim(20, 20000)
+    ax.set_ylim(-40, 20)
+    ax.set_xscale("log")
+    ax.grid()
+    return fig
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         title_md,
@@ -214,17 +245,23 @@ with gr.Blocks() as demo:
             audio_output = gr.Audio(
                 type="numpy", label="Output Audio", interactive=False, loop=True
             )
-            json_output = gr.JSON(label="Effect Settings", max_height=800, open=True)
+
+            peq_plot = gr.Plot(
+                plot_eq(), label="PEQ Frequency Response", elem_id="peq-plot"
+            )
+
+    with gr.Row():
+        json_output = gr.JSON(label="Effect Settings", max_height=800, open=True)
 
     render_button.click(
-        lambda *args: (lambda x: (x, model2json()))(inference(*args)),
+        lambda *args: (lambda x: (x, model2json(), plot_eq()))(inference(*args)),
         inputs=[
             audio_input,
             # random_rest_checkbox,
         ]
         # + sliders,
         ,
-        outputs=[audio_output, json_output],
+        outputs=[audio_output, json_output, peq_plot],
     )
 
     random_button.click(
@@ -235,16 +272,27 @@ with gr.Blocks() as demo:
         #     )(normalvariate(0, 1))
         #     for _ in range(len(xs))
         # ],
-        lambda i: (lambda x: x[:NUMBER_OF_PCS].tolist() + [x[i - 1].item()])(
-            z.normal_(0, 1).clip_(SLIDER_MIN, SLIDER_MAX)
+        # lambda i: (lambda x: x[:NUMBER_OF_PCS].tolist() + [x[i - 1].item()])(
+        #     z.normal_(0, 1).clip_(SLIDER_MIN, SLIDER_MAX)
+        # ),
+        chain_functions(
+            lambda i: (z.normal_(0, 1).clip_(SLIDER_MIN, SLIDER_MAX), i),
+            lambda args: args + (z2fx(),),
+            lambda args: args[0][:NUMBER_OF_PCS].tolist()
+            + [args[0][args[1] - 1].item(), plot_eq()],
         ),
         inputs=extra_pc_dropdown,
-        outputs=sliders + [extra_slider],
+        outputs=sliders + [extra_slider, peq_plot],
     )
     reset_button.click(
-        lambda *xs: (lambda _: [0 for _ in range(len(xs))])(z.zero_()),
-        inputs=sliders + [extra_slider],
-        outputs=sliders + [extra_slider],
+        # lambda: (lambda _: [0 for _ in range(NUMBER_OF_PCS + 1)])(z.zero_()),
+        lambda: chain_functions(
+            lambda _: z.zero_(),
+            lambda _: z2fx(),
+            lambda _: [0 for _ in range(NUMBER_OF_PCS + 1)] + [plot_eq()],
+        )(None),
+        # inputs=sliders + [extra_slider],
+        outputs=sliders + [extra_slider, peq_plot],
     )
 
     def update_z(s, i):
@@ -252,12 +300,26 @@ with gr.Blocks() as demo:
         return
 
     for i, slider in enumerate(sliders):
-        slider.change(partial(update_z, i=i), inputs=slider)
-    extra_slider.change(
-        lambda _, i: update_z(_, i - 1), inputs=[extra_slider, extra_pc_dropdown]
+        slider.input(
+            chain_functions(
+                partial(update_z, i=i),
+                lambda _: z2fx(),
+                lambda _: plot_eq(),
+            ),
+            inputs=slider,
+            outputs=peq_plot,
+        )
+    extra_slider.input(
+        lambda *xs: chain_functions(
+            lambda args: update_z(args[0], args[1] - 1),
+            lambda _: z2fx(),
+            lambda _: plot_eq(),
+        )(xs),
+        inputs=[extra_slider, extra_pc_dropdown],
+        outputs=peq_plot,
     )
 
-    extra_pc_dropdown.change(
+    extra_pc_dropdown.input(
         lambda i: z[i - 1].item(),
         inputs=extra_pc_dropdown,
         outputs=extra_slider,