Update app.py

app.py

@@ -53,11 +53,19 @@ def encode_audio(audio_file_path):
         # Convert tokens to NumPy and save to .owie file
         tokens_numpy = tokens.detach().cpu().numpy()
 
+        print(f"Original tokens shape: {tokens_numpy.shape}")
+
         # Ensure tokens_numpy is 2D
         if tokens_numpy.ndim == 1:
             tokens_numpy = tokens_numpy.reshape(1, -1)
-        elif tokens_numpy.ndim > 2:
-            raise ValueError("Tokens array must be 1D or 2D")
+        elif tokens_numpy.ndim == 2:
+            pass  # Already 2D
+        elif tokens_numpy.ndim == 3 and tokens_numpy.shape[0] == 1:
+            tokens_numpy = tokens_numpy.squeeze(0)
+        else:
+            raise ValueError(f"Unexpected tokens array shape: {tokens_numpy.shape}")
+
+        print(f"Reshaped tokens shape: {tokens_numpy.shape}")
 
         # Create temporary .owie file
         temp_fd, temp_file_path = tempfile.mkstemp(suffix=".owie")