feat: adjustable PEQ

app.py

@@ -37,7 +37,7 @@ For example:
 
 SLIDER_MAX = 3
 SLIDER_MIN = -3
-NUMBER_OF_PCS = 10
+NUMBER_OF_PCS = 4
 TEMPERATURE = 0.7
 CONFIG_PATH = "presets/rt_config.yaml"
 PCA_PARAM_FILE = "presets/internal/gaussian.npz"
@@ -56,14 +56,14 @@ pca_params = np.load(PCA_PARAM_FILE)
 mean = pca_params["mean"]
 cov = pca_params["cov"]
 eigvals, eigvecs = np.linalg.eigh(cov)
-eigvals = np.flip(eigvals, axis=0)[:75]
-eigvecs = np.flip(eigvecs, axis=1)[:, :75]
+eigvals = np.flip(eigvals, axis=0)
+eigvecs = np.flip(eigvecs, axis=1)
 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)
+z = torch.zeros_like(mean)
 
 with open(INFO_PATH) as f:
     info = json.load(f)
@@ -107,17 +107,13 @@ def z2fx():
     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 fx2z():
+    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
 
 
 @torch.no_grad()
@@ -166,12 +162,15 @@ def model2json():
         },
         "Cross Send (dB)": fx[7].params.sends_0.log10().mul(20).item(),
     }
