modify app

app.py

@@ -139,11 +139,31 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
 
         yield (args.sample_rate, current_output), ito_param_output, step, ito_log, pd.DataFrame(loss_values)
 
+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_generator = mastering_transfer.inference_time_optimization(
+        input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights
+    )
+    
+    all_results = []
+    for result in ito_generator:
+        all_results.append(result)
+    
+    min_loss_step = min(range(len(all_results)), key=lambda i: all_results[i]['loss'])
+    
+    loss_df = pd.DataFrame([(r['step'], r['loss']) for r in all_results], columns=['step', 'loss'])
+    
+    return all_results, min_loss_step, loss_df
 
-""" APP display """
-gr.Markdown("# ITO-Master: Inference Time Optimization for Mastering Style Transfer")
+def update_ito_output(all_results, selected_step):
+    selected_result = all_results[selected_step]
+    return (args.sample_rate, selected_result['audio']), selected_result['params'], selected_result['log']
 
+
+""" APP display """
 with gr.Blocks() as demo:
+    gr.Markdown("# ITO-Master: Inference Time Optimization for Mastering Style Transfer")
+
     gr.Markdown("# Step 1: Mastering Style Transfer")
 
     with gr.Tab("Upload Audio"):
@@ -207,6 +227,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             ito_output_audio = gr.Audio(label="ITO Output Audio")
             ito_param_output = gr.Textbox(label="ITO Predicted Parameters", lines=15)
+            ito_step_slider = gr.Slider(minimum=1, maximum=100, step=1, label="ITO Step", interactive=True)
         with gr.Column():
             ito_loss_plot = gr.LinePlot(
                 x="step",
@@ -219,39 +240,28 @@ with gr.Blocks() as demo:
             )
             ito_log = gr.Textbox(label="ITO Log", lines=10)
 
-    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_generator = perform_ito(
-            input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights
-        )
-        
-        # Initialize variables to store the final results
-        final_audio = None
-        final_params = None
-        final_steps = 0
-        final_log = ""
-        loss_df = None
-        
-        # Iterate through the generator to get the final results
-        for audio, params, steps, log, loss_data in ito_generator:
-            final_audio = audio
-            final_params = params
-            final_steps = steps
-            final_log = log
-            loss_df = loss_data
-
-        # Calculate y_min and y_max
-        y_min = loss_df['loss'].min()
-        y_max = loss_df['loss'].max()
-
-        # Return the plot configuration along with the data
-        return final_audio, final_params, final_log, loss_df
+    all_results = gr.State([])
+    min_loss_step = gr.State(0)
 
+    def on_ito_complete(results, min_step, loss_df):
+        all_results.value = results
+        min_loss_step.value = min_step
+        return loss_df, gr.update(maximum=len(results), value=min_step+1)
 
     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_log, ito_loss_plot]
+        outputs=[all_results, min_loss_step, ito_loss_plot, ito_step_slider]
+    ).then(
+        update_ito_output,
+        inputs=[all_results, ito_step_slider],
+        outputs=[ito_output_audio, ito_param_output, ito_log]
+    )
+
+    ito_step_slider.change(
+        update_ito_output,
+        inputs=[all_results, ito_step_slider],
+        outputs=[ito_output_audio, ito_param_output, ito_log]
     )
 
 demo.launch()

inference.py

@@ -60,6 +60,68 @@ class MasteringStyleTransfer:
             predicted_params = self.mastering_converter.get_last_predicted_params()
         return output_audio, predicted_params
 
