feat: add plotting functions for compressor, delay, reverb, and T60 responses

app.py

@@ -103,6 +103,7 @@ def z2fx():
     x = U @ z + mean
     # print(z)
     fx.load_state_dict(vec2dict(x), strict=False)
+    fx.apply(partial(clip_delay_eq_Q, Q=0.707))
     return
 
 
@@ -136,8 +137,6 @@ def inference(audio):
     if y.shape[1] != 1:
         y = y.mean(dim=1, keepdim=True)
 
-    fx.apply(partial(clip_delay_eq_Q, Q=0.707))
-
     rendered = fx(y).squeeze(0).T.numpy()
     if np.max(np.abs(rendered)) > 1:
         rendered = rendered / np.max(np.abs(rendered))
@@ -177,7 +176,7 @@ def model2json():
 
 @torch.no_grad()
 def plot_eq():
-    fig, ax = plt.subplots(figsize=(8, 4))
+    fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True)
     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):
@@ -192,6 +191,109 @@ def plot_eq():
     return fig
 
 
+@torch.no_grad()
+def plot_comp():
+    fig, ax = plt.subplots(figsize=(6, 5), constrained_layout=True)
+    comp = fx[6]
+    cmp_th = comp.params.cmp_th.item()
+    exp_th = comp.params.exp_th.item()
+    cmp_ratio = comp.params.cmp_ratio.item()
+    exp_ratio = comp.params.exp_ratio.item()
+    make_up = comp.params.make_up.item()
+    # print(cmp_ratio, cmp_th, exp_ratio, exp_th, make_up)
+
+    comp_in = np.linspace(-80, 0, 100)
+    comp_curve = np.where(
+        comp_in > cmp_th,
+        comp_in - (comp_in - cmp_th) * (cmp_ratio - 1) / cmp_ratio,
+        comp_in,
+    )
+    comp_out = (
+        np.where(
+            comp_curve < exp_th,
+            comp_curve - (exp_th - comp_curve) / exp_ratio,
+            comp_curve,
+        )
+        + make_up
+    )
+    ax.plot(comp_in, comp_out, c="black", linestyle="-")
+    ax.plot(comp_in, comp_in, c="r", alpha=0.5)
+    ax.set_xlabel("Input Level (dB)")
+    ax.set_ylabel("Output Level (dB)")
+    ax.set_xlim(-80, 0)
+    ax.set_ylim(-80, 0)
+    ax.grid()
+    return fig
+
+
+@torch.no_grad()
+def plot_delay():
+    fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True)
+    delay = fx[7].effects[0]
+    w, eq_log_mags = get_log_mags_from_eq([delay.eq])
+    log_gain = delay.params.gain.log10().item() * 20
+    d = delay.params.delay.item() / 1000
+    log_mag = sum(eq_log_mags)
+    ax.plot(w, log_mag + log_gain, color="black", linestyle="-")
+
+    log_feedback = delay.params.feedback.log10().item() * 20
+    for i in range(1, 10):
+        feedback_log_mag = log_mag * (i + 1) + log_feedback * i + log_gain
+        ax.plot(
+            w,
+            feedback_log_mag,
+            c="black",
+            alpha=max(0, (10 - i * d * 4) / 10),
+            linestyle="-",
+        )
+
+    ax.set_xscale("log")
+    ax.set_xlim(20, 20000)
+    ax.set_ylim(-80, 0)
+    ax.set_xlabel("Frequency (Hz)")
+    ax.set_ylabel("Magnitude (dB)")
+    ax.grid()
+    return fig
+
+
+@torch.no_grad()
+def plot_reverb():
+    fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True)
+    fdn = fx[7].effects[1]
+    w, eq_log_mags = get_log_mags_from_eq(fdn.eq)
+
+    bc = fdn.params.c.norm() * fdn.params.b.norm()
+    log_bc = torch.log10(bc).item() * 20
+    eq_log_mags = [x + log_bc / len(eq_log_mags) for x in eq_log_mags]
+    ax.plot(w, sum(eq_log_mags), color="black", linestyle="-")
+
+    ax.set_xlabel("Frequency (Hz)")
+    ax.set_ylabel("Magnitude (dB)")
+    ax.set_xlim(20, 20000)
+    ax.set_ylim(-40, 6)
+    ax.set_xscale("log")
+    ax.grid()
+    return fig
+
+
+@torch.no_grad()
+def plot_t60():
+    fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True)
+    fdn = fx[7].effects[1]
+    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
+    ax.plot(w, t60, color="black", linestyle="-")
+    ax.set_xlabel("Frequency (Hz)")
+    ax.set_ylabel("T60 (s)")
+    ax.set_xlim(20, 20000)
+    ax.set_ylim(0, 9)
+    ax.set_xscale("log")
+    ax.grid()
+    return fig
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         title_md,
@@ -249,19 +351,49 @@ with gr.Blocks() as demo:
             peq_plot = gr.Plot(
                 plot_eq(), label="PEQ Frequency Response", elem_id="peq-plot"
             )