-    return json.dumps(
-        {
-            "Direct": results,
-            "Sends": spatial_fx,
-        }
+    replace_neg_inf = lambda d: (
+        {k: (replace_neg_inf(v) if v != -np.inf else -1e500) for k, v in d.items()}
+        if isinstance(d, dict)
+        else d
     )
+    return {
+        "Direct": results,
+        "Sends": spatial_fx,
+    }
 
 
 @torch.no_grad()
@@ -283,7 +282,7 @@ def plot_t60():
     gamma = fdn.params.gamma.squeeze().numpy()
     delays = fdn.delays.numpy()
     w = np.linspace(0, 22050, gamma.size)
-    t60 = -60 / (20 * np.log10(gamma) / np.min(delays)) / 44100
+    t60 = -60 / (20 * np.log10(gamma + 1e-10) / np.min(delays)) / 44100
     ax.plot(w, t60, color="black", linestyle="-")
     ax.set_xlabel("Frequency (Hz)")
     ax.set_ylabel("T60 (s)")
@@ -294,6 +293,15 @@ def plot_t60():
     return fig
 
 
+@torch.no_grad()
+def upatePEQ(eq, attr_name, value):
+    match type(getattr(eq.params, attr_name)):
+        case torch.nn.Parameter:
+            getattr(eq.params, attr_name).data.copy_(value)
+        case _:
+            setattr(eq.params, attr_name, torch.tensor(value))
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         title_md,
@@ -328,10 +336,43 @@ with gr.Blocks() as demo:
             #     value=False,
             #     elem_id="randomise-checkbox",
             # )
-            sliders = get_important_pcs(NUMBER_OF_PCS, value=0)
+            # sliders = get_important_pcs(NUMBER_OF_PCS, value=0)
+            with gr.Row():
+                s1 = gr.Slider(
+                    minimum=SLIDER_MIN,
+                    maximum=SLIDER_MAX,
+                    label="PC 1",
+                    value=0,
+                    interactive=True,
+                )
+                s2 = gr.Slider(
+                    minimum=SLIDER_MIN,
+                    maximum=SLIDER_MAX,
+                    label="PC 2",
+                    value=0,
+                    interactive=True,
+                )
+
+            with gr.Row():
+                s3 = gr.Slider(
+                    minimum=SLIDER_MIN,
+                    maximum=SLIDER_MAX,
+                    label="PC 3",
+                    value=0,
+                    interactive=True,
+                )
+                s4 = gr.Slider(
+                    minimum=SLIDER_MIN,
+                    maximum=SLIDER_MAX,
+                    label="PC 4",
+                    value=0,
+                    interactive=True,
+                )
+
+            sliders = [s1, s2, s3, s4]
 
             extra_pc_dropdown = gr.Dropdown(
-                list(range(NUMBER_OF_PCS + 1, 76)),
+                list(range(NUMBER_OF_PCS + 1, mean.numel())),
                 label=f"PC > {NUMBER_OF_PCS}",
                 info="Select which extra PC to adjust",
                 interactive=True,
@@ -348,9 +389,85 @@ with gr.Blocks() as demo:
                 type="numpy", label="Output Audio", interactive=False, loop=True
             )
 
+            _ = gr.Markdown("## Parametric EQ")
             peq_plot = gr.Plot(
                 plot_eq(), label="PEQ Frequency Response", elem_id="peq-plot"
             )
+            with gr.Row():
+                with gr.Column():
+                    _ = gr.Markdown("Peak filter 1")
+                    pk1 = fx[0]
+                    pk1_freq = gr.Slider(
+                        minimum=33,
+                        maximum=5400,
+                        value=pk1.params.freq.item(),
+                        interactive=True,
+                    )
+                    pk1_gain = gr.Slider(
+                        minimum=-24,
+                        maximum=24,
+                        value=pk1.params.gain.item(),
+                        interactive=True,
+                    )
+                    pk1_q = gr.Slider(
+                        minimum=0.2,
+                        maximum=20,
+                        value=pk1.params.Q.item(),
+                        interactive=True,
+                    )
+                with gr.Column():
+                    _ = gr.Markdown("Peak filter 2")
+                    pk2 = fx[1]
+                    pk2_freq = gr.Slider(
+                        minimum=200,
+                        maximum=17500,
+                        value=pk2.params.freq.item(),
+                        interactive=True,
+                    )
+                    pk2_gain = gr.Slider(
+                        minimum=-24,
+                        maximum=24,
+                        value=pk2.params.gain.item(),
+                        interactive=True,
+                    )
+                    pk2_q = gr.Slider(
+                        minimum=0.2,
+                        maximum=20,
+                        value=pk2.params.Q.item(),
+                        interactive=True,
+                    )
+            with gr.Row():
+                with gr.Column():
+                    _ = gr.Markdown("Low Shelf")
+                    ls = fx[2]
+                    ls_freq = gr.Slider(
+                        minimum=30,
+                        maximum=200,
+                        value=ls.params.freq.item(),
+                        interactive=True,
+                    )
+                    ls_gain = gr.Slider(
+                        minimum=-24,
+                        maximum=24,
+                        value=ls.params.gain.item(),
+                        interactive=True,
+                    )
+                with gr.Column():
+                    _ = gr.Markdown("High Shelf")
+                    hs = fx[3]
+                    hs_freq = gr.Slider(
+                        minimum=750,
+                        maximum=8300,
+                        value=hs.params.freq.item(),
+                        interactive=True,
+                    )
+                    hs_gain = gr.Slider(
+                        minimum=-24,
+                        maximum=24,
+                        value=hs.params.gain.item(),
+                        interactive=True,
+                    )
+
             comp_plot = gr.Plot(
                 plot_comp(), label="Compressor Curve", elem_id="comp-plot"
             )
@@ -367,6 +484,36 @@ with gr.Blocks() as demo:
             model2json(), label="Effect Settings", max_height=800, open=True
         )
 
