modify app

app.py

@@ -35,14 +35,15 @@ def process_audio(input_audio, reference_audio):
     # Denormalize the audio to int16
     output_audio = denormalize_audio(output_audio, dtype=np.int16)
     
-    # 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 and output_audio.shape[1] > output_audio.shape[0]:
-        output_audio = output_audio.transpose(1,0)
     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)
+
     print(output_audio.shape)
     print(param_output)
     
@@ -80,13 +81,14 @@ 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 current_output.ndim == 1:
-            current_output = current_output.reshape(-1, 1)
-        elif current_output.ndim == 2 and current_output.shape[1] > current_output.shape[0]:
-            current_output = current_output.transpose(1,0)
-        elif current_output.ndim > 2:
-            current_output = current_output.squeeze()
+        if output_audio.shape[1] > output_audio.shape[0]:
+            output_audio = output_audio.transpose(1,0)
         
         yield (args.sample_rate, current_output), ito_param_output, step, ito_log