modify app

app.py

@@ -16,7 +16,17 @@ def process_audio(input_audio, reference_audio):
     
     param_output = mastering_transfer.get_param_output_string(predicted_params)
     
-    return "output_mastered.wav", param_output
+    # Convert output_audio to numpy array if it's a tensor
+    if isinstance(output_audio, torch.Tensor):
+        output_audio = output_audio.cpu().numpy()
+    
+    # Ensure the audio is in the correct shape (samples, channels)
+    if output_audio.ndim == 1:
+        output_audio = output_audio.reshape(-1, 1)
+    elif output_audio.ndim > 2:
+        output_audio = output_audio.squeeze()
+    
+    return (sr, output_audio), param_output
 
 def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
     if ito_reference_audio is None:
@@ -36,13 +46,24 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
 
     initial_reference_feature = mastering_transfer.get_reference_embedding(reference_tensor)
 
-    ito_output, ito_params, optimized_embedding, steps_taken, ito_log = mastering_transfer.inference_time_optimization(
+    ito_log = ""
+    for log_entry, current_output, current_params, step in mastering_transfer.inference_time_optimization(
         input_tensor, ito_reference_tensor, ito_config, initial_reference_feature
-    )
-
-    ito_param_output = mastering_transfer.get_param_output_string(ito_params)
-    
-    return "ito_output_mastered.wav", ito_param_output, steps_taken, ito_log
+    ):
+        ito_log += log_entry
+        ito_param_output = mastering_transfer.get_param_output_string(current_params)
+        
+        # Convert current_output to numpy array if it's a tensor
+        if isinstance(current_output, torch.Tensor):
+            current_output = current_output.cpu().numpy()
+        
+        # Ensure the audio is in the correct shape (samples, channels)
+        if current_output.ndim == 1:
+            current_output = current_output.reshape(-1, 1)
+        elif current_output.ndim > 2:
+            current_output = current_output.squeeze()
+        
+        yield (args.sample_rate, current_output), ito_param_output, step, ito_log
 
 with gr.Blocks() as demo:
     gr.Markdown("# Mastering Style Transfer Demo")
@@ -64,38 +85,37 @@ with gr.Blocks() as demo:
             outputs=[output_audio, param_output]
         )
 
-        gr.Markdown("## Inference Time Optimization (ITO)")
-        
-        with gr.Row():
-            with gr.Column(scale=2):
-                ito_reference_audio = gr.Audio(label="ITO Reference Audio (optional)")
-                num_steps = gr.Slider(minimum=1, maximum=1000, value=100, step=1, label="Number of Steps")
-                optimizer = gr.Dropdown(["Adam", "RAdam", "SGD"], value="RAdam", label="Optimizer")
-                learning_rate = gr.Slider(minimum=0.0001, maximum=0.1, value=0.001, step=0.0001, label="Learning Rate")
-                af_weights = gr.Textbox(label="AudioFeatureLoss Weights (comma-separated)", value="0.1,0.001,1.0,1.0,0.1")
-                
-                ito_button = gr.Button("Perform ITO")
-                
-                ito_output_audio = gr.Audio(label="ITO Output Audio")
-                ito_param_output = gr.Textbox(label="ITO Predicted Parameters", lines=10)
-                ito_steps_taken = gr.Number(label="ITO Steps Taken")
+    gr.Markdown("## Inference Time Optimization (ITO)")
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            ito_reference_audio = gr.Audio(label="ITO Reference Audio (optional)")
+            num_steps = gr.Slider(minimum=1, maximum=1000, value=100, step=1, label="Number of Steps")
+            optimizer = gr.Dropdown(["Adam", "RAdam", "SGD"], value="RAdam", label="Optimizer")
+            learning_rate = gr.Slider(minimum=0.0001, maximum=0.1, value=0.001, step=0.0001, label="Learning Rate")
+            af_weights = gr.Textbox(label="AudioFeatureLoss Weights (comma-separated)", value="0.1,0.001,1.0,1.0,0.1")
+            
+            ito_button = gr.Button("Perform ITO")
             
-            with gr.Column(scale=1):
-                ito_log = gr.Textbox(label="ITO Log", lines=30)
-
-        def run_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
-            af_weights = [float(w.strip()) for w in af_weights.split(',')]
-            ito_output, ito_params, steps_taken, log = perform_ito(
-                input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights
-            )
-            return ito_output, ito_params, steps_taken, log
-
-        ito_button.click(
-            run_ito,
-            inputs=[input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights],
-            outputs=[ito_output_audio, ito_param_output, ito_steps_taken, ito_log]
+            ito_output_audio = gr.Audio(label="ITO Output Audio")
+            ito_param_output = gr.Textbox(label="ITO Predicted Parameters", lines=10)
+            ito_steps_taken = gr.Number(label="ITO Steps Taken")
+        
+        with gr.Column(scale=1):
+            ito_log = gr.Textbox(label="ITO Log", lines=30)
+
+    def run_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
+        af_weights = [float(w.strip()) for w in af_weights.split(',')]
+        return perform_ito(
+            input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights
         )
 
+    ito_button.click(
+        run_ito,
+        inputs=[input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights],
+        outputs=[ito_output_audio, ito_param_output, ito_steps_taken, ito_log]
+    )
+
 demo.launch()
 
 

inference.py

@@ -110,7 +110,7 @@ class MasteringStyleTransfer:
                 initial_params = current_params
             top_10_diff = self.get_top_10_diff_string(initial_params, current_params)
             log_entry = f"Step {step + 1}, Loss: {total_loss.item():.4f}\n{top_10_diff}\n"
-            ito_log.append(log_entry)
+            yield log_entry, output_audio, current_params, step + 1
 
             if divergence_counter >= 10:
                 print(f"Optimization stopped early due to divergence at step {step}")
@@ -119,7 +119,7 @@ class MasteringStyleTransfer:
             total_loss.backward()
             optimizer.step()
 
-        return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1, "\n".join(ito_log)
+        return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1
 
     def preprocess_audio(self, audio, target_sample_rate=44100):
         sample_rate, data = audio