fantaxy/tango2
1
1import spaces2import yaml3import random4import inspect5import numpy as np6from tqdm import tqdm7 8import torch9import torch.nn as nn10import torch.nn.functional as F11 12from einops import repeat13from tools.torch_tools import wav_to_fbank14 15from audioldm.audio.stft import TacotronSTFT16from audioldm.variational_autoencoder import AutoencoderKL17from audioldm.utils import default_audioldm_config, get_metadata18 19from transformers import CLIPTokenizer, AutoTokenizer20from transformers import CLIPTextModel, T5EncoderModel, AutoModel21 22import sys23sys.path.insert(0, "diffusers/src")24 25import diffusers26from diffusers.utils import randn_tensor27from diffusers import DDPMScheduler, UNet2DConditionModel28from diffusers import AutoencoderKL as DiffuserAutoencoderKL29 30 31def build_pretrained_models(name):32 checkpoint = torch.load(get_metadata()[name]["path"], map_location="cpu")33 scale_factor = checkpoint["state_dict"]["scale_factor"].item()34 35 vae_state_dict = {k[18:]: v for k, v in checkpoint["state_dict"].items() if "first_stage_model." in k}36 37 config = default_audioldm_config(name)38 vae_config = config["model"]["params"]["first_stage_config"]["params"]39 vae_config["scale_factor"] = scale_factor40 41 vae = AutoencoderKL(**vae_config)42 vae.load_state_dict(vae_state_dict)43 44 fn_STFT = TacotronSTFT(45 config["preprocessing"]["stft"]["filter_length"],46 config["preprocessing"]["stft"]["hop_length"],47 config["preprocessing"]["stft"]["win_length"],48 config["preprocessing"]["mel"]["n_mel_channels"],49 config["preprocessing"]["audio"]["sampling_rate"],50 config["preprocessing"]["mel"]["mel_fmin"],51 config["preprocessing"]["mel"]["mel_fmax"],52 )53 54 vae.eval()55 fn_STFT.eval()56 return vae, fn_STFT57 58 59class AudioDiffusion(nn.Module):60 def __init__(61 self,62 text_encoder_name,63 scheduler_name,64 unet_model_name=None,65 unet_model_config_path=None,66 snr_gamma=None,67 freeze_text_encoder=True,68 uncondition=False,69 70 ):71 super().__init__()72 73 assert unet_model_name is not None or unet_model_config_path is not None, "Either UNet pretrain model name or a config file path is required"74 75 self.text_encoder_name = text_encoder_name76 self.scheduler_name = scheduler_name77 self.unet_model_name = unet_model_name78 self.unet_model_config_path = unet_model_config_path79 self.snr_gamma = snr_gamma80 self.freeze_text_encoder = freeze_text_encoder81 self.uncondition = uncondition82 83 # https://huggingface.co/docs/diffusers/v0.14.0/en/api/schedulers/overview84 self.noise_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")85 self.inference_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")86 87 if unet_model_config_path:88 unet_config = UNet2DConditionModel.load_config(unet_model_config_path)89 self.unet = UNet2DConditionModel.from_config(unet_config, subfolder="unet")90 self.set_from = "random"91 print("UNet initialized randomly.")92 else:93 self.unet = UNet2DConditionModel.from_pretrained(unet_model_name, subfolder="unet")94 self.set_from = "pre-trained"95 self.group_in = nn.Sequential(nn.Linear(8, 512), nn.Linear(512, 4))96 self.group_out = nn.Sequential(nn.Linear(4, 512), nn.Linear(512, 8))97 print("UNet initialized from stable diffusion checkpoint.")98 99 if "stable-diffusion" in self.text_encoder_name:100 self.tokenizer = CLIPTokenizer.from_pretrained(self.text_encoder_name, subfolder="tokenizer")101 self.text_encoder = CLIPTextModel.from_pretrained(self.text_encoder_name, subfolder="text_encoder")102 elif "t5" in self.text_encoder_name:103 self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)104 self.text_encoder = T5EncoderModel.from_pretrained(self.text_encoder_name)105 else:106 self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)107 self.text_encoder = AutoModel.from_pretrained(self.text_encoder_name)108 109 def compute_snr(self, timesteps):110 """111 Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849112 """113 alphas_cumprod = self.noise_scheduler.alphas_cumprod114 sqrt_alphas_cumprod = alphas_cumprod**0.5115 sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5116 117 # Expand the tensors.118 # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026119 sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()120 while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):121 sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]122 alpha = sqrt_alphas_cumprod.expand(timesteps.shape)123 124 sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()125 while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):126 sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]127 sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)128 129 # Compute SNR.130 snr = (alpha / sigma) ** 2131 return snr132 133 def encode_text(self, prompt):134 device = self.text_encoder.device135 batch = self.tokenizer(136 prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"137 )138 input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)139 140 if self.freeze_text_encoder:141 with torch.no_grad():142 encoder_hidden_states = self.text_encoder(143 input_ids=input_ids, attention_mask=attention_mask144 )[0]145 else:146 encoder_hidden_states = self.text_encoder(147 input_ids=input_ids, attention_mask=attention_mask148 )[0]149 150 boolean_encoder_mask = (attention_mask == 1).to(device)151 return encoder_hidden_states, boolean_encoder_mask152 153 def forward(self, latents, prompt):154 device = self.text_encoder.device155 num_train_timesteps = self.noise_scheduler.num_train_timesteps156 self.noise_scheduler.set_timesteps(num_train_timesteps, device=device)157 158 encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt)159 160 if self.uncondition:161 mask_indices = [k for k in range(len(prompt)) if random.random() < 0.1]162 if len(mask_indices) > 0:163 encoder_hidden_states[mask_indices] = 0164 165 bsz = latents.shape[0]166 # Sample a random timestep for each instance167 timesteps = torch.randint(0, self.noise_scheduler.num_train_timesteps, (bsz,), device=device)168 timesteps = timesteps.long()169 170 noise = torch.randn_like(latents)171 noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)172 173 # Get the target for loss depending on the prediction type174 if self.noise_scheduler.config.prediction_type == "epsilon":175 target = noise176 elif self.noise_scheduler.config.prediction_type == "v_prediction":177 target = self.noise_scheduler.get_velocity(latents, noise, timesteps)178 else:179 raise ValueError(f"Unknown prediction type {self.noise_scheduler.config.prediction_type}")180 181 if self.set_from == "random":182 model_pred = self.unet(183 noisy_latents, timesteps, encoder_hidden_states, 184 encoder_attention_mask=boolean_encoder_mask185 ).sample186 187 elif self.set_from == "pre-trained":188 compressed_latents = self.group_in(noisy_latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()189 model_pred = self.unet(190 compressed_latents, timesteps, encoder_hidden_states, 191 encoder_attention_mask=boolean_encoder_mask192 ).sample193 model_pred = self.group_out(model_pred.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()194 195 if self.snr_gamma is None:196 loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")197 else:198 # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.199 # Adaptef from huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py200 snr = self.compute_snr(timesteps)201 mse_loss_weights = (202 torch.stack([snr, self.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr203 )204 loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")205 loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights206 loss = loss.mean()207 208 return loss209 210 @torch.no_grad()211 def inference(self, prompt, inference_scheduler, num_steps=20, guidance_scale=3, num_samples_per_prompt=1, 212 disable_progress=True):213 device = self.text_encoder.device214 classifier_free_guidance = guidance_scale > 1.0215 batch_size = len(prompt) * num_samples_per_prompt216 217 if classifier_free_guidance:218 prompt_embeds, boolean_prompt_mask = self.encode_text_classifier_free(prompt, num_samples_per_prompt)219 else:220 prompt_embeds, boolean_prompt_mask = self.encode_text(prompt)221 prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)222 boolean_prompt_mask = boolean_prompt_mask.repeat_interleave(num_samples_per_prompt, 0)223 224 inference_scheduler.set_timesteps(num_steps, device=device)225 timesteps = inference_scheduler.timesteps226 227 num_channels_latents = self.unet.in_channels228 latents = self.prepare_latents(batch_size, inference_scheduler, num_channels_latents, prompt_embeds.dtype, device)229 230 num_warmup_steps = len(timesteps) - num_steps * inference_scheduler.order231 progress_bar = tqdm(range(num_steps), disable=disable_progress)232 233 for i, t in enumerate(timesteps):234 # expand the latents if we are doing classifier free guidance235 latent_model_input = torch.cat([latents] * 2) if classifier_free_guidance else latents236 latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)237 238 noise_pred = self.unet(239 latent_model_input, t, encoder_hidden_states=prompt_embeds,240 encoder_attention_mask=boolean_prompt_mask241 ).sample242 243 # perform guidance244 if classifier_free_guidance:245 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)246 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)247 248 # compute the previous noisy sample x_t -> x_t-1249 latents = inference_scheduler.step(noise_pred, t, latents).prev_sample250 251 # call the callback, if provided252 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):253 progress_bar.update(1)254 255 if self.set_from == "pre-trained":256 latents = self.group_out(latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()257 return latents258 259 def prepare_latents(self, batch_size, inference_scheduler, num_channels_latents, dtype, device):260 shape = (batch_size, num_channels_latents, 256, 16)261 latents = randn_tensor(shape, generator=None, device=device, dtype=dtype)262 # scale the initial noise by the standard deviation required by the scheduler263 latents = latents * inference_scheduler.init_noise_sigma264 return latents265 266 def encode_text_classifier_free(self, prompt, num_samples_per_prompt):267 device = self.text_encoder.device268 batch = self.tokenizer(269 prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"270 )271 input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)272 273 with torch.no_grad():274 prompt_embeds = self.text_encoder(275 input_ids=input_ids, attention_mask=attention_mask276 )[0]277 278 prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)279 attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)280 281 # get unconditional embeddings for classifier free guidance282 uncond_tokens = [""] * len(prompt)283 284 max_length = prompt_embeds.shape[1]285 uncond_batch = self.tokenizer(286 uncond_tokens, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt",287 )288 uncond_input_ids = uncond_batch.input_ids.to(device)289 uncond_attention_mask = uncond_batch.attention_mask.to(device)290 291 with torch.no_grad():292 negative_prompt_embeds = self.text_encoder(293 input_ids=uncond_input_ids, attention_mask=uncond_attention_mask294 )[0]295 296 negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)297 uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)298 299 # For classifier free guidance, we need to do two forward passes.300 # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes301 prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])302 prompt_mask = torch.cat([uncond_attention_mask, attention_mask])303 boolean_prompt_mask = (prompt_mask == 1).to(device)304 305 return prompt_embeds, boolean_prompt_mask