feat: enhance chain_functions for improved function composition and readability

app.py

@@ -7,15 +7,23 @@ import json
 import pyloudnorm as pyln
 from hydra.utils import instantiate
 from soxr import resample
-from functools import partial
+from functools import partial, reduce
 from torchcomp import coef2ms, ms2coef
 from copy import deepcopy
 
-from modules.utils import chain_functions, vec2statedict, get_chunks
+from modules.utils import vec2statedict, get_chunks
 from modules.fx import clip_delay_eq_Q
 from plot_utils import get_log_mags_from_eq
 
 
+def chain_functions(*functions):
+    return lambda *initial_args: reduce(
+        lambda xs, f: f(*xs) if isinstance(xs, tuple) else f(xs),
+        functions,
+        initial_args,
+    )
+
+
 title_md = "# Vocal Effects Generator"
 description_md = """
 This is a demo of the paper [DiffVox: A Differentiable Model for Capturing and Analysing Professional Effects Distributions](https://arxiv.org/abs/2504.14735), accepted at DAFx 2025.
@@ -211,11 +219,11 @@ def plot_eq(fx):
 def plot_comp(fx):
     fig, ax = plt.subplots(figsize=(6, 5), constrained_layout=True)
     comp = fx[6]
-    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()
+    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)
@@ -246,13 +254,13 @@ def plot_comp(fx):
 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([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
+    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 = fx[7].effects[0].params.feedback.log10().item() * 20
+    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(
@@ -631,11 +639,11 @@ with gr.Blocks() as demo:
     ]
 
     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_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)] * 2
+    )
     delay_attr_names = [
         "params.delay",
         "params.feedback",
@@ -669,20 +677,22 @@ with gr.Blocks() as demo:
 
     for idx, s, attr_name in zip(peq_indices, peq_sliders, peq_attr_names):
         s.input(
-            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],
+            chain_functions(
+                lambda x, s, extra_pc_idx, *all_s: (
+                    assign_fx_params(vec2fx(x), *all_s),
+                    s,
+                    extra_pc_idx,
                 ),
-                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
+                lambda fx, s, extra_pc_idx, idx=idx, attr_name=attr_name: (
+                    update_param(fx[idx].params, attr_name, s),
+                    fx,
+                    extra_pc_idx,
+                ),
+                lambda _, fx, extra_pc_idx: (fx2x(fx), fx, extra_pc_idx),
+                lambda x, fx, extra_pc_idx: (x2z(x), x, fx, extra_pc_idx),
+                lambda z, x, fx, extra_pc_idx: [z, x]
+                + [model2json(fx), plot_eq(fx)]
+                + update_pc(z, extra_pc_idx),
             ),
             inputs=[fx_params, s, extra_pc_dropdown] + all_effect_sliders,
             outputs=[z, fx_params, json_output, peq_plot] + update_pc_outputs,
@@ -690,19 +700,23 @@ with gr.Blocks() as demo:
 
     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: (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],
+            chain_functions(
+                lambda x, s, e_pc_i, *all_s: (
+                    assign_fx_params(vec2fx(x), *all_s),
+                    s,
+                    e_pc_i,
+                ),
+                lambda fx, s, e_pc_i, attr_name=attr_name, f=f: (
+                    f(fx[6].params, attr_name, s),
+                    fx,
+                    e_pc_i,
                 ),
-                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),
+                lambda _, fx, e_pc_i: (fx2x(fx), fx, e_pc_i),
+                lambda x, fx, e_pc_i: (x2z(x), x, fx, e_pc_i),
+                lambda z, x, fx, e_pc_i: [z, x]
+                + [model2json(fx), plot_comp(fx)]
+                + update_pc(z, e_pc_i),
+            ),
             inputs=[fx_params, s, extra_pc_dropdown] + all_effect_sliders,
             outputs=[z, fx_params, json_output, comp_plot] + update_pc_outputs,
         )
@@ -711,24 +725,25 @@ with gr.Blocks() as demo:
         delay_update_funcs, delay_sliders, delay_attr_names, delay_update_plot_flag
     ):
         s.input(
-            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],
+            chain_functions(
+                lambda x, s, e_pc_i, *all_s: (
+                    assign_fx_params(vec2fx(x), *all_s),
+                    s,
+                    e_pc_i,
+                ),
+                lambda fx, s, e_pc_i, f=f, attr_name=attr_name: (
+                    f(*get_last_attribute(fx[7].effects[0], attr_name), s),
+                    fx,
+                    e_pc_i,
                 ),
-                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])
+                lambda _, fx, e_pc_i: (fx2x(fx), fx, e_pc_i),
+                lambda x, fx, e_pc_i: (x2z(x), x, fx, e_pc_i),
+                lambda z, x, fx, e_pc_i, update_plot=update_plot: (
+                    [z, x]
+                    + [model2json(fx)]
+                    + ([plot_delay(fx)] if update_plot else [])
+                    + update_pc(z, e_pc_i)
                 ),
-            )(
-                args
             ),
             inputs=[fx_params, s, extra_pc_dropdown] + all_effect_sliders,
             outputs=[z, fx_params]
@@ -738,18 +753,10 @@ with gr.Blocks() as demo:
         )
 
     render_button.click(
-        # lambda *args: (
-        #     lambda x: (
-        #         x,
-        #         model2json(),
-        #     )
-        # )(inference(*args)),
-        # 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),
+        chain_functions(
+            lambda audio, x, *all_s: (audio, assign_fx_params(vec2fx(x), *all_s)),
+            inference,
+        ),
         inputs=[
             audio_input,
             fx_params,
@@ -841,22 +848,21 @@ with gr.Blocks() as demo:
     random_button.click(
         chain_functions(
             lambda i: (torch.randn_like(mean).clip(SLIDER_MIN, SLIDER_MAX), i),
-            lambda args: (args[0], z2x(args[0]), args[1]),
-            lambda args: [args[0], args[1], vec2fx(args[1]), args[2]],
-            lambda args: update_all(args[0], args[2], args[3]) + args[:2],
+            lambda z, i: (z, z2x(z), i),
+            lambda z, x, i: (z, x, vec2fx(x), i),
+            lambda z, x, fx, i: [z, x] + update_all(z, fx, i),
         ),
         inputs=extra_pc_dropdown,
-        outputs=update_all_outputs + [z, fx_params],
+        outputs=[z, fx_params] + update_all_outputs,
     )
     reset_button.click(
-        # lambda: (lambda _: [0 for _ in range(NUMBER_OF_PCS + 1)])(z.zero_()),
-        lambda: chain_functions(
-            lambda _: torch.zeros_like(mean),
+        chain_functions(
+            lambda: torch.zeros_like(mean),
             lambda z: (z, z2x(z)),
-            lambda args: [*args[:2], vec2fx(args[1])],
-            lambda args: update_all(args[0], args[2], NUMBER_OF_PCS) + args[:2],
-        )(None),
-        outputs=update_all_outputs + [z, fx_params],
+            lambda z, x: (z, x, vec2fx(x)),
+            lambda z, x, fx: [z, x] + update_all(z, fx, NUMBER_OF_PCS),
+        ),
+        outputs=[z, fx_params] + update_all_outputs,
     )
 
     def update_z(z, s, i):
@@ -865,36 +871,28 @@ with gr.Blocks() as demo:
 
     for i, slider in enumerate(sliders):
         slider.input(
-            lambda *args, i=i: chain_functions(
-                lambda args: update_z(args[0], args[1], i),
+            chain_functions(
+                lambda z, s, i=i: update_z(z, s, i),
                 lambda z: (z, z2x(z)),
-                lambda args: [args[0], args[1], vec2fx(args[1])],
-                lambda args: args[:2]
-                + update_fx(args[2])
-                + update_plots(args[2])
-                + [model2json(args[2])],
-            )(args),
+                lambda z, x: (z, x, vec2fx(x)),
+                lambda z, x, fx: [z, x, model2json(fx)]
+                + update_fx(fx)
+                + update_plots(fx),
+            ),
             inputs=[z, slider],
-            outputs=[z, fx_params]
+            outputs=[z, fx_params, json_output]
             + update_fx_outputs
-            + update_plots_outputs
-            + [json_output],
+            + update_plots_outputs,
         )
     extra_slider.input(
-        lambda *xs: chain_functions(
-            lambda args: update_z(args[0], args[1], args[2]),
+        chain_functions(
+            lambda z, s, i: update_z(z, s, i - 1),
             lambda z: (z, z2x(z)),
-            lambda args: [args[0], args[1], vec2fx(args[1])],
-            lambda args: args[:2]
-            + update_fx(args[2])
-            + update_plots(args[2])
-            + [model2json(args[2])],
-        )(xs),
+            lambda z, x: (z, x, vec2fx(x)),
+            lambda z, x, fx: [z, x, model2json(fx)] + update_fx(fx) + update_plots(fx),
+        ),
         inputs=[z, extra_slider, extra_pc_dropdown],
-        outputs=[z, fx_params]
-        + update_fx_outputs
-        + update_plots_outputs
-        + [json_output],
+        outputs=[z, fx_params, json_output] + update_fx_outputs + update_plots_outputs,
     )
 
     extra_pc_dropdown.input(