feat: refactor global effect handling, enable user session states

app.py

@@ -9,6 +9,7 @@ from hydra.utils import instantiate
 from soxr import resample
 from functools import partial
 from torchcomp import coef2ms, ms2coef
+from copy import deepcopy
 
 from modules.utils import chain_functions, vec2statedict, get_chunks
 from modules.fx import clip_delay_eq_Q
@@ -50,8 +51,8 @@ with open(CONFIG_PATH) as fp:
     fx_config = yaml.safe_load(fp)["model"]
 
 # Global effect
-fx = instantiate(fx_config)
-fx.eval()
+global_fx = instantiate(fx_config)
+global_fx.eval()
 
 pca_params = np.load(PCA_PARAM_FILE)
 mean = pca_params["mean"]
@@ -64,7 +65,7 @@ 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_like(mean)
+# z = torch.zeros_like(mean)
 
 with open(INFO_PATH) as f:
     info = json.load(f)
@@ -91,35 +92,40 @@ vec2dict = partial(
         )
     ),
 )
-fx.load_state_dict(vec2dict(mean), strict=False)
+global_fx.load_state_dict(vec2dict(mean), strict=False)
 
 
 meter = pyln.Meter(44100)
 
 
 @torch.no_grad()
-def z2fx():
+def z2x(z):
     # 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)
-    fx.apply(partial(clip_delay_eq_Q, Q=0.707))
-    return
+    # # print(z)
+    # fx.load_state_dict(vec2dict(x), strict=False)
+    # fx.apply(partial(clip_delay_eq_Q, Q=0.707))
+    return x
 
 
 @torch.no_grad()
-def fx2z():
+def fx2x(fx):
     plt.close("all")
     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
+    return x
+
+
+@torch.no_grad()
+def x2z(x):
+    z = U.T @ (x - mean)
+    return z
 
 
 @torch.no_grad()
-def inference(audio):
+def inference(audio, fx):
     sr, y = audio
     if sr != 44100:
         y = resample(y, sr, 44100)
@@ -163,7 +169,7 @@ def get_important_pcs(n=10, **kwargs):
     return sliders
 
 
-def model2json():
+def model2json(fx):
     fx_names = ["PK1", "PK2", "LS", "HS", "LP", "HP", "DRC"]
     results = {k: v.toJSON() for k, v in zip(fx_names, fx)} | {
         "Panner": fx[7].pan.toJSON()
@@ -190,7 +196,7 @@ def model2json():
 
 
 @torch.no_grad()
-def plot_eq():
+def plot_eq(fx):
     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="-")
@@ -207,14 +213,14 @@ def plot_eq():
 
 
 @torch.no_grad()
-def plot_comp():
+def plot_comp(fx):
     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()
+    cmp_th = fx[6].params.cmp_th.item()
+    exp_th = fx[6].params.exp_th.item()
+    cmp_ratio = fx[6].params.cmp_ratio.item()
+    exp_ratio = fx[6].params.exp_ratio.item()
+    make_up = fx[6].params.make_up.item()
     # print(cmp_ratio, cmp_th, exp_ratio, exp_th, make_up)
 
     comp_in = np.linspace(-80, 0, 100)
@@ -242,16 +248,16 @@ def plot_comp():
 
 
 @torch.no_grad()
-def plot_delay():
+def plot_delay(fx):
     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
+    w, eq_log_mags = get_log_mags_from_eq([fx[7].effects[0].eq])
+    log_gain = fx[7].effects[0].params.gain.log10().item() * 20
+    d = fx[7].effects[0].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
+    log_feedback = fx[7].effects[0].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(
@@ -272,7 +278,7 @@ def plot_delay():
 
 
 @torch.no_grad()
-def plot_reverb():
+def plot_reverb(fx):
     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)
@@ -292,7 +298,7 @@ def plot_reverb():
 
 
 @torch.no_grad()
-def plot_t60():
+def plot_t60(fx):
     fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True)
     fdn = fx[7].effects[1]
     gamma = fdn.params.gamma.squeeze().numpy()