+    # def inference_time_optimization(self, input_tensor, reference_tensor, ito_config, initial_reference_feature):
+    #     fit_embedding = torch.nn.Parameter(initial_reference_feature)
+    #     optimizer = getattr(torch.optim, ito_config['optimizer'])([fit_embedding], lr=ito_config['learning_rate'])
+
+    #     af_loss = AudioFeatureLoss(
+    #         weights=ito_config['af_weights'],
+    #         sample_rate=ito_config['sample_rate'],
+    #         stem_separation=False,
+    #         use_clap=False
+    #     )
+
+    #     min_loss = float('inf')
+    #     min_loss_step = 0
+    #     min_loss_output = None
+    #     min_loss_params = None
+    #     min_loss_embedding = None
+
+    #     loss_history = []
+    #     divergence_counter = 0
+    #     ito_log = []
+
+    #     for step in range(ito_config['num_steps']):
+    #         optimizer.zero_grad()
+
+    #         output_audio = self.mastering_converter(input_tensor, fit_embedding)
+    #         current_params = self.mastering_converter.get_last_predicted_params()
+
+    #         losses = af_loss(output_audio, reference_tensor)
+    #         total_loss = sum(losses.values())
+
+    #         loss_history.append(total_loss.item())
+
+    #         if total_loss < min_loss:
+    #             min_loss = total_loss.item()
+    #             min_loss_step = step
+    #             min_loss_output = output_audio.detach()
+    #             min_loss_params = current_params
+    #             min_loss_embedding = fit_embedding.detach().clone()
+
+    #         # Check for divergence
+    #         if len(loss_history) > 10 and total_loss > loss_history[-11]:
+    #             divergence_counter += 1
+    #         else:
+    #             divergence_counter = 0
+
+    #         # Log top 5 parameter differences
+    #         if step == 0:
+    #             initial_params = current_params
+    #         top_5_diff = self.get_top_n_diff_string(initial_params, current_params, top_n=5)
+    #         log_entry = f"Step {step + 1}\n   Loss: {total_loss.item():.4f}\n{top_5_diff}\n"
+
+    #         if divergence_counter >= 10:
+    #             print(f"Optimization stopped early due to divergence at step {step}")
+    #             break
+
+    #         total_loss.backward()
+    #         optimizer.step()
+
+    #         yield log_entry, output_audio.detach(), current_params, step + 1, total_loss.item()
+
+    #     return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1
+
     def inference_time_optimization(self, input_tensor, reference_tensor, ito_config, initial_reference_feature):
         fit_embedding = torch.nn.Parameter(initial_reference_feature)
         optimizer = getattr(torch.optim, ito_config['optimizer'])([fit_embedding], lr=ito_config['learning_rate'])
@@ -73,13 +135,7 @@ class MasteringStyleTransfer:
 
         min_loss = float('inf')
         min_loss_step = 0
-        min_loss_output = None
-        min_loss_params = None
-        min_loss_embedding = None
-
-        loss_history = []
-        divergence_counter = 0
-        ito_log = []
+        all_results = []
 
         for step in range(ito_config['num_steps']):
             optimizer.zero_grad()
@@ -90,20 +146,9 @@ class MasteringStyleTransfer:
             losses = af_loss(output_audio, reference_tensor)
             total_loss = sum(losses.values())
 
-            loss_history.append(total_loss.item())
-
             if total_loss < min_loss:
                 min_loss = total_loss.item()
                 min_loss_step = step
-                min_loss_output = output_audio.detach()
-                min_loss_params = current_params
-                min_loss_embedding = fit_embedding.detach().clone()
-
-            # Check for divergence
-            if len(loss_history) > 10 and total_loss > loss_history[-11]:
-                divergence_counter += 1
-            else:
-                divergence_counter = 0
 
             # Log top 5 parameter differences
             if step == 0:
@@ -111,16 +156,20 @@ class MasteringStyleTransfer:
             top_5_diff = self.get_top_n_diff_string(initial_params, current_params, top_n=5)
             log_entry = f"Step {step + 1}\n   Loss: {total_loss.item():.4f}\n{top_5_diff}\n"
 
-            if divergence_counter >= 10:
-                print(f"Optimization stopped early due to divergence at step {step}")
-                break
+            all_results.append({
+                'step': step + 1,
+                'loss': total_loss.item(),
+                'audio': output_audio.detach(),
+                'params': current_params,
+                'log': log_entry
+            })
 
             total_loss.backward()
             optimizer.step()
 
-            yield log_entry, output_audio.detach(), current_params, step + 1, total_loss.item()
+            yield all_results[-1]
 
-        return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1
+        return all_results, min_loss_step
 
     def preprocess_audio(self, audio, target_sample_rate=44100):
         sample_rate, data = audio