Update app.py

app.py

@@ -103,7 +103,7 @@ def main():
             wav, sample_rate = torchaudio.load("test.wav")
             st.markdown("Before unsquueze wav")
             st.markdown(wav)
-
+            file_extension_ori =".wav"
             #Unsqueeze for line 176
             wav= wav.unsqueeze(0)
                 
@@ -123,6 +123,7 @@ def main():
             audio.export("test.wav", format="wav")
             wav3, sample_rate = torchaudio.load("test.wav")
             wav= wav3.unsqueeze(0)
+            file_extension_ori =".mp3"
             file_extension =".wav"
 
             #RuntimeError: Could not infer dtype of numpy.float32
@@ -160,7 +161,7 @@ def main():
                         #st.markdown(shape)
 
                         #st.markdown(squeeze)
-                    if file_extension in [".wav", ".flac"]:
+                    if file_extension_ori in [".wav", ".flac"]:
                         torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
                         watermarked_wav = torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
 
@@ -178,9 +179,13 @@ def main():
                             )
 
 
-                    elif file_extension == ".mp3":
-                        torchaudio.save("output.mp3", squeeze, default_sr)
-                        watermarked_mp3 = torchaudio.save("output.mp3", squeeze, default_sr)
+                    elif file_extension_ori == ".mp3":
+                        torchaudio.save("output.wav", squeeze, default_sr)
+                        watermarked_mp3 = torchaudio.save("output.wav", squeeze, default_sr)
+                        audio = AudioSegment.from_wav("output.wav")
+
+                        # Export as MP3
+                        audio.export("output.mp3", format="mp3")
                         st.audio("output.mp3", format="audio/mpeg")
                         
                         with open("output.mp3", "rb") as file: