update app.py

app.py

@@ -73,6 +73,9 @@ EXAMPLE_PATH = "eleanor_erased.wav"
 with open(CONFIG_PATH["approx"]) as fp:
     fx_config = yaml.safe_load(fp)["model"]
 
+with open(CONFIG_PATH["realtime"]) as fp:
+    rt_config = yaml.safe_load(fp)["model"]
+
 
 def load_presets(preset_folder: Path) -> Tensor:
     raw_params = torch.from_numpy(np.load(preset_folder / PARAMS_PATH))
@@ -136,8 +139,18 @@ global_fx.load_state_dict(vec2dict(internal_mean), strict=False)
 meter = pyln.Meter(44100)
 
 
-@torch.no_grad()
-def inference(audio, ratio, fx):
+def inference(
+    audio,
+    ratio,
+    method,
+    dataset,
+    embedding,
+    remove_approx,
+    steps,
+    prior_weight,
+    optimiser,
+    lr,
+):
     sr, y = audio
     if sr != 44100:
         y = resample(y, sr, 44100)
@@ -153,7 +166,26 @@ def inference(audio, ratio, fx):
     if y.shape[1] != 1:
         y = y.mean(dim=1, keepdim=True)
 
-    direct, wet = fx(y)
+    fx = deepcopy(global_fx)
+    fx.train()
+
+    match method:
+        case "Mean":
+            vec = gaussian_params_dict[dataset][0]
+        case _:
+            vec = internal_mean.clone()
+
+    if remove_approx:
+        infer_fx = instantiate(rt_config)
+    else:
+        infer_fx = fx
+
+    infer_fx.load_state_dict(vec2dict(vec), strict=False)
+    # fx.apply(partial(clip_delay_eq_Q, Q=0.707))
+    infer_fx.eval()
+
+    with torch.no_grad():
+        direct, wet = fx(y)
     direct = direct.squeeze(0).T.numpy()
     wet = wet.squeeze(0).T.numpy()
     angle = ratio * math.pi * 0.5
@@ -327,8 +359,8 @@ def vec2fx(x):
 
 
 with gr.Blocks() as demo:
-    fx_params = gr.State(internal_mean)
-    fx = vec2fx(fx_params.value)
+    # fx_params = gr.State(internal_mean)
+    # fx = vec2fx(fx_params.value)
     # sr, y = read(EXAMPLE_PATH)
 
     default_pc_slider = partial(
@@ -357,14 +389,10 @@ with gr.Blocks() as demo:
                 label="Input Audio",
                 # value=(sr, y)
             )
-            with gr.Row():
-                reset_button = gr.Button(
-                    "Reset",
-                    elem_id="reset-button",
-                )
-                render_button = gr.Button(
-                    "Run", elem_id="render-button", variant="primary"
-                )
+            audio_reference = default_audio_block(
+                sources="upload",
+                label="Reference Audio",
+            )
 
         with gr.Column():
             audio_output = default_audio_block(label="Output Audio", interactive=False)
@@ -378,6 +406,10 @@ with gr.Blocks() as demo:
             direct_output = default_audio_block(label="Direct Audio", interactive=False)
             wet_output = default_audio_block(label="Wet Audio", interactive=False)
 
+    with gr.Row():
+        reset_button = gr.Button("Reset", elem_id="reset-button")
+        render_button = gr.Button("Run", elem_id="render-button", variant="primary")
+
     _ = gr.Markdown("## Common Parameters")
     with gr.Row():
         method_dropdown = gr.Dropdown(
@@ -387,10 +419,10 @@ with gr.Blocks() as demo:
             interactive=True,
         )
         dataset_dropdown = gr.Dropdown(
-            ["Internal", "MedleyDB"],
+            [("Internal", "internal"), ("MedleyDB", "medleydb")],
             label="Prior Distribution",
             info="When using the Regression method, this parameter has no effect as the model is trained on the internal dataset.",
-            value="Internal",
+            value="internal",
             interactive=True,
         )
         embedding_dropdown = gr.Dropdown(
@@ -400,6 +432,12 @@ with gr.Blocks() as demo:
             value="AFx-Rep",
             interactive=True,
         )
+        remove_approx_checkbox = gr.Checkbox(
+            label="Use Real-time Effects",
+            info="Use real-time delay and reverb effects instead of approximated ones.",
+            value=False,
+            interactive=True,
+        )
 
     _ = gr.Markdown("## Parameters for ST-ITO Method")
     with gr.Row():
@@ -435,5 +473,42 @@ with gr.Blocks() as demo:
             label="Optimiser",
             interactive=True,
         )
+        lr_slider = gr.Slider(
+            minimum=1e-6,
+            maximum=1.0,
+            value=1e-3,
+            label="Learning Rate",
+            interactive=True,
+        )
+
+    render_button.click(
+        chain_functions(
+            # lambda audio, ratio, x, *all_s: (
+            #     audio,
+            #     ratio,
+            #     # assign_fx_params(vec2fx(x), *all_s),
+            # ),
+            inference,
+        ),
+        inputs=[
+            audio_input,
+            dry_wet_ratio,
+            method_dropdown,
+            dataset_dropdown,
+            embedding_dropdown,
+            remove_approx_checkbox,
+            optimisation_steps,
+            prior_weight,
+            optimiser_dropdown,
+            lr_slider,
+            # fx_params,
+        ],
+        outputs=[
+            audio_output,
+            direct_output,
+            wet_output,
+        ],
+    )
+
 
 demo.launch()