+    for eq, s, attr_name in zip(
+        [fx[0]] * 3 + [fx[1]] * 3 + [fx[2]] * 2 + [fx[3]] * 2,
+        [
+            pk1_freq,
+            pk1_gain,
+            pk1_q,
+            pk2_freq,
+            pk2_gain,
+            pk2_q,
+            ls_freq,
+            ls_gain,
+            hs_freq,
+            hs_gain,
+        ],
+        ["freq", "gain", "Q"] * 2 + ["freq", "gain"] * 2,
+    ):
+        s.input(
+            lambda *args, eq=eq, attr_name=attr_name: chain_functions(  # chain_functions(
+                lambda args: (upatePEQ(eq, attr_name, args[0]), args[1]),
+                lambda args: (fx2z(), args[1]),
+                lambda args: [plot_eq()]
+                + z[:NUMBER_OF_PCS].tolist()
+                + [z[args[1] - 1].item(), model2json()],
+            )(
+                args
+            ),
+            inputs=[s, extra_pc_dropdown],
+            outputs=[peq_plot] + sliders + [extra_slider, json_output],
+        )
+
     render_button.click(
         lambda *args: (
             lambda x: (
@@ -419,6 +566,18 @@ with gr.Blocks() as demo:
                 plot_delay(),
                 plot_reverb(),
                 plot_t60(),
+            ]
+            + [
+                pk1.params.freq.item(),
+                pk1.params.gain.item(),
+                pk1.params.Q.item(),
+                pk2.params.freq.item(),
+                pk2.params.gain.item(),
+                pk2.params.Q.item(),
+                ls.params.freq.item(),
+                ls.params.gain.item(),
+                hs.params.freq.item(),
+                hs.params.gain.item(),
             ],
         ),
         inputs=extra_pc_dropdown,
@@ -431,6 +590,18 @@ with gr.Blocks() as demo:
             delay_plot,
             reverb_plot,
             t60_plot,
+        ]
+        + [
+            pk1_freq,
+            pk1_gain,
+            pk1_q,
+            pk2_freq,
+            pk2_gain,
+            pk2_q,
+            ls_freq,
+            ls_gain,
+            hs_freq,
+            hs_gain,
         ],
     )
     reset_button.click(
@@ -446,6 +617,18 @@ with gr.Blocks() as demo:
                 plot_delay(),
                 plot_reverb(),
                 plot_t60(),
+            ]
+            + [
+                pk1.params.freq.item(),
+                pk1.params.gain.item(),
+                pk1.params.Q.item(),
+                pk2.params.freq.item(),
+                pk2.params.gain.item(),
+                pk2.params.Q.item(),
+                ls.params.freq.item(),
+                ls.params.gain.item(),
+                hs.params.freq.item(),
+                hs.params.gain.item(),
             ],
         )(None),
         # inputs=sliders + [extra_slider],
@@ -458,6 +641,18 @@ with gr.Blocks() as demo:
             delay_plot,
             reverb_plot,
             t60_plot,
+        ]
+        + [
+            pk1_freq,
+            pk1_gain,
+            pk1_q,
+            pk2_freq,
+            pk2_gain,
+            pk2_q,
+            ls_freq,
+            ls_gain,
+            hs_freq,
+            hs_gain,
         ],
     )
 
@@ -477,6 +672,16 @@ with gr.Blocks() as demo:
                     plot_delay(),
                     plot_reverb(),
                     plot_t60(),
+                    pk1.params.freq.item(),
+                    pk1.params.gain.item(),
+                    pk1.params.Q.item(),
+                    pk2.params.freq.item(),
+                    pk2.params.gain.item(),
+                    pk2.params.Q.item(),
+                    ls.params.freq.item(),
+                    ls.params.gain.item(),
+                    hs.params.freq.item(),
+                    hs.params.gain.item(),
                 ),
             ),
             inputs=slider,
@@ -487,6 +692,16 @@ with gr.Blocks() as demo:
                 delay_plot,
                 reverb_plot,
                 t60_plot,
+                pk1_freq,
+                pk1_gain,
+                pk1_q,
+                pk2_freq,
+                pk2_gain,
+                pk2_q,
+                ls_freq,
+                ls_gain,
+                hs_freq,
+                hs_gain,
             ],
         )
     extra_slider.input(
@@ -500,10 +715,37 @@ with gr.Blocks() as demo:
                 plot_delay(),
                 plot_reverb(),
                 plot_t60(),
+                pk1.params.freq.item(),
+                pk1.params.gain.item(),
+                pk1.params.Q.item(),
+                pk2.params.freq.item(),
+                pk2.params.gain.item(),
+                pk2.params.Q.item(),
+                ls.params.freq.item(),
+                ls.params.gain.item(),
+                hs.params.freq.item(),
+                hs.params.gain.item(),
             ),
         )(xs),
         inputs=[extra_slider, extra_pc_dropdown],
-        outputs=[json_output, peq_plot, comp_plot, delay_plot, reverb_plot, t60_plot],
+        outputs=[
+            json_output,
+            peq_plot,
+            comp_plot,
+            delay_plot,
+            reverb_plot,
+            t60_plot,
+            pk1_freq,
+            pk1_gain,
+            pk1_q,
+            pk2_freq,
+            pk2_gain,
+            pk2_q,
+            ls_freq,
+            ls_gain,
+            hs_freq,
+            hs_gain,
+        ],
     )
 
     extra_pc_dropdown.input(