@@ -311,19 +317,39 @@ def plot_t60():
 
 @torch.no_grad()
 def update_param(m, attr_name, value):
-    match type(getattr(m.params, attr_name)):
+    match type(getattr(m, attr_name)):
         case torch.nn.Parameter:
-            getattr(m.params, attr_name).data.copy_(value)
+            getattr(m, attr_name).data.copy_(value)
         case _:
-            setattr(m.params, attr_name, torch.tensor(value))
+            setattr(m, attr_name, torch.tensor(value))
 
 
 @torch.no_grad()
 def update_atrt(comp, attr_name, value):
-    setattr(comp.params, attr_name, ms2coef(torch.tensor(value), 44100))
+    setattr(comp, attr_name, ms2coef(torch.tensor(value), 44100))
+
+
+def vec2fx(x):
+    fx = deepcopy(global_fx)
+    fx.load_state_dict(vec2dict(x), strict=False)
+    fx.apply(partial(clip_delay_eq_Q, Q=0.707))
+    return fx
+
+
+get_last_attribute = lambda m, attr_name: (
+    (m, attr_name)
+    if "." not in attr_name
+    else (lambda x, *remain: get_last_attribute(getattr(m, x), ".".join(remain)))(
+        *attr_name.split(".")
+    )
+)
 
 
 with gr.Blocks() as demo:
