Refactor

app.py

@@ -60,7 +60,7 @@ eigvecs = np.flip(eigvecs, axis=1)[:, :75]
 U = eigvecs * np.sqrt(eigvals)
 U = torch.from_numpy(U).float()
 mean = torch.from_numpy(mean).float()
-
+z = torch.zeros(75)
 
 with open(INFO_PATH) as f:
     info = json.load(f)
@@ -93,7 +93,7 @@ meter = pyln.Meter(44100)
 
 
 @torch.no_grad()
-def inference(audio, randomise_rest, *pcs):
+def inference(audio):
     sr, y = audio
     if sr != 44100:
         y = resample(y, sr, 44100)
@@ -109,18 +109,19 @@ def inference(audio, randomise_rest, *pcs):
     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))
-            ),
-        ]
-    )
+    # 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))
@@ -179,7 +180,7 @@ with gr.Blocks() as demo:
             audio_input = gr.Audio(type="numpy", sources="upload", label="Input Audio")
             with gr.Row():
                 random_button = gr.Button(
-                    f"Randomise the first {NUMBER_OF_PCS} PCs",
+                    f"Randomise PCs",
                     elem_id="randomise-button",
                 )
                 reset_button = gr.Button(
@@ -189,12 +190,26 @@ 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",
-            )
+            # 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)
+
+            extra_pc_dropdown = gr.Dropdown(
+                list(range(NUMBER_OF_PCS + 1, 76)),
+                label=f"PC > {NUMBER_OF_PCS}",
+                info="Select which extra PC to adjust",
+                interactive=True,
+            )
+            extra_slider = gr.Slider(
+                minimum=SLIDER_MIN,
+                maximum=SLIDER_MAX,
+                label="Extra PC",
+                value=0,
+            )
+
         with gr.Column():
             audio_output = gr.Audio(
                 type="numpy", label="Output Audio", interactive=False
@@ -205,27 +220,47 @@ with gr.Blocks() as demo:
         lambda *args: (lambda x: (x, model2json()))(inference(*args)),
         inputs=[
             audio_input,
-            random_rest_checkbox,
+            # random_rest_checkbox,
         ]
-        + sliders,
+        # + sliders,
+        ,
         outputs=[audio_output, json_output],
     )
 
     random_button.click(
-        lambda *xs: [
-            chain_functions(
-                partial(max, SLIDER_MIN),
-                partial(min, SLIDER_MAX),
-            )(normalvariate(0, 1))
-            for _ in range(len(xs))
-        ],
-        inputs=sliders,
-        outputs=sliders,
+        # lambda *xs: [
+        #     chain_functions(
+        #         partial(max, SLIDER_MIN),
+        #         partial(min, SLIDER_MAX),
+        #     )(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)
+        ),
+        inputs=extra_pc_dropdown,
+        outputs=sliders + [extra_slider],
     )
     reset_button.click(
-        lambda *xs: [0 for _ in range(len(xs))],
-        inputs=sliders,
-        outputs=sliders,
+        lambda *xs: (lambda _: [0 for _ in range(len(xs))])(z.zero_()),
+        inputs=sliders + [extra_slider],
+        outputs=sliders + [extra_slider],
+    )
+
+    def update_z(s, i):
+        z[i] = s
+        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]
+    )
+
+    extra_pc_dropdown.change(
+        lambda i: z[i - 1].item(),
+        inputs=extra_pc_dropdown,
+        outputs=extra_slider,
     )
 
 demo.launch()