import os import random import pandas as pd import torch import librosa import numpy as np import soundfile as sf from tqdm import tqdm from .utils import scale_shift_re def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): """ Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 """ std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) # rescale the results from guidance (fixes overexposure) noise_pred_rescaled = noise_cfg * (std_text / std_cfg) # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg return noise_cfg @torch.no_grad() def inference(autoencoder, unet, gt, gt_mask, tokenizer, text_encoder, params, noise_scheduler, text_raw, neg_text=None, audio_frames=500, guidance_scale=3, guidance_rescale=0.0, ddim_steps=50, eta=1, random_seed=2024, device='cuda', ): if neg_text is None: neg_text = [""] if tokenizer is not None: text_batch = tokenizer(text_raw, max_length=params['text_encoder']['max_length'], padding="max_length", truncation=True, return_tensors="pt") text, text_mask = text_batch.input_ids.to(device), text_batch.attention_mask.to(device).bool() text = text_encoder(input_ids=text, attention_mask=text_mask).last_hidden_state uncond_text_batch = tokenizer(neg_text, max_length=params['text_encoder']['max_length'], padding="max_length", truncation=True, return_tensors="pt") uncond_text, uncond_text_mask = uncond_text_batch.input_ids.to(device), uncond_text_batch.attention_mask.to(device).bool() uncond_text = text_encoder(input_ids=uncond_text, attention_mask=uncond_text_mask).last_hidden_state else: text, text_mask = None, None guidance_scale = None codec_dim = params['model']['out_chans'] unet.eval() if random_seed is not None: generator = torch.Generator(device=device).manual_seed(random_seed) else: generator = torch.Generator(device=device) generator.seed() noise_scheduler.set_timesteps(ddim_steps) # init noise noise = torch.randn((1, codec_dim, audio_frames), generator=generator, device=device) latents = noise for t in noise_scheduler.timesteps: latents = noise_scheduler.scale_model_input(latents, t) if guidance_scale: latents_combined = torch.cat([latents, latents], dim=0) text_combined = torch.cat([text, uncond_text], dim=0) text_mask_combined = torch.cat([text_mask, uncond_text_mask], dim=0) if gt is not None: gt_combined = torch.cat([gt, gt], dim=0) gt_mask_combined = torch.cat([gt_mask, gt_mask], dim=0) else: gt_combined = None gt_mask_combined = None output_combined, _ = unet(latents_combined, t, text_combined, context_mask=text_mask_combined, cls_token=None, gt=gt_combined, mae_mask_infer=gt_mask_combined) output_text, output_uncond = torch.chunk(output_combined, 2, dim=0) output_pred = output_uncond + guidance_scale * (output_text - output_uncond) if guidance_rescale > 0.0: output_pred = rescale_noise_cfg(output_pred, output_text, guidance_rescale=guidance_rescale) else: output_pred, mae_mask = unet(latents, t, text, context_mask=text_mask, cls_token=None, gt=gt, mae_mask_infer=gt_mask) latents = noise_scheduler.step(model_output=output_pred, timestep=t, sample=latents, eta=eta, generator=generator).prev_sample pred = scale_shift_re(latents, params['autoencoder']['scale'], params['autoencoder']['shift']) if gt is not None: pred[~gt_mask] = gt[~gt_mask] pred_wav = autoencoder(embedding=pred) return pred_wav @torch.no_grad() def eval_udit(autoencoder, unet, tokenizer, text_encoder, params, noise_scheduler, val_df, subset, audio_frames, mae=False, guidance_scale=3, guidance_rescale=0.0, ddim_steps=50, eta=1, random_seed=2023, device='cuda', epoch=0, save_path='logs/eval/', val_num=5): val_df = pd.read_csv(val_df) val_df = val_df[val_df['split'] == subset] if mae: val_df = val_df[val_df['audio_length'] != 0] save_path = save_path + str(epoch) + '/' os.makedirs(save_path, exist_ok=True) for i in tqdm(range(len(val_df))): row = val_df.iloc[i] text = [row['caption']] if mae: audio_path = params['data']['val_dir'] + str(row['audio_path']) gt, sr = librosa.load(audio_path, sr=params['data']['sr']) gt = gt / (np.max(np.abs(gt)) + 1e-9) sf.write(save_path + text[0] + '_gt.wav', gt, samplerate=params['data']['sr']) num_samples = 10 * sr if len(gt) < num_samples: padding = num_samples - len(gt) gt = np.pad(gt, (0, padding), 'constant') else: gt = gt[:num_samples] gt = torch.tensor(gt).unsqueeze(0).unsqueeze(1).to(device) gt = autoencoder(audio=gt) B, D, L = gt.shape mask_len = int(L * 0.2) gt_mask = torch.zeros(B, D, L).to(device) for _ in range(2): start = random.randint(0, L - mask_len) gt_mask[:, :, start:start + mask_len] = 1 gt_mask = gt_mask.bool() else: gt = None gt_mask = None pred = inference(autoencoder, unet, gt, gt_mask, tokenizer, text_encoder, params, noise_scheduler, text, neg_text=None, audio_frames=audio_frames, guidance_scale=guidance_scale, guidance_rescale=guidance_rescale, ddim_steps=ddim_steps, eta=eta, random_seed=random_seed, device=device) pred = pred.cpu().numpy().squeeze(0).squeeze(0) sf.write(save_path + text[0] + '.wav', pred, samplerate=params['data']['sr']) if i + 1 >= val_num: break