Format output

app.py

@@ -1,4 +1,8 @@
 import gradio as gr
+import json
+import torch
+import numpy as np
+import soundfile as sf
 from src.config import BaseConfig
 from src.reverb import BaseFDN
 from flamo.optimize.trainer import Trainer
@@ -76,16 +80,53 @@ def process_fdn(N, delay_lengths):
             # Train the model
             print("Starting training...")
             trainer.train(train_loader, valid_loader)
+            est_param = model.get_params()
+            
+            # Convert parameters to JSON format
+            # Assuming est_param is a dict or can be converted to one
+            param_dict = {}
+            for key, value in est_param.items():
+                # Convert tensors to lists for JSON serialization
+                if hasattr(value, 'cpu'):
+                    param_dict[key] = value.cpu().detach().numpy().tolist()
+                else:
+                    param_dict[key] = value
+            
+            # Save to JSON file
+            output_path = "estimated_parameters.json"
+            with open(output_path, 'w') as f:
+                json.dump(param_dict, f, indent=2)
 
-            return result
+            ir = model.shell.get_time_response()
+            
+            # Convert ir to audio format for Gradio
+            ir_audio = ir.cpu().detach().numpy()
+            
+            # Ensure proper shape (1D array)
+            if ir_audio.ndim > 1:
+                ir_audio = ir_audio.squeeze()
+            
+            # Normalize to [-1, 1] range to prevent overflow
+            max_val = np.abs(ir_audio).max()
+            if max_val > 0:
+                ir_audio = ir_audio / max_val
+            
+            # Get the sample rate from config
+            sample_rate = getattr(config, 'fs', 48000)
+            
+            # Save audio to file using soundfile (avoids Gradio's conversion issues)
+            audio_path = "impulse_response.wav"
+            sf.write(audio_path, ir_audio, sample_rate)
+            
+            return result, output_path, audio_path
         else:
-            return "Error: No delay lengths provided"
+            return "Error: No delay lengths provided", None, None
             
     except Exception as e:
         print(f"Error: {e}")
         import traceback
         traceback.print_exc()
-        return f"Error processing inputs: {str(e)}"
+        return f"Error processing inputs: {str(e)}", None, None
 
 demo = gr.Interface(
     fn=process_fdn,
@@ -99,9 +140,13 @@ demo = gr.Interface(
             label="Delay Lengths (N integer values)"
         )
     ],
-    outputs=gr.Textbox(label="Output"),
-    title="Feedback Delay Network Configuration",
-    description="Configure your feedback delay network by specifying N (number of delay lines) and their corresponding delay lengths."
+    outputs=[
+        gr.Textbox(label="Output"),
+        gr.File(label="Estimated Parameters (JSON)"),
+        gr.Audio(label="Impulse Response", type="numpy")
+    ],
+    title="Feedback Delay Network Optimization",
+    description="Configure your homogeneous feedback delay network by specifying N (number of delay lines) and their corresponding delay lengths. Submit the values to run optimization and obtain estimated parameters and playback the resulting impulse response."
 )
 
 demo.launch(debug=True)

src/reverb.py

@@ -104,7 +104,7 @@ class BaseFDN(nn.Module):
         - The returned parameters can be used to recreate or modify the FDN
         """
         core = self.shell.get_core()
-
+        map_matrix = core.branchA.feedback_loop.feedback.mixing_matrix.map
         params = OrderedDict()
         params["delays"] = self.delay_lengths.cpu().numpy().tolist()
         params["onset_time"] = self.onset
@@ -112,23 +112,24 @@ class BaseFDN(nn.Module):
             core.branchB.early_reflections.param.cpu().detach().numpy().tolist()
         )
         params["input_gains"] = (
-            core.branchA.input_gain.param.cpu().detach().numpy().tolist()
+            core.branchA.input_gain.param.cpu().squeeze().detach().numpy().tolist()
         )
         params["output_gains"] = (
-            core.branchA.output_gain.param[0].cpu().detach().numpy().tolist()
+            core.branchA.output_gain.param[0].cpu().squeeze().detach().numpy().tolist()
         )
         params["feedback_matrix"] = (
-            core.branchA.feedback_loop.feedback.mixing_matrix.param.cpu()
-            .detach()
-            .numpy()
-            .tolist()
-        )
-        params["attenuation"] = (
-            core.branchA.feedback_loop.feedback.attenuation.param.cpu()
+            map_matrix(core.branchA.feedback_loop.feedback.mixing_matrix.param).cpu()
             .detach()
+            .squeeze()
             .numpy()
             .tolist()
         )
+        # params["attenuation"] = (
+        #     core.branchA.feedback_loop.feedback.attenuation.param.cpu()
+        #     .detach()
+        #     .numpy()
+        #     .tolist()
+        # )
         return params
 
     def _validate_delays(self, config: BaseConfig, delay_lengths: List[int]) -> None: