1, 2 or 3 output files

app.py

@@ -49,12 +49,12 @@ class Tango:
         self.scheduler = DDPMScheduler.from_pretrained(main_config["scheduler_name"], subfolder="scheduler")
         
     def chunks(self, lst, n):
-        """ Yield successive n-sized chunks from a list. """
+        # Yield successive n-sized chunks from a list
         for i in range(0, len(lst), n):
             yield lst[i:i + n]
         
     def generate(self, prompt, steps=100, guidance=3, samples=3, disable_progress=True):
-        """ Genrate audio for a single prompt string. """
+        # Genrate audio for a single prompt string
         with torch.no_grad():
             latents = self.model.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
             mel = self.vae.decode_first_stage(latents)
@@ -62,7 +62,7 @@ class Tango:
         return wave
     
     def generate_for_batch(self, prompts, steps=200, guidance=3, samples=1, batch_size=8, disable_progress=True):
-        """ Genrate audio for a list of prompt strings. """
+        # Genrate audio for a list of prompt strings
         outputs = []
         for k in tqdm(range(0, len(prompts), batch_size)):
             batch = prompts[k: k+batch_size]
@@ -84,24 +84,42 @@ tango.stft.to(device_type)
 tango.model.to(device_type)
     
 @spaces.GPU(duration=120)
-def gradio_generate(prompt, output_format, steps, guidance):
-    output_wave = tango.generate(prompt, steps, guidance)
+def gradio_generate(
+    prompt,
+    output_format,
+    output_number,
+    steps,
+    guidance
+):
+    output_wave = tango.generate(prompt, steps, guidance, output_number)
     # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
 
     output_filename_1 = "tmp1.wav"
-    wavio.write(output_filename_1, output_wave[0], rate=16000, sampwidth=2)
-    output_filename_2 = "tmp2.wav"
-    wavio.write(output_filename_2, output_wave[1], rate=16000, sampwidth=2)
-    output_filename_3 = "tmp3.wav"
-    wavio.write(output_filename_3, output_wave[2], rate=16000, sampwidth=2)
+    wavio.write(output_filename_1, output_wave[0], rate = 16000, sampwidth = 2)
 
     if (output_format == "mp3"):
         AudioSegment.from_wav("tmp1.wav").export("tmp1.mp3", format = "mp3")
         output_filename_1 = "tmp1.mp3"
-        AudioSegment.from_wav("tmp2.wav").export("tmp2.mp3", format = "mp3")
-        output_filename_2 = "tmp2.mp3"
-        AudioSegment.from_wav("tmp3.wav").export("tmp3.mp3", format = "mp3")
-        output_filename_3 = "tmp3.mp3"
+
+    if (2 <= output_number):
+        output_filename_2 = "tmp2.wav"
+        wavio.write(output_filename_2, output_wave[1], rate = 16000, sampwidth = 2)
+
+        if (output_format == "mp3"):
+            AudioSegment.from_wav("tmp2.wav").export("tmp2.mp3", format = "mp3")
+            output_filename_2 = "tmp2.mp3"
+    else:
+        output_filename_2 = None
+
+    if (output_number == 3):
+        output_filename_3 = "tmp3.wav"
+        wavio.write(output_filename_3, output_wave[2], rate = 16000, sampwidth = 2)
+
+        if (output_format == "mp3"):
+            AudioSegment.from_wav("tmp3.wav").export("tmp3.mp3", format = "mp3")
+            output_filename_3 = "tmp3.mp3"
+    else:
+        output_filename_3 = None
 
     return [output_filename_1, output_filename_2, output_filename_3]
 
@@ -133,16 +151,17 @@ 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 download", choices = ["mp3", "wav"], value = "wav")
+output_number = gr.Slider(label = "Number of generations", info = "1, 2 or 3 output files", minimum = 1, maximum = 3, value = 3, step = 1, interactive = True)
 output_audio_1 = gr.Audio(label="Generated Audio #1/3", type="filepath")
 output_audio_2 = gr.Audio(label="Generated Audio #2/3", type="filepath")
 output_audio_3 = gr.Audio(label="Generated Audio #3/3", type="filepath")
-denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
+denoising_steps = gr.Slider(minimum=10, 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)
 
 # Gradio interface
 gr_interface = gr.Interface(
     fn=gradio_generate,
-    inputs=[input_text, output_format, denoising_steps, guidance_scale],
+    inputs=[input_text, output_format, output_number, denoising_steps, guidance_scale],
     outputs=[output_audio_1, output_audio_2, output_audio_3],
     title="Tango 2: Aligning Diffusion-based Text-to-Audio Generations through Direct Preference Optimization",
     description=description_text,