Automatic device detection

app.py

@@ -2,18 +2,27 @@ import gradio as gr
 import json
 import torch
 import wavio
+import spaces
+
 from tqdm import tqdm
 from huggingface_hub import snapshot_download
 from models import AudioDiffusion, DDPMScheduler
 from audioldm.audio.stft import TacotronSTFT
 from audioldm.variational_autoencoder import AutoencoderKL
 from gradio import Markdown
-import spaces
+
+# Automatic device detection
+if torch.cuda.is_available():
+    device_type = "cuda"
+    device_selection = "cuda:0"
+else:
+    device_type = "cpu"
+    device_selection = "cpu"
 
 class Tango:
-    def __init__(self, name="declare-lab/tango2", device="cpu"):
+    def __init__(self, name = "declare-lab/tango2", device = device_selection):
         
-        path = snapshot_download(repo_id=name)
+        path = snapshot_download(repo_id = name)
         
         vae_config = json.load(open("{}/vae_config.json".format(path)))
         stft_config = json.load(open("{}/stft_config.json".format(path)))
@@ -23,9 +32,9 @@ class Tango:
         self.stft = TacotronSTFT(**stft_config).to(device)
         self.model = AudioDiffusion(**main_config).to(device)
         
-        vae_weights = torch.load("{}/pytorch_model_vae.bin".format(path), map_location=device)
-        stft_weights = torch.load("{}/pytorch_model_stft.bin".format(path), map_location=device)
-        main_weights = torch.load("{}/pytorch_model_main.bin".format(path), map_location=device)
+        vae_weights = torch.load("{}/pytorch_model_vae.bin".format(path), map_location = device)
+        stft_weights = torch.load("{}/pytorch_model_stft.bin".format(path), map_location = device)
+        main_weights = torch.load("{}/pytorch_model_main.bin".format(path), map_location = device)
         
         self.vae.load_state_dict(vae_weights)
         self.stft.load_state_dict(stft_weights)
@@ -37,92 +46,84 @@ class Tango:
         self.stft.eval()
         self.model.eval()
         
-        self.scheduler = DDPMScheduler.from_pretrained(main_config["scheduler_name"], subfolder="scheduler")
+        self.scheduler = DDPMScheduler.from_pretrained(main_config["scheduler_name"], subfolder = "scheduler")
         
     def chunks(self, lst, n):
         """ 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=1, disable_progress=True):
-        """ Genrate audio for a single prompt string. """
+    def generate(self, prompt, steps = 100, guidance = 3, samples = 1, disable_progress = True):
+        """ Generate audio for a single prompt string. """
         with torch.no_grad():
-            latents = self.model.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
+            latents = self.model.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress = disable_progress)
             mel = self.vae.decode_first_stage(latents)
             wave = self.vae.decode_to_waveform(mel)
         return wave[0]
     
-    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. """
+    def generate_for_batch(self, prompts, steps = 200, guidance = 3, samples = 1, batch_size = 8, disable_progress = True):
+        """ Generate audio for a list of prompt strings. """
         outputs = []
         for k in tqdm(range(0, len(prompts), batch_size)):
-            batch = prompts[k: k+batch_size]
+            batch = prompts[k: k + batch_size]
             with torch.no_grad():
