declare-lab/tango2
92
1import spaces2import gradio as gr3import json4import torch5import wavio6from tqdm import tqdm7from huggingface_hub import snapshot_download8from models import AudioDiffusion, DDPMScheduler9from audioldm.audio.stft import TacotronSTFT10from audioldm.variational_autoencoder import AutoencoderKL11from pydub import AudioSegment12from gradio import Markdown13 14import torch15#from diffusers.models.autoencoder_kl import AutoencoderKL16from diffusers.models.unet_2d_condition import UNet2DConditionModel17from diffusers import DiffusionPipeline,AudioPipelineOutput18from transformers import CLIPTextModel, T5EncoderModel, AutoModel, T5Tokenizer, T5TokenizerFast19from typing import Union20from diffusers.utils.torch_utils import randn_tensor21from tqdm import tqdm22 23 24 25 26 27class Tango2Pipeline(DiffusionPipeline):28 29 30 def __init__(31 self,32 vae: AutoencoderKL,33 text_encoder: T5EncoderModel,34 tokenizer: Union[T5Tokenizer, T5TokenizerFast],35 unet: UNet2DConditionModel,36 scheduler: DDPMScheduler37 ):38 39 super().__init__()40 41 self.register_modules(vae=vae,42 text_encoder=text_encoder,43 tokenizer=tokenizer,44 unet=unet,45 scheduler=scheduler46 )47 48 49 def _encode_prompt(self, prompt):50 device = self.text_encoder.device51 52 batch = self.tokenizer(53 prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"54 )55 input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)56 57 58 encoder_hidden_states = self.text_encoder(59 input_ids=input_ids, attention_mask=attention_mask60 )[0]61 62 boolean_encoder_mask = (attention_mask == 1).to(device)63 64 return encoder_hidden_states, boolean_encoder_mask65 66 def _encode_text_classifier_free(self, prompt, num_samples_per_prompt):67 device = self.text_encoder.device68 batch = self.tokenizer(69 prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"70 )71 input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)72 73 with torch.no_grad():74 prompt_embeds = self.text_encoder(75 input_ids=input_ids, attention_mask=attention_mask76 )[0]77 78 prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)79 attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)80 81 # get unconditional embeddings for classifier free guidance82 uncond_tokens = [""] * len(prompt)83 84 max_length = prompt_embeds.shape[1]85 uncond_batch = self.tokenizer(86 uncond_tokens, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt",87 )88 uncond_input_ids = uncond_batch.input_ids.to(device)89 uncond_attention_mask = uncond_batch.attention_mask.to(device)90 91 with torch.no_grad():92 negative_prompt_embeds = self.text_encoder(93 input_ids=uncond_input_ids, attention_mask=uncond_attention_mask94 )[0]95 96 negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)97 uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)98 99 # For classifier free guidance, we need to do two forward passes.100 # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes101 prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])102 prompt_mask = torch.cat([uncond_attention_mask, attention_mask])103 boolean_prompt_mask = (prompt_mask == 1).to(device)104 105 return prompt_embeds, boolean_prompt_mask106 107 def prepare_latents(self, batch_size, inference_scheduler, num_channels_latents, dtype, device):108 shape = (batch_size, num_channels_latents, 256, 16)109 latents = randn_tensor(shape, generator=None, device=device, dtype=dtype)110 # scale the initial noise by the standard deviation required by the scheduler111 latents = latents * inference_scheduler.init_noise_sigma112 return latents113 114 @torch.no_grad()115 def inference(self, prompt, inference_scheduler, num_steps=20, guidance_scale=3, num_samples_per_prompt=1, 116 disable_progress=True):117 device = self.text_encoder.device118 classifier_free_guidance = guidance_scale > 1.0119 batch_size = len(prompt) * num_samples_per_prompt120 121 if classifier_free_guidance:122 prompt_embeds, boolean_prompt_mask = self._encode_text_classifier_free(prompt, num_samples_per_prompt)123 else:124 prompt_embeds, boolean_prompt_mask = self._encode_text(prompt)125 prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)126 boolean_prompt_mask = boolean_prompt_mask.repeat_interleave(num_samples_per_prompt, 0)127 128 inference_scheduler.set_timesteps(num_steps, device=device)129 timesteps = inference_scheduler.timesteps130 131 num_channels_latents = self.unet.config.in_channels132 latents = self.prepare_latents(batch_size, inference_scheduler, num_channels_latents, prompt_embeds.dtype, device)133 134 num_warmup_steps = len(timesteps) - num_steps * inference_scheduler.order135 progress_bar = tqdm(range(num_steps), disable=disable_progress)136 137 for i, t in enumerate(timesteps):138 # expand the latents if we are doing classifier free guidance139 latent_model_input = torch.cat([latents] * 2) if classifier_free_guidance else latents140 latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)141 142 noise_pred = self.unet(143 latent_model_input, t, encoder_hidden_states=prompt_embeds,144 encoder_attention_mask=boolean_prompt_mask145 ).sample146 147 # perform guidance148 if classifier_free_guidance:149 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)150 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)151 152 # compute the previous noisy sample x_t -> x_t-1153 latents = inference_scheduler.step(noise_pred, t, latents).prev_sample154 155 # call the callback, if provided156 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):157 progress_bar.update(1)158 159 return latents160 161 @torch.no_grad()162 def __call__(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):163 """ Genrate audio for a single prompt string. """164 with torch.no_grad():165 latents = self.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)166 mel = self.vae.decode_first_stage(latents)167 wave = self.vae.decode_to_waveform(mel)168 169 170 return AudioPipelineOutput(audios=wave)171 172 173# Automatic device detection174if torch.cuda.is_available():175 device_type = "cuda"176 device_selection = "cuda:0"177else:178 device_type = "cpu"179 device_selection = "cpu"180 181class Tango:182 def __init__(self, name="declare-lab/tango2", device=device_selection):183 184 path = snapshot_download(repo_id=name)185 186 vae_config = json.load(open("{}/vae_config.json".format(path)))187 stft_config = json.load(open("{}/stft_config.json".format(path)))188 main_config = json.load(open("{}/main_config.json".format(path)))189 190 self.vae = AutoencoderKL(**vae_config).to(device)191 self.stft = TacotronSTFT(**stft_config).to(device)192 self.model = AudioDiffusion(**main_config).to(device)193 194 vae_weights = torch.load("{}/pytorch_model_vae.bin".format(path), map_location=device)195 stft_weights = torch.load("{}/pytorch_model_stft.bin".format(path), map_location=device)196 main_weights = torch.load("{}/pytorch_model_main.bin".format(path), map_location=device)197 198 self.vae.load_state_dict(vae_weights)199 self.stft.load_state_dict(stft_weights)200 self.model.load_state_dict(main_weights)201 202 print ("Successfully loaded checkpoint from:", name)203 204 self.vae.eval()205 self.stft.eval()206 self.model.eval()207 208 self.scheduler = DDPMScheduler.from_pretrained(main_config["scheduler_name"], subfolder="scheduler")209 210 def chunks(self, lst, n):211 """ Yield successive n-sized chunks from a list. """212 for i in range(0, len(lst), n):213 yield lst[i:i + n]214 215 def generate(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):216 """ Genrate audio for a single prompt string. """217 with torch.no_grad():218 latents = self.model.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)219 mel = self.vae.decode_first_stage(latents)220 wave = self.vae.decode_to_waveform(mel)221 return wave[0]222 223 def generate_for_batch(self, prompts, steps=200, guidance=3, samples=1, batch_size=8, disable_progress=True):224 """ Genrate audio for a list of prompt strings. """225 outputs = []226 for k in tqdm(range(0, len(prompts), batch_size)):227 batch = prompts[k: k+batch_size]228 with torch.no_grad():229 latents = self.model.inference(batch, self.scheduler, steps, guidance, samples, disable_progress=disable_progress)230 mel = self.vae.decode_first_stage(latents)231 wave = self.vae.decode_to_waveform(mel)232 outputs += [item for item in wave]233 if samples == 1:234 return outputs235 else:236 return list(self.chunks(outputs, samples))237 238# Initialize TANGO239 240tango = Tango(device="cpu")241tango.vae.to(device_type)242tango.stft.to(device_type)243tango.model.to(device_type)244 245pipe = Tango2Pipeline(vae=tango.vae,246 text_encoder=tango.model.text_encoder,247 tokenizer=tango.model.tokenizer,248 unet=tango.model.unet,249 scheduler=tango.scheduler250 )251 252 253@spaces.GPU(duration=60)254def gradio_generate(prompt, output_format, steps, guidance):255 output_wave = pipe(prompt,steps,guidance) ## Using pipeliine automatically uses flash attention for torch2.0 above256 #output_wave = tango.generate(prompt, steps, guidance)257 # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"258 output_wave = output_wave.audios[0]259 output_filename = "temp.wav"260 wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)261 262 if (output_format == "mp3"):263 AudioSegment.from_wav("temp.wav").export("temp.mp3", format = "mp3")264 output_filename = "temp.mp3"265 266 return output_filename267 268# description_text = """269# <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/>270# Generate audio using TANGO by providing a text prompt.271# <br/><br/>Limitations: TANGO is trained on the small AudioCaps dataset so it may not generate good audio \272# samples related to concepts that it has not seen in training (e.g. singing). For the same reason, TANGO \273# is not always able to finely control its generations over textual control prompts. For example, \274# the generations from TANGO for prompts Chopping tomatoes on a wooden table and Chopping potatoes \275# on a metal table are very similar. \276# <br/><br/>We are currently training another version of TANGO on larger datasets to enhance its generalization, \277# compositional and controllable generation ability.278# <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.279# <br/><br/>280# <h1> ChatGPT-enhanced audio generation</h1>281# <br/>282# 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.283# 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.284# Using this ChatGPT-generated description of the sound, TANGO provides superior results.285# <p/>286# """287description_text = """288<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/>289Generate 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>290<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>291<p/>292"""293# Gradio input and output components294input_text = gr.Textbox(lines=2, label="Prompt")295output_format = gr.Radio(label = "Output format", info = "The file you can dowload", choices = ["mp3", "wav"], value = "wav")296output_audio = gr.Audio(label="Generated Audio", type="filepath")297denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)298guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)299 300# Gradio interface301gr_interface = gr.Interface(302 fn=gradio_generate,303 inputs=[input_text, output_format, denoising_steps, guidance_scale],304 outputs=[output_audio],305 title="Tango 2: Aligning Diffusion-based Text-to-Audio Generations through Direct Preference Optimization",306 description=description_text,307 allow_flagging=False,308 examples=[309 ["Quiet speech and then and airplane flying away"],310 ["A bicycle peddling on dirt and gravel followed by a man speaking then laughing"],311 ["Ducks quack and water splashes with some animal screeching in the background"],312 ["Describe the sound of the ocean"],313 ["A woman and a baby are having a conversation"],314 ["A man speaks followed by a popping noise and laughter"],315 ["A cup is filled from a faucet"],316 ["An audience cheering and clapping"],317 ["Rolling thunder with lightning strikes"],318 ["A dog barking and a cat mewing and a racing car passes by"],319 ["Gentle water stream, birds chirping and sudden gun shot"],320 ["A man talking followed by a goat baaing then a metal gate sliding shut as ducks quack and wind blows into a microphone."],321 ["A dog barking"],322 ["A cat meowing"],323 ["Wooden table tapping sound while water pouring"],324 ["Applause from a crowd with distant clicking and a man speaking over a loudspeaker"],325 ["two gunshots followed by birds flying away while chirping"],326 ["Whistling with birds chirping"],327 ["A person snoring"],328 ["Motor vehicles are driving with loud engines and a person whistles"],329 ["People cheering in a stadium while thunder and lightning strikes"],330 ["A helicopter is in flight"],331 ["A dog barking and a man talking and a racing car passes by"],332 ],333 cache_examples="lazy", # Turn on to cache.334)335 336# Launch Gradio app337gr_interface.queue(10).launch()