fantaxy/tango2
1
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 269# Gradio input and output components270input_text = gr.Textbox(lines=2, label="Prompt")271output_format = gr.Radio(label = "Output format", info = "The file you can dowload", choices = ["mp3", "wav"], value = "wav")272output_audio = gr.Audio(label="Generated Audio", type="filepath")273denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)274guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)275 276# Gradio interface277gr_interface = gr.Interface(theme="Nymbo/Nymbo_Theme",278 fn=gradio_generate,279 inputs=[input_text, output_format, denoising_steps, guidance_scale],280 outputs=[output_audio],281 title="T2: Text to SoundFX",282 283 allow_flagging=False,284 examples=[285 ["Quiet speech and then and airplane flying away"],286 ["A bicycle peddling on dirt and gravel followed by a man speaking then laughing"],287 ["Ducks quack and water splashes with some animal screeching in the background"],288 ["Describe the sound of the ocean"],289 ["A woman and a baby are having a conversation"],290 ["A man speaks followed by a popping noise and laughter"],291 ["A cup is filled from a faucet"],292 ["An audience cheering and clapping"],293 ["Rolling thunder with lightning strikes"],294 ["A dog barking and a cat mewing and a racing car passes by"],295 ["Gentle water stream, birds chirping and sudden gun shot"],296 ["A man talking followed by a goat baaing then a metal gate sliding shut as ducks quack and wind blows into a microphone."],297 ["A dog barking"],298 ["A cat meowing"],299 ["Wooden table tapping sound while water pouring"],300 ["Applause from a crowd with distant clicking and a man speaking over a loudspeaker"],301 ["two gunshots followed by birds flying away while chirping"],302 ["Whistling with birds chirping"],303 ["A person snoring"],304 ["Motor vehicles are driving with loud engines and a person whistles"],305 ["People cheering in a stadium while thunder and lightning strikes"],306 ["A helicopter is in flight"],307 ["A dog barking and a man talking and a racing car passes by"],308 ],309 cache_examples="lazy", # Turn on to cache.310)311 312# Launch Gradio app313gr_interface.queue(10).launch()