Update app.py

app.py

@@ -49,9 +49,11 @@ def encode_audio(audio_file_path):
         # Convert tokens to NumPy and save to .owie file
         tokens_numpy = tokens.detach().cpu().numpy()
 
-        # Ensure tokens_numpy is 1D
-        if tokens_numpy.ndim != 1:
-            raise ValueError("Tokens array must be 1D")
+        # 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")
 
         # Create temporary .owie file
         temp_fd, temp_file_path = tempfile.mkstemp(suffix=".owie")
@@ -67,10 +69,18 @@ def encode_audio(audio_file_path):
 
     except Exception as e:
         print(f"Encoding error: {e}")
-        return str(e)  # Return error message as string
+        return None  # Return None instead of the error message
 
     finally:
         cancel_encode = False  # Reset cancel flag after encoding
+        if 'temp_wav_file_path' in locals():
+            os.remove(temp_wav_file_path)  # Clean up temporary WAV file
+
+# Add this function to handle the output
+def handle_encode_output(file_path):
+    if file_path is None:
+        return None, gr.Markdown("Encoding failed. Please check the input file and try again.", visible=True)
+    return file_path, gr.Markdown(visible=False)
 
 @spaces.GPU(duration=250)
 def decode_audio(encoded_file_path):
@@ -156,8 +166,13 @@ with gr.Blocks() as demo:
         encode_button = gr.Button("Encode")
         cancel_encode_button = gr.Button("Cancel")
         encoded_output = gr.File(label="Encoded File (.owie)", type="filepath")  # Using "filepath" mode
+        encode_error_message = gr.Markdown(visible=False)
 
-        encode_button.click(encode_audio, inputs=input_audio, outputs=encoded_output)
+        encode_button.click(
+            lambda x: handle_encode_output(encode_audio(x)),
+            inputs=input_audio,
+            outputs=[encoded_output, encode_error_message]
+        )
         cancel_encode_button.click(lambda: globals().update(cancel_encode=True), outputs=None)  # Set cancel_encode flag
 
     with gr.Tab("Decode"):
@@ -178,4 +193,4 @@ with gr.Blocks() as demo:
         stream_button.click(stream_decode_audio, inputs=input_encoded_stream, outputs=audio_output)
         cancel_stream_button.click(lambda: globals().update(cancel_stream=True), outputs=None)  # Set cancel_stream flag
 
-demo.queue().launch()
+demo.queue().launch()