modify app

app.py

@@ -81,14 +81,11 @@ 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)
         
-        if output_audio.ndim == 1:
-            output_audio = output_audio.reshape(-1, 1)
-        elif output_audio.ndim > 2:
-            output_audio = output_audio.squeeze()
-
         # Ensure the audio is in the correct shape (samples, channels)
-        if output_audio.shape[1] > output_audio.shape[0]:
-            output_audio = output_audio.transpose(1,0)
+        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
 
@@ -135,9 +132,24 @@ with gr.Blocks() as demo:
 
     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(
+        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 = ""
+        
+        # Iterate through the generator to get the final results
+        for audio, params, steps, log in ito_generator:
+            final_audio = audio
+            final_params = params
+            final_steps = steps
+            final_log = log
+        
+        return final_audio, final_params, final_steps, final_log
 
     ito_button.click(
         run_ito,