modify app

app.py

@@ -112,16 +112,19 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
 
     initial_reference_feature = mastering_transfer.get_reference_embedding(reference_tensor)
 
+    all_results, min_loss_step = mastering_transfer.inference_time_optimization(
+        input_tensor, ito_reference_tensor, ito_config, initial_reference_feature
+    )
+
     ito_log = ""
     loss_values = []
-    all_results = []
-    for log_entry, current_output, current_params, step, loss in mastering_transfer.inference_time_optimization(
-        input_tensor, ito_reference_tensor, ito_config, initial_reference_feature
-    ):
-        ito_log += log_entry
-        ito_param_output = mastering_transfer.get_param_output_string(current_params)
-        loss_values.append({"step": int(step), "loss": loss})
+    for result in all_results:
+        ito_log += result['log']
+        loss_values.append({"step": result['step'], "loss": result['loss']})
         
+        current_output = result['audio']
+        ito_param_output = mastering_transfer.get_param_output_string(result['params'])
+
         # Convert current_output to numpy array if it's a tensor
         if isinstance(current_output, torch.Tensor):
             current_output = current_output.cpu().numpy()
@@ -139,15 +142,7 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
         # Denormalize the audio to int16
         current_output = denormalize_audio(current_output, dtype=np.int16)
 
-        all_results.append({
-            'step': step,
-            'audio': current_output,
-            'params': ito_param_output,
-            'log': log_entry,
-            'loss': loss
-        })
-
-        yield (args.sample_rate, current_output), ito_param_output, step, ito_log, pd.DataFrame(loss_values), all_results
+        yield (args.sample_rate, current_output), ito_param_output, result['step'], ito_log, pd.DataFrame(loss_values), all_results
 
 def update_ito_output(all_results, selected_step):
     print(all_results)