import os import torch import numpy as np from tqdm import tqdm import os try: from audioldm.utils import default, instantiate_from_config, save_wave from audioldm.latent_diffusion.ddpm import DDPM from audioldm.variational_autoencoder.distributions import DiagonalGaussianDistribution from audioldm.latent_diffusion.util import noise_like from audioldm.latent_diffusion.ddim import DDIMSampler except ModuleNotFoundError: from .utils import default, instantiate_from_config, save_wave from .latent_diffusion.ddpm import DDPM from .variational_autoencoder.distributions import DiagonalGaussianDistribution from .latent_diffusion.util import noise_like from .latent_diffusion.ddim import DDIMSampler def disabled_train(self, mode=True): """Overwrite model.train with this function to make sure train/eval mode does not change anymore.""" return self class LatentDiffusion(DDPM): """main class""" def __init__( self, device="cuda", first_stage_config=None, cond_stage_config=None, num_timesteps_cond=None, cond_stage_key="image", cond_stage_trainable=False, concat_mode=True, cond_stage_forward=None, conditioning_key=None, scale_factor=1.0, scale_by_std=False, base_learning_rate=None, *args, **kwargs, ): self.device = device self.learning_rate = base_learning_rate self.num_timesteps_cond = default(num_timesteps_cond, 1) self.scale_by_std = scale_by_std assert self.num_timesteps_cond <= kwargs["timesteps"] # for backwards compatibility after implementation of DiffusionWrapper if conditioning_key is None: conditioning_key = "concat" if concat_mode else "crossattn" if cond_stage_config == "__is_unconditional__": conditioning_key = None ckpt_path = kwargs.pop("ckpt_path", None) ignore_keys = kwargs.pop("ignore_keys", []) super().__init__(conditioning_key=conditioning_key, *args, **kwargs) self.concat_mode = concat_mode self.cond_stage_trainable = cond_stage_trainable self.cond_stage_key = cond_stage_key self.cond_stage_key_orig = cond_stage_key try: self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 except: self.num_downs = 0 if not scale_by_std: self.scale_factor = scale_factor else: self.register_buffer("scale_factor", torch.tensor(scale_factor)) self.instantiate_first_stage(first_stage_config) self.instantiate_cond_stage(cond_stage_config) self.cond_stage_forward = cond_stage_forward self.clip_denoised = False def make_cond_schedule( self, ): self.cond_ids = torch.full( size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long, ) ids = torch.round( torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond) ).long() self.cond_ids[: self.num_timesteps_cond] = ids def register_schedule( self, given_betas=None, beta_schedule="linear", timesteps=1000, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3, ): super().register_schedule( given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s ) self.shorten_cond_schedule = self.num_timesteps_cond > 1 if self.shorten_cond_schedule: self.make_cond_schedule() def instantiate_first_stage(self, config): model = instantiate_from_config(config) self.first_stage_model = model.eval() self.first_stage_model.train = disabled_train for param in self.first_stage_model.parameters(): param.requires_grad = False def instantiate_cond_stage(self, config): if not self.cond_stage_trainable: if config == "__is_first_stage__": print("Using first stage also as cond stage.") self.cond_stage_model = self.first_stage_model elif config == "__is_unconditional__": print(f"Training {self.__class__.__name__} as an unconditional model.") self.cond_stage_model = None # self.be_unconditional = True else: model = instantiate_from_config(config) self.cond_stage_model = model.eval() self.cond_stage_model.train = disabled_train for param in self.cond_stage_model.parameters(): param.requires_grad = False else: assert config != "__is_first_stage__" assert config != "__is_unconditional__" model = instantiate_from_config(config) self.cond_stage_model = model self.cond_stage_model = self.cond_stage_model.to(self.device) def get_first_stage_encoding(self, encoder_posterior): if isinstance(encoder_posterior, DiagonalGaussianDistribution): z = encoder_posterior.sample() elif isinstance(encoder_posterior, torch.Tensor): z = encoder_posterior else: raise NotImplementedError( f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented" ) return self.scale_factor * z def get_learned_conditioning(self, c): if self.cond_stage_forward is None: if hasattr(self.cond_stage_model, "encode") and callable( self.cond_stage_model.encode ): c = self.cond_stage_model.encode(c) if isinstance(c, DiagonalGaussianDistribution): c = c.mode() else: # Text input is list if type(c) == list and len(c) == 1: c = self.cond_stage_model([c[0], c[0]]) c = c[0:1] else: c = self.cond_stage_model(c) else: assert hasattr(self.cond_stage_model, self.cond_stage_forward) c = getattr(self.cond_stage_model, self.cond_stage_forward)(c) return c @torch.no_grad() def get_input( self, batch, k, return_first_stage_encode=True, return_first_stage_outputs=False, force_c_encode=False, cond_key=None, return_original_cond=False, bs=None, ): x = super().get_input(batch, k) if bs is not None: x = x[:bs] x = x.to(self.device) if return_first_stage_encode: encoder_posterior = self.encode_first_stage(x) z = self.get_first_stage_encoding(encoder_posterior).detach() else: z = None if self.model.conditioning_key is not None: if cond_key is None: cond_key = self.cond_stage_key if cond_key != self.first_stage_key: if cond_key in ["caption", "coordinates_bbox"]: xc = batch[cond_key] elif cond_key == "class_label": xc = batch else: # [bs, 1, 527] xc = super().get_input(batch, cond_key) if type(xc) == torch.Tensor: xc = xc.to(self.device) else: xc = x if not self.cond_stage_trainable or force_c_encode: if isinstance(xc, dict) or isinstance(xc, list): c = self.get_learned_conditioning(xc) else: c = self.get_learned_conditioning(xc.to(self.device)) else: c = xc if bs is not None: c = c[:bs] else: c = None xc = None if self.use_positional_encodings: pos_x, pos_y = self.compute_latent_shifts(batch) c = {"pos_x": pos_x, "pos_y": pos_y} out = [z, c] if return_first_stage_outputs: xrec = self.decode_first_stage(z) out.extend([x, xrec]) if return_original_cond: out.append(xc) return out @torch.no_grad() def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): if predict_cids: if z.dim() == 4: z = torch.argmax(z.exp(), dim=1).long() z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) z = rearrange(z, "b h w c -> b c h w").contiguous() z = 1.0 / self.scale_factor * z return self.first_stage_model.decode(z) def mel_spectrogram_to_waveform(self, mel): # Mel: [bs, 1, t-steps, fbins] if len(mel.size()) == 4: mel = mel.squeeze(1) mel = mel.permute(0, 2, 1) waveform = self.first_stage_model.vocoder(mel) waveform = waveform.cpu().detach().numpy() return waveform @torch.no_grad() def encode_first_stage(self, x): return self.first_stage_model.encode(x) def apply_model(self, x_noisy, t, cond, return_ids=False): if isinstance(cond, dict): # hybrid case, cond is exptected to be a dict pass else: if not isinstance(cond, list): cond = [cond] if self.model.conditioning_key == "concat": key = "c_concat" elif self.model.conditioning_key == "crossattn": key = "c_crossattn" else: key = "c_film" cond = {key: cond} x_recon = self.model(x_noisy, t, **cond) if isinstance(x_recon, tuple) and not return_ids: return x_recon[0] else: return x_recon def p_mean_variance( self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False, return_x0=False, score_corrector=None, corrector_kwargs=None, ): t_in = t model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids) if score_corrector is not None: assert self.parameterization == "eps" model_out = score_corrector.modify_score( self, model_out, x, t, c, **corrector_kwargs ) if return_codebook_ids: model_out, logits = model_out if self.parameterization == "eps": x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) elif self.parameterization == "x0": x_recon = model_out else: raise NotImplementedError() if clip_denoised: x_recon.clamp_(-1.0, 1.0) if quantize_denoised: x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon) model_mean, posterior_variance, posterior_log_variance = self.q_posterior( x_start=x_recon, x_t=x, t=t ) if return_codebook_ids: return model_mean, posterior_variance, posterior_log_variance, logits elif return_x0: return model_mean, posterior_variance, posterior_log_variance, x_recon else: return model_mean, posterior_variance, posterior_log_variance @torch.no_grad() def p_sample( self, x, c, t, clip_denoised=False, repeat_noise=False, return_codebook_ids=False, quantize_denoised=False, return_x0=False, temperature=1.0, noise_dropout=0.0, score_corrector=None, corrector_kwargs=None, ): b, *_, device = *x.shape, x.device outputs = self.p_mean_variance( x=x, c=c, t=t, clip_denoised=clip_denoised, return_codebook_ids=return_codebook_ids, quantize_denoised=quantize_denoised, return_x0=return_x0, score_corrector=score_corrector, corrector_kwargs=corrector_kwargs, ) if return_codebook_ids: raise DeprecationWarning("Support dropped.") model_mean, _, model_log_variance, logits = outputs elif return_x0: model_mean, _, model_log_variance, x0 = outputs else: model_mean, _, model_log_variance = outputs noise = noise_like(x.shape, device, repeat_noise) * temperature if noise_dropout > 0.0: noise = torch.nn.functional.dropout(noise, p=noise_dropout) # no noise when t == 0 nonzero_mask = ( (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))).contiguous() ) if return_codebook_ids: return model_mean + nonzero_mask * ( 0.5 * model_log_variance ).exp() * noise, logits.argmax(dim=1) if return_x0: return ( model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0, ) else: return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise @torch.no_grad() def progressive_denoising( self, cond, shape, verbose=True, callback=None, quantize_denoised=False, img_callback=None, mask=None, x0=None, temperature=1.0, noise_dropout=0.0, score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None, log_every_t=None, ): if not log_every_t: log_every_t = self.log_every_t timesteps = self.num_timesteps if batch_size is not None: b = batch_size if batch_size is not None else shape[0] shape = [batch_size] + list(shape) else: b = batch_size = shape[0] if x_T is None: img = torch.randn(shape, device=self.device) else: img = x_T intermediates = [] if cond is not None: if isinstance(cond, dict): cond = { key: cond[key][:batch_size] if not isinstance(cond[key], list) else list(map(lambda x: x[:batch_size], cond[key])) for key in cond } else: cond = ( [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] ) if start_T is not None: timesteps = min(timesteps, start_T) iterator = ( tqdm( reversed(range(0, timesteps)), desc="Progressive Generation", total=timesteps, ) if verbose else reversed(range(0, timesteps)) ) if type(temperature) == float: temperature = [temperature] * timesteps for i in iterator: ts = torch.full((b,), i, device=self.device, dtype=torch.long) if self.shorten_cond_schedule: assert self.model.conditioning_key != "hybrid" tc = self.cond_ids[ts].to(cond.device) cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) img, x0_partial = self.p_sample( img, cond, ts, clip_denoised=self.clip_denoised, quantize_denoised=quantize_denoised, return_x0=True, temperature=temperature[i], noise_dropout=noise_dropout, score_corrector=score_corrector, corrector_kwargs=corrector_kwargs, ) if mask is not None: assert x0 is not None img_orig = self.q_sample(x0, ts) img = img_orig * mask + (1.0 - mask) * img if i % log_every_t == 0 or i == timesteps - 1: intermediates.append(x0_partial) if callback: callback(i) if img_callback: img_callback(img, i) return img, intermediates @torch.no_grad() def p_sample_loop( self, cond, shape, return_intermediates=False, x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, mask=None, x0=None, img_callback=None, start_T=None, log_every_t=None, ): if not log_every_t: log_every_t = self.log_every_t device = self.betas.device b = shape[0] if x_T is None: img = torch.randn(shape, device=device) else: img = x_T intermediates = [img] if timesteps is None: timesteps = self.num_timesteps if start_T is not None: timesteps = min(timesteps, start_T) iterator = ( tqdm(reversed(range(0, timesteps)), desc="Sampling t", total=timesteps) if verbose else reversed(range(0, timesteps)) ) if mask is not None: assert x0 is not None assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match for i in iterator: ts = torch.full((b,), i, device=device, dtype=torch.long) if self.shorten_cond_schedule: assert self.model.conditioning_key != "hybrid" tc = self.cond_ids[ts].to(cond.device) cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) img = self.p_sample( img, cond, ts, clip_denoised=self.clip_denoised, quantize_denoised=quantize_denoised, ) if mask is not None: img_orig = self.q_sample(x0, ts) img = img_orig * mask + (1.0 - mask) * img if i % log_every_t == 0 or i == timesteps - 1: intermediates.append(img) if callback: callback(i) if img_callback: img_callback(img, i) if return_intermediates: return img, intermediates return img @torch.no_grad() def sample( self, cond, batch_size=16, return_intermediates=False, x_T=None, verbose=True, timesteps=None, quantize_denoised=False, mask=None, x0=None, shape=None, **kwargs, ): if shape is None: shape = (batch_size, self.channels, self.latent_t_size, self.latent_f_size) if cond is not None: if isinstance(cond, dict): cond = { key: cond[key][:batch_size] if not isinstance(cond[key], list) else list(map(lambda x: x[:batch_size], cond[key])) for key in cond } else: cond = ( [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] ) return self.p_sample_loop( cond, shape, return_intermediates=return_intermediates, x_T=x_T, verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised, mask=mask, x0=x0, **kwargs, ) @torch.no_grad() def sample_log( self, cond, batch_size, ddim, ddim_steps, unconditional_guidance_scale=1.0, unconditional_conditioning=None, use_plms=False, mask=None, **kwargs, ): if mask is not None: shape = (self.channels, mask.size()[-2], mask.size()[-1]) else: shape = (self.channels, self.latent_t_size, self.latent_f_size) intermediate = None if ddim and not use_plms: # print("Use ddim sampler") ddim_sampler = DDIMSampler(self) samples, intermediates = ddim_sampler.sample( ddim_steps, batch_size, shape, cond, verbose=False, unconditional_guidance_scale=unconditional_guidance_scale, unconditional_conditioning=unconditional_conditioning, mask=mask, **kwargs, ) else: # print("Use DDPM sampler") samples, intermediates = self.sample( cond=cond, batch_size=batch_size, return_intermediates=True, unconditional_guidance_scale=unconditional_guidance_scale, mask=mask, unconditional_conditioning=unconditional_conditioning, **kwargs, ) return samples, intermediate @torch.no_grad() def generate_sample( self, batchs, ddim_steps=200, ddim_eta=1.0, x_T=None, n_candidate_gen_per_text=1, unconditional_guidance_scale=1.0, unconditional_conditioning=None, name="waveform", use_plms=False, save=False, **kwargs, ): # Generate n_candidate_gen_per_text times and select the best # Batch: audio, text, fnames assert x_T is None try: batchs = iter(batchs) except TypeError: raise ValueError("The first input argument should be an iterable object") if use_plms: assert ddim_steps is not None use_ddim = ddim_steps is not None # waveform_save_path = os.path.join(self.get_log_dir(), name) # os.makedirs(waveform_save_path, exist_ok=True) # print("Waveform save path: ", waveform_save_path) with self.ema_scope("Generate"): for batch in batchs: z, c = self.get_input( batch, self.first_stage_key, cond_key=self.cond_stage_key, return_first_stage_outputs=False, force_c_encode=True, return_original_cond=False, bs=None, ) text = super().get_input(batch, "text") # Generate multiple samples batch_size = z.shape[0] * n_candidate_gen_per_text c = torch.cat([c] * n_candidate_gen_per_text, dim=0) text = text * n_candidate_gen_per_text if unconditional_guidance_scale != 1.0: unconditional_conditioning = ( self.cond_stage_model.get_unconditional_condition(batch_size) ) samples, _ = self.sample_log( cond=c, batch_size=batch_size, x_T=x_T, ddim=use_ddim, ddim_steps=ddim_steps, eta=ddim_eta, unconditional_guidance_scale=unconditional_guidance_scale, unconditional_conditioning=unconditional_conditioning, use_plms=use_plms, ) if(torch.max(torch.abs(samples)) > 1e2): samples = torch.clip(samples, min=-10, max=10) mel = self.decode_first_stage(samples) waveform = self.mel_spectrogram_to_waveform(mel) if waveform.shape[0] > 1: similarity = self.cond_stage_model.cos_similarity( torch.FloatTensor(waveform).squeeze(1), text ) best_index = [] for i in range(z.shape[0]): candidates = similarity[i :: z.shape[0]] max_index = torch.argmax(candidates).item() best_index.append(i + max_index * z.shape[0]) waveform = waveform[best_index] # print("Similarity between generated audio and text", similarity) # print("Choose the following indexes:", best_index) return waveform @torch.no_grad() def generate_sample_masked( self, batchs, ddim_steps=200, ddim_eta=1.0, x_T=None, n_candidate_gen_per_text=1, unconditional_guidance_scale=1.0, unconditional_conditioning=None, name="waveform", use_plms=False, time_mask_ratio_start_and_end=(0.25, 0.75), freq_mask_ratio_start_and_end=(0.75, 1.0), save=False, **kwargs, ): # Generate n_candidate_gen_per_text times and select the best # Batch: audio, text, fnames assert x_T is None try: batchs = iter(batchs) except TypeError: raise ValueError("The first input argument should be an iterable object") if use_plms: assert ddim_steps is not None use_ddim = ddim_steps is not None # waveform_save_path = os.path.join(self.get_log_dir(), name) # os.makedirs(waveform_save_path, exist_ok=True) # print("Waveform save path: ", waveform_save_path) with self.ema_scope("Generate"): for batch in batchs: z, c = self.get_input( batch, self.first_stage_key, cond_key=self.cond_stage_key, return_first_stage_outputs=False, force_c_encode=True, return_original_cond=False, bs=None, ) text = super().get_input(batch, "text") # Generate multiple samples batch_size = z.shape[0] * n_candidate_gen_per_text _, h, w = z.shape[0], z.shape[2], z.shape[3] mask = torch.ones(batch_size, h, w).to(self.device) mask[:, int(h * time_mask_ratio_start_and_end[0]) : int(h * time_mask_ratio_start_and_end[1]), :] = 0 mask[:, :, int(w * freq_mask_ratio_start_and_end[0]) : int(w * freq_mask_ratio_start_and_end[1])] = 0 mask = mask[:, None, ...] c = torch.cat([c] * n_candidate_gen_per_text, dim=0) text = text * n_candidate_gen_per_text if unconditional_guidance_scale != 1.0: unconditional_conditioning = ( self.cond_stage_model.get_unconditional_condition(batch_size) ) samples, _ = self.sample_log( cond=c, batch_size=batch_size, x_T=x_T, ddim=use_ddim, ddim_steps=ddim_steps, eta=ddim_eta, unconditional_guidance_scale=unconditional_guidance_scale, unconditional_conditioning=unconditional_conditioning, use_plms=use_plms, mask=mask, x0=torch.cat([z] * n_candidate_gen_per_text) ) mel = self.decode_first_stage(samples) waveform = self.mel_spectrogram_to_waveform(mel) if waveform.shape[0] > 1: similarity = self.cond_stage_model.cos_similarity( torch.FloatTensor(waveform).squeeze(1), text ) best_index = [] for i in range(z.shape[0]): candidates = similarity[i :: z.shape[0]] max_index = torch.argmax(candidates).item() best_index.append(i + max_index * z.shape[0]) waveform = waveform[best_index] # print("Similarity between generated audio and text", similarity) # print("Choose the following indexes:", best_index) return waveform