+    z = gr.State(torch.zeros_like(mean))
+    fx_params = gr.State(mean)
+    fx = vec2fx(fx_params.value)
+
     gr.Markdown(
         title_md,
         elem_id="title",
@@ -352,12 +378,6 @@ with gr.Blocks() as demo:
                 render_button = gr.Button(
                     "Run", elem_id="render-button", variant="primary"
                 )
-            # random_rest_checkbox = gr.Checkbox(
-            #     label=f"Randomise PCs > {NUMBER_OF_PCS} (default to zeros)",
-            #     value=False,
-            #     elem_id="randomise-checkbox",
-            # )
-            # sliders = get_important_pcs(NUMBER_OF_PCS, value=0)
             with gr.Row():
                 s1 = gr.Slider(
                     minimum=SLIDER_MIN,
@@ -417,7 +437,7 @@ with gr.Blocks() as demo:
             )
 
     _ = gr.Markdown("## Parametric EQ")
-    peq_plot = gr.Plot(plot_eq(), label="PEQ Frequency Response", elem_id="peq-plot")
+    peq_plot = gr.Plot(plot_eq(fx), label="PEQ Frequency Response", elem_id="peq-plot")
     with gr.Row():
         with gr.Column(min_width=160):
             _ = gr.Markdown("High Pass")
@@ -425,14 +445,14 @@ with gr.Blocks() as demo:
             hp_freq = gr.Slider(
                 minimum=16,
                 maximum=5300,
-                value=hp.params.freq.item(),
+                value=fx[5].params.freq.item(),
                 interactive=True,
                 label="Frequency (Hz)",
             )
             hp_q = gr.Slider(
                 minimum=0.5,
                 maximum=10,
-                value=hp.params.Q.item(),
+                value=fx[5].params.Q.item(),
                 interactive=True,
                 label="Q",
             )
@@ -443,14 +463,14 @@ with gr.Blocks() as demo:
             ls_freq = gr.Slider(
                 minimum=30,
                 maximum=200,
-                value=ls.params.freq.item(),
+                value=fx[2].params.freq.item(),
                 interactive=True,
                 label="Frequency (Hz)",
             )
             ls_gain = gr.Slider(
                 minimum=-12,
                 maximum=12,
-                value=ls.params.gain.item(),
+                value=fx[2].params.gain.item(),
                 interactive=True,
                 label="Gain (dB)",
             )
@@ -461,21 +481,21 @@ with gr.Blocks() as demo:
             pk1_freq = gr.Slider(
                 minimum=33,
                 maximum=5400,
-                value=pk1.params.freq.item(),
+                value=fx[0].params.freq.item(),
                 interactive=True,
                 label="Frequency (Hz)",
             )
             pk1_gain = gr.Slider(
                 minimum=-12,
                 maximum=12,
-                value=pk1.params.gain.item(),
+                value=fx[0].params.gain.item(),
                 interactive=True,
                 label="Gain (dB)",
             )
             pk1_q = gr.Slider(
                 minimum=0.2,
                 maximum=20,
-                value=pk1.params.Q.item(),
+                value=fx[0].params.Q.item(),
                 interactive=True,
                 label="Q",
             )
@@ -485,21 +505,21 @@ with gr.Blocks() as demo:
             pk2_freq = gr.Slider(
                 minimum=200,
                 maximum=17500,
-                value=pk2.params.freq.item(),
+                value=fx[1].params.freq.item(),
                 interactive=True,
                 label="Frequency (Hz)",
             )
             pk2_gain = gr.Slider(
                 minimum=-12,
                 maximum=12,
-                value=pk2.params.gain.item(),
+                value=fx[1].params.gain.item(),
                 interactive=True,
                 label="Gain (dB)",
             )
             pk2_q = gr.Slider(
                 minimum=0.2,
                 maximum=20,
-                value=pk2.params.Q.item(),
+                value=fx[1].params.Q.item(),
                 interactive=True,
                 label="Q",
             )
@@ -510,14 +530,14 @@ with gr.Blocks() as demo:
             hs_freq = gr.Slider(
                 minimum=750,
                 maximum=8300,
-                value=hs.params.freq.item(),
+                value=fx[3].params.freq.item(),
                 interactive=True,
                 label="Frequency (Hz)",
             )
             hs_gain = gr.Slider(
                 minimum=-12,
                 maximum=12,
-                value=hs.params.gain.item(),
+                value=fx[3].params.gain.item(),
                 interactive=True,
                 label="Gain (dB)",
             )
@@ -527,14 +547,14 @@ with gr.Blocks() as demo:
             lp_freq = gr.Slider(
                 minimum=200,
                 maximum=18000,
-                value=lp.params.freq.item(),
+                value=fx[4].params.freq.item(),
                 interactive=True,
                 label="Frequency (Hz)",
             )
             lp_q = gr.Slider(
                 minimum=0.5,
                 maximum=10,
-                value=lp.params.Q.item(),
+                value=fx[4].params.Q.item(),
                 interactive=True,
                 label="Q",
             )
@@ -546,55 +566,55 @@ with gr.Blocks() as demo:
             cmp_th = gr.Slider(
                 minimum=-60,
                 maximum=0,
-                value=comp.params.cmp_th.item(),
+                value=fx[6].params.cmp_th.item(),
                 interactive=True,
-                label="Comp. Threshold (dB)",
+                label="fx[6]. Threshold (dB)",
             )
             cmp_ratio = gr.Slider(
                 minimum=1,
                 maximum=20,
-                value=comp.params.cmp_ratio.item(),
+                value=fx[6].params.cmp_ratio.item(),
                 interactive=True,
-                label="Comp. Ratio",
+                label="fx[6]. Ratio",
             )
             make_up = gr.Slider(
                 minimum=-12,
                 maximum=12,
-                value=comp.params.make_up.item(),
+                value=fx[6].params.make_up.item(),
                 interactive=True,
                 label="Make Up (dB)",
             )
             attack_time = gr.Slider(
                 minimum=0.1,
                 maximum=100,
-                value=coef2ms(comp.params.at, 44100).item(),
+                value=coef2ms(fx[6].params.at, 44100).item(),
                 interactive=True,
                 label="Attack Time (ms)",
             )
             release_time = gr.Slider(
                 minimum=50,
                 maximum=1000,
-                value=coef2ms(comp.params.rt, 44100).item(),
+                value=coef2ms(fx[6].params.rt, 44100).item(),
                 interactive=True,
                 label="Release Time (ms)",
             )
             exp_ratio = gr.Slider(
                 minimum=0,
                 maximum=1,
-                value=comp.params.exp_ratio.item(),
+                value=fx[6].params.exp_ratio.item(),
                 interactive=True,
                 label="Exp. Ratio",
             )
             exp_th = gr.Slider(
                 minimum=-80,
                 maximum=0,
-                value=comp.params.exp_th.item(),
+                value=fx[6].params.exp_th.item(),
                 interactive=True,
                 label="Exp. Threshold (dB)",
             )
         with gr.Column():
             comp_plot = gr.Plot(
-                plot_comp(), label="Compressor Curve", elem_id="comp-plot"
+                plot_comp(fx), label="Compressor Curve", elem_id="comp-plot"
             )
 
     _ = gr.Markdown("## Ping-Pong Delay")
@@ -604,160 +624,215 @@ with gr.Blocks() as demo:
             delay_time = gr.Slider(
                 minimum=100,
                 maximum=1000,
-                value=delay.params.delay.item(),
+                value=fx[7].effects[0].params.delay.item(),
                 interactive=True,
                 label="Delay Time (ms)",
             )
             feedback = gr.Slider(
                 minimum=0,
                 maximum=1,
-                value=delay.params.feedback.item(),
+                value=fx[7].effects[0].params.feedback.item(),
                 interactive=True,
                 label="Feedback",
             )
             delay_gain = gr.Slider(
                 minimum=-80,
                 maximum=0,
-                value=delay.params.gain.log10().item() * 20,
+                value=fx[7].effects[0].params.gain.log10().item() * 20,
                 interactive=True,
                 label="Gain (dB)",
             )
             odd_pan = gr.Slider(
                 minimum=-100,
                 maximum=100,
-                value=delay.odd_pan.params.pan.item() * 200 - 100,
+                value=fx[7].effects[0].odd_pan.params.pan.item() * 200 - 100,
                 interactive=True,
                 label="Odd Delay Pan",
             )
             even_pan = gr.Slider(
                 minimum=-100,
                 maximum=100,
-                value=delay.even_pan.params.pan.item() * 200 - 100,
+                value=fx[7].effects[0].even_pan.params.pan.item() * 200 - 100,
                 interactive=True,
                 label="Even Delay Pan",
             )
             delay_lp_freq = gr.Slider(
                 minimum=200,
                 maximum=16000,
-                value=delay.eq.params.freq.item(),
+                value=fx[7].effects[0].eq.params.freq.item(),
                 interactive=True,
                 label="Low Pass Frequency (Hz)",
             )
         with gr.Column():
             delay_plot = gr.Plot(
-                plot_delay(), label="Delay Frequency Response", elem_id="delay-plot"
+                plot_delay(fx), label="Delay Frequency Response", elem_id="delay-plot"
             )
 
     with gr.Row():
         reverb_plot = gr.Plot(
-            plot_reverb(),
+            plot_reverb(fx),
             label="Reverb Tone Correction PEQ",
             elem_id="reverb-plot",
             min_width=160,
         )
         t60_plot = gr.Plot(
-            plot_t60(), label="Reverb T60", elem_id="t60-plot", min_width=160
+            plot_t60(fx), label="Reverb T60", elem_id="t60-plot", min_width=160
         )
 
     with gr.Row():
         json_output = gr.JSON(
-            model2json(), label="Effect Settings", max_height=800, open=True
+            model2json(fx), label="Effect Settings", max_height=800, open=True
         )
 
-    update_pc = lambda i: z[:NUMBER_OF_PCS].tolist() + [z[i - 1].item()]
+    update_pc = lambda z, i: z[:NUMBER_OF_PCS].tolist() + [z[i - 1].item()]
     update_pc_outputs = sliders + [extra_slider]
 
-    for eq, s, attr_name in zip(
-        [fx[0]] * 3
-        + [fx[1]] * 3
-        + [fx[2]] * 2
-        + [fx[3]] * 2
-        + [fx[4]] * 2
-        + [fx[5]] * 2,
-        [
-            pk1_freq,
-            pk1_gain,
-            pk1_q,
-            pk2_freq,
-            pk2_gain,
-            pk2_q,
-            ls_freq,
-            ls_gain,
-            hs_freq,
-            hs_gain,
-            lp_freq,
-            lp_q,
-            hp_freq,
-            hp_q,
-        ],
-        ["freq", "gain", "Q"] * 2 + ["freq", "gain"] * 2 + ["freq", "Q"] * 2,
-    ):
+    peq_sliders = [
+        pk1_freq,
+        pk1_gain,
+        pk1_q,
+        pk2_freq,
+        pk2_gain,
+        pk2_q,
+        ls_freq,
+        ls_gain,
+        hs_freq,
+        hs_gain,
+        lp_freq,
+        lp_q,
+        hp_freq,
+        hp_q,
+    ]
+    peq_attr_names = (
+        ["freq", "gain", "Q"] * 2 + ["freq", "gain"] * 2 + ["freq", "Q"] * 2
+    )
+    peq_indices = [0] * 3 + [1] * 3 + [2] * 2 + [3] * 2 + [4] * 2 + [5] * 2
+
+    cmp_sliders = [
+        cmp_th,
+        cmp_ratio,
+        make_up,
+        exp_ratio,
+        exp_th,
+        attack_time,
+        release_time,
+    ]
+    cmp_update_funcs = [update_param] * 5 + [update_atrt] * 2
+    cmp_attr_names = [
+        "cmp_th",
+        "cmp_ratio",
+        "make_up",
+        "exp_ratio",
+        "exp_th",
+        "at",
+        "rt",
+    ]
+
+    delay_sliders = [delay_time, feedback, delay_lp_freq, delay_gain, odd_pan, even_pan]
+    delay_update_funcs = [update_param] * 3 + [
+        lambda m, a, v: update_param(m, a, 10 ** (v / 20)),
+        lambda m, a, v: update_param(m, a, (v + 100) / 200),
+        lambda m, a, v: update_param(m, a, (v + 100) / 200),
+    ]
+    delay_attr_names = [
+        "params.delay",
+        "params.feedback",
+        "eq.params.freq",
+        "params.gain",
+        "odd_pan.params.pan",
+        "even_pan.params.pan",
+    ]
+    delay_update_plot_flag = [True] * 4 + [False] * 2
+
+    all_effect_sliders = peq_sliders + cmp_sliders + delay_sliders
+    split_sizes = [len(peq_sliders), len(cmp_sliders), len(delay_sliders)]
+
+    def assign_fx_params(fx, *args):
+        peq_sliders, cmp_sliders, delay_sliders = (
+            args[: split_sizes[0]],
+            args[split_sizes[0] : sum(split_sizes[:2])],
+            args[sum(split_sizes[:2]) :],
+        )
+        for idx, s, attr_name in zip(peq_indices, peq_sliders, peq_attr_names):
+            update_param(fx[idx].params, attr_name, s)
+
+        for f, s, attr_name in zip(cmp_update_funcs, cmp_sliders, cmp_attr_names):
+            f(fx[6].params, attr_name, s)
+
+        for f, s, attr_name in zip(delay_update_funcs, delay_sliders, delay_attr_names):
+            m, name = get_last_attribute(fx[7].effects[0], attr_name)
+            f(m, name, s)
+
+        return fx
+
+    for idx, s, attr_name in zip(peq_indices, peq_sliders, peq_attr_names):
         s.input(
-            lambda *args, eq=eq, attr_name=attr_name: chain_functions(  # chain_functions(
-                lambda args: (update_param(eq, attr_name, args[0]), args[1]),
-                lambda args: (fx2z(), args[1]),
-                lambda args: args[1],
-                lambda i: update_pc(i) + [model2json(), plot_eq()],
+            lambda *args, idx=idx, attr_name=attr_name: chain_functions(  # chain_functions(
+                lambda args: (assign_fx_params(vec2fx(args[0]), *args[3:]), *args[1:3]),
+                lambda args: (
+                    update_param(args[0][idx].params, attr_name, args[1]),
+                    args[0],
+                    args[2],
+                ),
+                lambda args: (fx2x(args[1]), *args[1:]),
+                lambda args: [x2z(args[0]), *args],
+                lambda args: args[:2]
+                + [model2json(args[2]), plot_eq(args[2])]
+                + update_pc(args[0], args[3]),
             )(
                 args
             ),
-            inputs=[s, extra_pc_dropdown],
-            outputs=update_pc_outputs + [json_output, peq_plot],
+            inputs=[fx_params, s, extra_pc_dropdown] + all_effect_sliders,
+            outputs=[z, fx_params, json_output, peq_plot] + update_pc_outputs,
         )
 
-    for f, s, attr_name in zip(
-        [update_param] * 5 + [update_atrt] * 2,
-        [
-            cmp_th,
-            cmp_ratio,
-            make_up,
-            exp_ratio,
-            exp_th,
-            attack_time,
-            release_time,
-        ],
-        ["cmp_th", "cmp_ratio", "make_up", "exp_ratio", "exp_th", "at", "rt"],
-    ):
+    for f, s, attr_name in zip(cmp_update_funcs, cmp_sliders, cmp_attr_names):
         s.input(
             lambda *args, attr_name=attr_name, f=f: chain_functions(
-                lambda args: (f(comp, attr_name, args[0]), args[1]),
-                lambda args: (fx2z(), args[1]),
-                lambda args: args[1],
-                lambda i: update_pc(i) + [model2json(), plot_comp()],
+                lambda args: (assign_fx_params(vec2fx(args[0]), *args[3:]), *args[1:3]),
+                lambda args: (
+                    f(args[0][6].params, attr_name, args[1]),
+                    args[0],
+                    args[2],
+                ),
+                lambda args: (fx2x(args[1]), *args[1:]),
+                lambda args: [x2z(args[0]), *args],
+                lambda args: args[:2]
+                + [model2json(args[2]), plot_comp(args[2])]
+                + update_pc(args[0], args[3]),
             )(args),
-            inputs=[s, extra_pc_dropdown],
-            outputs=update_pc_outputs + [json_output, comp_plot],
+            inputs=[fx_params, s, extra_pc_dropdown] + all_effect_sliders,
+            outputs=[z, fx_params, json_output, comp_plot] + update_pc_outputs,
         )
 
-    for m, f, s, attr_name, update_plot in zip(
-        [delay] * 2 + [delay.eq] + [delay, delay.odd_pan, delay.even_pan],
-        [update_param] * 3
-        + [
-            lambda m, a, v: update_param(m, a, 10 ** (v / 20)),
-            lambda m, a, v: update_param(m, a, (v + 100) / 200),
-            lambda m, a, v: update_param(m, a, (v + 100) / 200),
-        ],
-        [delay_time, feedback, delay_lp_freq, delay_gain, odd_pan, even_pan],
-        ["delay", "feedback", "freq", "gain", "pan", "pan"],
-        [True] * 4 + [False] * 2,
+    for f, s, attr_name, update_plot in zip(
+        delay_update_funcs, delay_sliders, delay_attr_names, delay_update_plot_flag
     ):
         s.input(
-            lambda *args, f=f, m=m, attr_name=attr_name, update_plot=update_plot: chain_functions(
-                lambda args: (f(m, attr_name, args[0]), args[1]),
-                lambda args: (fx2z(), args[1]),
-                lambda args: args[1],
-                lambda i: (
-                    update_pc(i)
-                    + [model2json()]
-                    + ([plot_delay()] if update_plot else [])
+            lambda *args, f=f, attr_name=attr_name, update_plot=update_plot: chain_functions(
+                lambda args: (assign_fx_params(vec2fx(args[0]), *args[3:]), *args[1:3]),
+                lambda args: (
+                    # f(args[0][7].effects[0], attr_name, args[1]),
+                    f(*get_last_attribute(args[0][7].effects[0], attr_name), args[1]),
+                    args[0],
+                    args[2],
+                ),
+                lambda args: (fx2x(args[1]), *args[1:]),
+                lambda args: [x2z(args[0]), *args],
+                lambda args: (
+                    args[:2]
+                    + [model2json(args[2])]
+                    + ([plot_delay(args[2])] if update_plot else [])
+                    + update_pc(args[0], args[3])
                 ),
             )(
                 args
             ),
-            inputs=[s, extra_pc_dropdown],
-            outputs=update_pc_outputs
+            inputs=[fx_params, s, extra_pc_dropdown] + all_effect_sliders,
+            outputs=[z, fx_params]
             + [json_output]
-            + ([delay_plot] if update_plot else []),
+            + ([delay_plot] if update_plot else [])
+            + update_pc_outputs,
         )
 
     render_button.click(
@@ -767,10 +842,17 @@ with gr.Blocks() as demo:
         #         model2json(),
         #     )
         # )(inference(*args)),
-        inference,
+        # inference,
+        # lambda audio, x: inference(audio, vec2fx(x)),
+        lambda audio, *args: chain_functions(
+            lambda args: assign_fx_params(vec2fx(args[0]), *args[1:]),
+            partial(inference, audio),
+        )(args),
         inputs=[
             audio_input,
-        ],
+            fx_params,
+        ]
+        + all_effect_sliders,
         outputs=[
             audio_output,
             direct_output,
@@ -778,34 +860,34 @@ with gr.Blocks() as demo:
         ],
     )
 
-    update_fx = lambda: [
-        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(),
-        lp.params.freq.item(),
-        lp.params.Q.item(),
-        hp.params.freq.item(),
-        hp.params.Q.item(),
-        comp.params.cmp_th.item(),
-        comp.params.cmp_ratio.item(),
-        comp.params.make_up.item(),
-        comp.params.exp_th.item(),
-        comp.params.exp_ratio.item(),
-        coef2ms(comp.params.at, 44100).item(),
-        coef2ms(comp.params.rt, 44100).item(),
-        delay.params.delay.item(),
-        delay.params.feedback.item(),
-        delay.params.gain.log10().item() * 20,
-        delay.eq.params.freq.item(),
-        delay.odd_pan.params.pan.item() * 200 - 100,
-        delay.even_pan.params.pan.item() * 200 - 100,
+    update_fx = lambda fx: [
+        fx[0].params.freq.item(),
+        fx[0].params.gain.item(),
+        fx[0].params.Q.item(),
+        fx[1].params.freq.item(),
+        fx[1].params.gain.item(),
+        fx[1].params.Q.item(),
+        fx[2].params.freq.item(),
+        fx[2].params.gain.item(),
+        fx[3].params.freq.item(),
+        fx[3].params.gain.item(),
+        fx[4].params.freq.item(),
+        fx[4].params.Q.item(),
+        fx[5].params.freq.item(),
+        fx[5].params.Q.item(),
+        fx[6].params.cmp_th.item(),
+        fx[6].params.cmp_ratio.item(),
+        fx[6].params.make_up.item(),
+        fx[6].params.exp_th.item(),
+        fx[6].params.exp_ratio.item(),
+        coef2ms(fx[6].params.at, 44100).item(),
+        coef2ms(fx[6].params.rt, 44100).item(),
+        fx[7].effects[0].params.delay.item(),
+        fx[7].effects[0].params.feedback.item(),
+        fx[7].effects[0].params.gain.log10().item() * 20,
+        fx[7].effects[0].eq.params.freq.item(),
+        fx[7].effects[0].odd_pan.params.pan.item() * 200 - 100,
+        fx[7].effects[0].even_pan.params.pan.item() * 200 - 100,
     ]
     update_fx_outputs = [
         pk1_freq,
@@ -836,12 +918,12 @@ with gr.Blocks() as demo:
         odd_pan,
         even_pan,
     ]
-    update_plots = lambda: [
-        plot_eq(),
-        plot_comp(),
-        plot_delay(),
-        plot_reverb(),
-        plot_t60(),
+    update_plots = lambda fx: [
+        plot_eq(fx),
+        plot_comp(fx),
+        plot_delay(fx),
+        plot_reverb(fx),
+        plot_t60(fx),
     ]
     update_plots_outputs = [
         peq_plot,
@@ -851,56 +933,61 @@ with gr.Blocks() as demo:
         t60_plot,
     ]
 
-    update_all = lambda i: update_pc(i) + update_fx() + update_plots()
+    update_all = lambda z, fx, i: update_pc(z, i) + update_fx(fx) + update_plots(fx)
     update_all_outputs = update_pc_outputs + update_fx_outputs + update_plots_outputs
 
     random_button.click(
         chain_functions(
-            lambda i: (z.normal_(0, 1).clip_(SLIDER_MIN, SLIDER_MAX), i),
-            lambda args: (z2fx(), args[1]),
-            lambda args: args[1],
-            update_all,
+            lambda i: (torch.randn_like(mean).clip(SLIDER_MIN, SLIDER_MAX), i),
+            lambda args: (args[0], vec2fx(z2x(args[0])), args[1]),
+            lambda args: update_all(*args) + [args[0]],
         ),
         inputs=extra_pc_dropdown,
-        outputs=update_all_outputs,
+        outputs=update_all_outputs + [z],
     )
     reset_button.click(
         # lambda: (lambda _: [0 for _ in range(NUMBER_OF_PCS + 1)])(z.zero_()),
         lambda: chain_functions(
-            lambda _: z.zero_(),
-            lambda _: z2fx(),
-            lambda _: update_all(NUMBER_OF_PCS),
+            lambda _: torch.zeros_like(mean),
+            lambda z: (z, vec2fx(z2x(z))),
+            lambda args: update_all(args[0], args[1], NUMBER_OF_PCS) + [args[0]],
         )(None),
-        outputs=update_all_outputs,
+        outputs=update_all_outputs + [z],
     )
 
-    def update_z(s, i):
+    def update_z(z, s, i):
         z[i] = s
-        return
+        return z
 
     for i, slider in enumerate(sliders):
         slider.input(
-            chain_functions(
-                partial(update_z, i=i),
-                lambda _: z2fx(),
-                lambda _: update_fx() + update_plots() + [model2json()],
-            ),
-            inputs=slider,
-            outputs=update_fx_outputs + update_plots_outputs + [json_output],
+            lambda *args, i=i: chain_functions(
+                lambda args: update_z(args[0], args[1], i),
+                lambda z: (z, vec2fx(z2x(z))),
+                lambda args: [args[0]]
+                + update_fx(args[1])
+                + update_plots(args[1])
+                + [model2json(args[1])],
+            )(args),
+            inputs=[z, slider],
+            outputs=[z] + update_fx_outputs + update_plots_outputs + [json_output],
         )
     extra_slider.input(
         lambda *xs: chain_functions(
-            lambda args: update_z(args[0], args[1] - 1),
-            lambda _: z2fx(),
-            lambda _: update_fx() + update_plots() + [model2json()],
+            lambda args: update_z(args[0], args[1], args[2]),
+            lambda z: (z, vec2fx(z2x(z))),
+            lambda args: [args[0]]
+            + update_fx(args[1])
+            + update_plots(args[1])
+            + [model2json(args[1])],
         )(xs),
-        inputs=[extra_slider, extra_pc_dropdown],
-        outputs=update_fx_outputs + update_plots_outputs + [json_output],
+        inputs=[z, extra_slider, extra_pc_dropdown],
+        outputs=[z] + update_fx_outputs + update_plots_outputs + [json_output],
     )
 
     extra_pc_dropdown.input(
-        lambda i: z[i - 1].item(),
-        inputs=extra_pc_dropdown,
+        lambda z, i: z[i - 1].item(),
+        inputs=[z, extra_pc_dropdown],
         outputs=extra_slider,
     )