Spaces:
Running
on
A10G
Running
on
A10G
| import os | |
| import torch | |
| import numpy as np | |
| from tqdm import tqdm | |
| 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 | |
| import os | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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, | |
| ) | |
| 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 | |
| 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 | |
| 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 |