From d1d1ae32a145b2cdbe94e2cc1a17b8e33173a707 Mon Sep 17 00:00:00 2001 From: James Betker Date: Thu, 10 Feb 2022 22:55:46 -0700 Subject: [PATCH] audio diffusion frechet distance measurement! --- .../gen/use_diffuse_voice_translation.py | 25 ++-- codes/trainer/eval/audio_diffusion_fid.py | 115 ++++++++++++++++++ codes/utils/util.py | 10 ++ 3 files changed, 140 insertions(+), 10 deletions(-) create mode 100644 codes/trainer/eval/audio_diffusion_fid.py diff --git a/codes/scripts/audio/gen/use_diffuse_voice_translation.py b/codes/scripts/audio/gen/use_diffuse_voice_translation.py index e5d8c209..61bf9dce 100644 --- a/codes/scripts/audio/gen/use_diffuse_voice_translation.py +++ b/codes/scripts/audio/gen/use_diffuse_voice_translation.py @@ -9,14 +9,8 @@ from data.util import find_files_of_type, is_audio_file from scripts.audio.gen.speech_synthesis_utils import do_spectrogram_diffusion, \ load_discrete_vocoder_diffuser, wav_to_mel, convert_mel_to_codes from utils.audio import plot_spectrogram -from utils.util import load_model_from_config - - -def ceil_multiple(base, multiple): - res = base % multiple - if res == 0: - return base - return base + (multiple - res) +from utils.util import load_model_from_config, ceil_multiple +import torch.nn.functional as F def get_ctc_codes_for(src_clip_path): @@ -34,6 +28,18 @@ def get_ctc_codes_for(src_clip_path): return torch.argmax(logits, dim=-1), clip +def determine_output_size(codes, base_sample_rate) + aligned_codes_compression_factor = base_sample_rate * 221 // 11025 + output_size = codes.shape[-1]*aligned_codes_compression_factor + padded_size = ceil_multiple(output_size, 2048) + padding_added = padded_size - output_size + padding_needed_for_codes = padding_added // aligned_codes_compression_factor + if padding_needed_for_codes > 0: + codes = F.pad(codes, (0, padding_needed_for_codes)) + output_shape = (1, 1, padded_size) + return output_shape, codes + + if __name__ == '__main__': provided_voices = { # Male @@ -70,7 +76,6 @@ if __name__ == '__main__': diffusion = load_model_from_config(args.opt, args.diffusion_model_name, also_load_savepoint=False, load_path=args.diffusion_model_path, device='cpu').eval() diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=args.diffusion_steps, schedule='cosine') - aligned_codes_compression_factor = base_sample_rate * 221 // 11025 sr_diffusion = load_model_from_config(args.sr_opt, args.sr_diffusion_model_name, also_load_savepoint=False, load_path=args.sr_diffusion_model_path, device='cpu').eval() sr_diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=args.diffusion_steps, schedule='linear') @@ -90,7 +95,7 @@ if __name__ == '__main__': torchaudio.save(os.path.join(args.output_path, f'{e}_source_clip.wav'), src_clip.unsqueeze(0).cpu(), 16000) print("Performing initial diffusion..") - output_shape = (1, 1, ceil_multiple(aligned_codes.shape[-1]*aligned_codes_compression_factor, 2048)) + output_shape, aligned_codes = determine_output_size(aligned_codes, base_sample_rate) diffusion = diffusion.cuda() output_base = diffuser.p_sample_loop(diffusion, output_shape, noise=torch.zeros(output_shape, device=args.device), model_kwargs={'tokens': aligned_codes, diff --git a/codes/trainer/eval/audio_diffusion_fid.py b/codes/trainer/eval/audio_diffusion_fid.py new file mode 100644 index 00000000..c38d6c7a --- /dev/null +++ b/codes/trainer/eval/audio_diffusion_fid.py @@ -0,0 +1,115 @@ +import os +import os.path as osp +import torch +import torchaudio +import torchvision +from pytorch_fid import fid_score +from pytorch_fid.fid_score import calculate_frechet_distance +from torch import distributed +from tqdm import tqdm +from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor +import torch.nn.functional as F +import numpy as np + +import trainer.eval.evaluator as evaluator +from data.audio.paired_voice_audio_dataset import load_tsv_aligned_codes +from data.audio.unsupervised_audio_dataset import load_audio +from scripts.audio.gen.speech_synthesis_utils import load_discrete_vocoder_diffuser + + +class StyleTransferEvaluator(evaluator.Evaluator): + """ + Evaluator produces generate from a diffusion model, then uses a pretrained wav2vec model to compute a frechet + distance between real and fake samples. + """ + def __init__(self, model, opt_eval, env): + super().__init__(model, opt_eval, env, uses_all_ddp=True) + self.real_path = opt_eval['eval_tsv'] + self.data = load_tsv_aligned_codes(self.real_path) + if distributed.is_initialized() and distributed.get_world_size() > 1: + self.skip = distributed.get_world_size() # One batch element per GPU. + else: + self.skip = 1 + diffusion_steps = opt_get(opt_eval, ['diffusion_steps'], 50) + diffusion_schedule = opt_get(opt_eval, ['diffusion_schedule'], 'cosine') + self.diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=diffusion_steps, schedule=diffusion_schedule) + self.dev = self.env['device'] + + def perform_diffusion(self, audio, codes, sample_rate=5500): + real_resampled = torchaudio.functional.resample(audio, 22050, sample_rate).unsqueeze(0) + aligned_codes_compression_factor = sample_rate * 221 // 11025 + output_size = codes.shape[-1]*aligned_codes_compression_factor + padded_size = ceil_multiple(output_size, 2048) + padding_added = padded_size - output_size + padding_needed_for_codes = padding_added // aligned_codes_compression_factor + if padding_needed_for_codes > 0: + codes = F.pad(codes, (0, padding_needed_for_codes)) + output_shape = (1, 1, padded_size) + gen = self.diffuser.p_sample_loop(self.model, output_shape, + model_kwargs={'tokens': codes.unsqueeze(0), + 'conditioning_input': real_resampled}) + return gen, real_resampled, sample_rate + + def project(self, projector, sample, sample_rate): + sample = torchaudio.functional.resample(sample, sample_rate, 16000) + sample = (sample - sample.mean()) / torch.sqrt(sample.var() + 1e-7) + return projector(sample.squeeze(1), output_hidden_states=True).hidden_states[-1].squeeze(0) # Getting rid of the batch dimension means it's just [seq_len,hidden_states] + + def compute_frechet_distance(self, proj1, proj2): + # I really REALLY FUCKING HATE that this is going to numpy. Why does "pytorch_fid" operate in numpy land. WHY? + proj1 = proj1.cpu().numpy() + proj2 = proj2.cpu().numpy() + mu1 = np.mean(proj1, axis=0) + mu2 = np.mean(proj2, axis=0) + sigma1 = np.cov(proj1, rowvar=False) + sigma2 = np.cov(proj2, rowvar=False) + return torch.tensor(calculate_frechet_distance(mu1, sigma1, mu2, sigma2)) + + def perform_eval(self): + save_path = osp.join(self.env['base_path'], "../", "audio_eval", str(self.env["step"])) + os.makedirs(save_path, exist_ok=True) + + projector = Wav2Vec2ForCTC.from_pretrained(f"facebook/wav2vec2-large").to(self.dev) + projector.eval() + + # Attempt to fix the random state as much as possible. RNG state will be restored before returning. + rng_state = torch.get_rng_state() + torch.manual_seed(5) + self.model.eval() + + with torch.no_grad(): + gen_projections = [] + real_projections = [] + for i in tqdm(list(range(0, len(self.data), self.skip))): + path, text, codes = self.data[i + self.env['rank']] + audio = load_audio(path, 22050).to(self.dev) + codes = codes.to(self.dev) + sample, ref, sample_rate = self.perform_diffusion(audio, codes) + + gen_projections.append(self.project(projector, sample, sample_rate).cpu()) # Store on CPU to avoid wasting GPU memory. + real_projections.append(self.project(projector, ref, sample_rate).cpu()) + + torchaudio.save(os.path.join(save_path, f"{self.env['rank']}_{i}_gen.wav"), sample.squeeze(0).cpu(), sample_rate) + torchaudio.save(os.path.join(save_path, f"{self.env['rank']}_{i}_real.wav"), ref.squeeze(0).cpu(), sample_rate) + gen_projections = torch.cat(gen_projections, dim=0) + real_projections = torch.cat(real_projections, dim=0) + fid = self.compute_frechet_distance(gen_projections, real_projections) + + if distributed.is_initialized() and distributed.get_world_size() > 1: + fid = distributed.all_reduce(fid) / distributed.get_world_size() + + self.model.train() + torch.set_rng_state(rng_state) + + return {"fid": fid} + + +if __name__ == '__main__': + from utils.util import load_model_from_config, ceil_multiple, opt_get + + diffusion = load_model_from_config('X:\\dlas\\experiments\\train_diffusion_tts5_medium.yml', 'generator', + also_load_savepoint=False, load_path='X:\\dlas\\experiments\\train_diffusion_tts5_medium\\models\\73000_generator_ema.pth').cuda() + opt_eval = {'eval_tsv': 'Y:\\libritts\\test-clean\\transcribed-brief-w2v.tsv', 'diffusion_steps': 50} + env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval', 'step': 500, 'device': 'cuda'} + eval = StyleTransferEvaluator(diffusion, opt_eval, env) + eval.perform_eval() \ No newline at end of file diff --git a/codes/utils/util.py b/codes/utils/util.py index 92b84cf4..f8d9fb26 100644 --- a/codes/utils/util.py +++ b/codes/utils/util.py @@ -498,3 +498,13 @@ def map_cuda_to_correct_device(storage, loc): return storage.cuda(torch.cuda.current_device()) else: return storage.cpu() + + +def ceil_multiple(base, multiple): + """ + Returns the next closest multiple >= base. + """ + res = base % multiple + if res == 0: + return base + return base + (multiple - res)