Update app.py

app.py

@@ -214,10 +214,14 @@ def main():
                         #st.markdown(shape)
 
                         #st.markdown(squeeze)
-                        
+                    if file_extension in [".wav", ".flac"]:
                         torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
                         st.audio("output.wav", format="audio/wav")
 
+                    elif file_extension == ".mp3":
+                        torchaudio.save("output.mp3", squeeze, default_sr, bits_per_sample=16)
+                        st.audio("output.wav", format="audio/mpeg")
+
                         # st.download_button(
                         #     label="Download Watermarked audio",
                         #     data="output.wav",