+            comp_plot = gr.Plot(
+                plot_comp(), label="Compressor Curve", elem_id="comp-plot"
+            )
+            delay_plot = gr.Plot(
+                plot_delay(), label="Delay Frequency Response", elem_id="delay-plot"
+            )
+            reverb_plot = gr.Plot(
+                plot_reverb(), label="Reverb Tone Correction PEQ", elem_id="reverb-plot"
+            )
+            t60_plot = gr.Plot(plot_t60(), label="Reverb T60", elem_id="t60-plot")
 
     with gr.Row():
-        json_output = gr.JSON(label="Effect Settings", max_height=800, open=True)
+        json_output = gr.JSON(
+            model2json(), label="Effect Settings", max_height=800, open=True
+        )
 
     render_button.click(
-        lambda *args: (lambda x: (x, model2json(), plot_eq()))(inference(*args)),
+        lambda *args: (
+            lambda x: (
+                x,
+                model2json(),
+                plot_eq(),
+                plot_comp(),
+                plot_delay(),
+                plot_reverb(),
+                plot_t60(),
+            )
+        )(inference(*args)),
         inputs=[
             audio_input,
             # random_rest_checkbox,
         ]
         # + sliders,
         ,
-        outputs=[audio_output, json_output, peq_plot],
+        outputs=[
+            audio_output,
+            json_output,
+            peq_plot,
+            comp_plot,
+            delay_plot,
+            reverb_plot,
+            t60_plot,
+        ],
     )
 
     random_button.click(
@@ -279,20 +411,54 @@ with gr.Blocks() as demo:
             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()],
+            + [
+                args[0][args[1] - 1].item(),
+                model2json(),
+                plot_eq(),
+                plot_comp(),
+                plot_delay(),
+                plot_reverb(),
+                plot_t60(),
+            ],
         ),
         inputs=extra_pc_dropdown,
-        outputs=sliders + [extra_slider, peq_plot],
+        outputs=sliders
+        + [
+            extra_slider,
+            json_output,
+            peq_plot,
+            comp_plot,
+            delay_plot,
+            reverb_plot,
+            t60_plot,
+        ],
     )
     reset_button.click(
         # 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()],
+            lambda _: [0 for _ in range(NUMBER_OF_PCS + 1)]
+            + [
+                model2json(),
+                plot_eq(),
+                plot_comp(),
+                plot_delay(),
+                plot_reverb(),
+                plot_t60(),
+            ],
         )(None),
         # inputs=sliders + [extra_slider],
-        outputs=sliders + [extra_slider, peq_plot],
+        outputs=sliders
+        + [
+            extra_slider,
+            json_output,
+            peq_plot,
+            comp_plot,
+            delay_plot,
+            reverb_plot,
+            t60_plot,
+        ],
     )
 
     def update_z(s, i):
@@ -304,19 +470,40 @@ with gr.Blocks() as demo:
             chain_functions(
                 partial(update_z, i=i),
                 lambda _: z2fx(),
-                lambda _: plot_eq(),
+                lambda _: (
+                    model2json(),
+                    plot_eq(),
+                    plot_comp(),
+                    plot_delay(),
+                    plot_reverb(),
+                    plot_t60(),
+                ),
             ),
             inputs=slider,
-            outputs=peq_plot,
+            outputs=[
+                json_output,
+                peq_plot,
+                comp_plot,
+                delay_plot,
+                reverb_plot,
+                t60_plot,
+            ],
         )
     extra_slider.input(
         lambda *xs: chain_functions(
             lambda args: update_z(args[0], args[1] - 1),
             lambda _: z2fx(),
-            lambda _: plot_eq(),
+            lambda _: (
+                model2json(),
+                plot_eq(),
+                plot_comp(),
+                plot_delay(),
+                plot_reverb(),
+                plot_t60(),
+            ),
         )(xs),
         inputs=[extra_slider, extra_pc_dropdown],
-        outputs=peq_plot,
+        outputs=[json_output, peq_plot, comp_plot, delay_plot, reverb_plot, t60_plot],
     )
 
     extra_pc_dropdown.input(