Format choice

Allow the user to choose between wav or mp3

app.py

@@ -7,6 +7,7 @@ from huggingface_hub import snapshot_download
 from models import AudioDiffusion, DDPMScheduler
 from audioldm.audio.stft import TacotronSTFT
 from audioldm.variational_autoencoder import AutoencoderKL
+from pydub import AudioSegment
 from gradio import Markdown
 import spaces
 
@@ -83,12 +84,16 @@ tango.stft.to(device_type)
 tango.model.to(device_type)
     
 @spaces.GPU(duration=60)
-def gradio_generate(prompt, steps, guidance):
+def gradio_generate(prompt, output_format, steps, guidance):
     output_wave = tango.generate(prompt, steps, guidance)
     # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
     output_filename = "temp.wav"
     wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)
-    
+
+    if (output_format == "mp3"):
+        AudioSegment.from_wav("temp.wav").export("temp.mp3", format = "mp3")
+        output_filename = "temp.mp3"
+
     return output_filename
 
 # description_text = """
@@ -118,6 +123,7 @@ Generate audio using Tango2 by providing a text prompt. Tango2 was built from Ta
 """
 # Gradio input and output components
 input_text = gr.Textbox(lines=2, label="Prompt")
+output_format = gr.Radio(label = "Output format", info = "The file you can dowload", choices = ["mp3", "wav"], value = "wav")
 output_audio = gr.Audio(label="Generated Audio", type="filepath")
 denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
 guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)
@@ -125,7 +131,7 @@ guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guid
 # Gradio interface
 gr_interface = gr.Interface(
     fn=gradio_generate,
-    inputs=[input_text, denoising_steps, guidance_scale],
+    inputs=[input_text, output_format, denoising_steps, guidance_scale],
     outputs=[output_audio],
     title="Tango 2: Aligning Diffusion-based Text-to-Audio Generations through Direct Preference Optimization",
     description=description_text,