-                latents = self.model.inference(batch, self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
+                latents = self.model.inference(batch, self.scheduler, steps, guidance, samples, disable_progress = disable_progress)
                 mel = self.vae.decode_first_stage(latents)
                 wave = self.vae.decode_to_waveform(mel)
                 outputs += [item for item in wave]
         if samples == 1:
             return outputs
-        else:
-            return list(self.chunks(outputs, samples))
+        return list(self.chunks(outputs, samples))
 
 # Initialize TANGO
 
-tango = Tango(device="cpu")
-tango.vae.to("cpu")
-tango.stft.to("cpu")
-tango.model.to("cpu")
+tango = Tango(device = "cpu")
+tango.vae.to(device_type)
+tango.stft.to(device_type)
+tango.model.to(device_type)
     
-@spaces.GPU(duration=60)
+@spaces.GPU(duration = 60)
 def gradio_generate(prompt, 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)
+    wavio.write(output_filename, output_wave, rate = 16000, sampwidth = 2)
     
     return output_filename
 
-# description_text = """
-# <p><a href="https://huggingface.co/spaces/declare-lab/tango/blob/main/app.py?duplicate=true"> <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> For faster inference without waiting in queue, you may duplicate the space and upgrade to a GPU in the settings. <br/><br/>
-# Generate audio using TANGO by providing a text prompt.
-# <br/><br/>Limitations: TANGO is trained on the small AudioCaps dataset so it may not generate good audio \
-# samples related to concepts that it has not seen in training (e.g. singing). For the same reason, TANGO \
-# is not always able to finely control its generations over textual control prompts. For example, \
-# the generations from TANGO for prompts Chopping tomatoes on a wooden table and Chopping potatoes \
-# on a metal table are very similar. \
-# <br/><br/>We are currently training another version of TANGO on larger datasets to enhance its generalization, \
-# compositional and controllable generation ability.
-# <br/><br/>We recommend using a guidance scale of 3. The default number of steps is set to 100. More steps generally lead to better quality of generated audios but will take longer.
-# <br/><br/>
-# <h1> ChatGPT-enhanced audio generation</h1>
-# <br/>
-# As TANGO consists of an instruction-tuned LLM, it is able to process complex sound descriptions allowing us to provide more detailed instructions to improve the generation quality.
-# For example, ``A boat is moving on the sea'' vs ``The sound of the water lapping against the hull of the boat or splashing as you move through the waves''. The latter is obtained by prompting ChatGPT to explain the sound generated when a boat moves on the sea.
-# Using this ChatGPT-generated description of the sound, TANGO provides superior results.
-# <p/>
-# """
 description_text = """
-<p><a href="https://huggingface.co/spaces/declare-lab/tango2/blob/main/app.py?duplicate=true"> <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> For faster inference without waiting in queue, you may duplicate the space and upgrade to a GPU in the settings. <br/><br/>
-Generate audio using Tango2 by providing a text prompt. Tango2 was built from Tango and was trained on <a href="https://huggingface.co/datasets/declare-lab/audio-alpaca">Audio-alpaca</a>
-<br/><br/> This is the demo for Tango2 for text to audio generation: <a href="https://arxiv.org/abs/2404.09956">Read our paper.</a>
-<p/>
-"""
+        <p style="text-align: center;">
+        <b><big><big><big>Text-to-Audio</big></big></big></b>
+        <br/>Generates an audio file, freely, without account, without watermark, that you can download.
+        </p>
+        <br/>
+        <br/>
+        🚀 Powered by <i>Tango 2</i> AI.
+        <br/>
+        <ul>
+        <li>If you need to generate <b>music</b>, I recommend you to use <i>MusicGen</i>,</li>
+        </ul>
+        <br/>
+        🐌 Slow process... Your computer must <b><u>not</u></b> enter into standby mode.<br/>You can duplicate this space on a free account, it works on CPU.<br/>
+        <a href='https://huggingface.co/spaces/Fabrice-TIERCELIN/Text-to-Audio?duplicate=true'><img src='https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14'></a>
+        <br/>
+        ⚖️ You can use, modify and share the generated sounds but not for commercial uses.
+        """
 # Gradio input and output components
-input_text = gr.Textbox(lines=2, label="Prompt")
-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)
+input_text = gr.Textbox(label = "Prompt", value = "Snort of a horse", lines = 2, autofocus = True)
+output_audio = gr.Audio(label = "Generated Audio", type = "filepath")
+denoising_steps = gr.Slider(label = "Steps", minimum = 100, maximum = 200, value = 100, step = 1, interactive = True)
+guidance_scale = gr.Slider(label = "Guidance Scale", minimum = 1, maximum = 10, value = 3, step = 0.1, interactive = True)
 
 # Gradio interface
 gr_interface = gr.Interface(
-    fn=gradio_generate,
-    inputs=[input_text, denoising_steps, guidance_scale],
-    outputs=[output_audio],
-    title="Tango 2: Aligning Diffusion-based Text-to-Audio Generations through Direct Preference Optimization",
-    description=description_text,
-    allow_flagging=False,
-    examples=[
+    fn = gradio_generate,
+    inputs = [input_text, denoising_steps, guidance_scale],
+    outputs = [output_audio],
+    title = "",
+    description = description_text,
+    allow_flagging = False,
+    examples = [
         ["Quiet speech and then and airplane flying away"],
         ["A bicycle peddling on dirt and gravel followed by a man speaking then laughing"],
         ["Ducks quack and water splashes with some animal screeching in the background"],
@@ -147,7 +148,7 @@ gr_interface = gr.Interface(
         ["A helicopter is in flight"],
         ["A dog barking and a man talking and a racing car passes by"],
     ],
-    cache_examples="lazy", # Turn on to cache.
+    cache_examples = "lazy", # Turn on to cache.
 )
 
 # Launch Gradio app