From bac9a8b728190afc687f1863639b235d30c298b9 Mon Sep 17 00:00:00 2001 From: James Betker Date: Mon, 4 Jul 2022 08:38:47 -0600 Subject: [PATCH] Make MDF compatible with ar_prior models --- codes/trainer/eval/music_diffusion_fid.py | 52 +++++++++++++++++----- codes/trainer/injectors/audio_injectors.py | 2 +- codes/utils/music_utils.py | 8 ++++ 3 files changed, 50 insertions(+), 12 deletions(-) diff --git a/codes/trainer/eval/music_diffusion_fid.py b/codes/trainer/eval/music_diffusion_fid.py index a94ab7a0..184550ef 100644 --- a/codes/trainer/eval/music_diffusion_fid.py +++ b/codes/trainer/eval/music_diffusion_fid.py @@ -21,9 +21,9 @@ from models.clip.contrastive_audio import ContrastiveAudio from models.diffusion.gaussian_diffusion import get_named_beta_schedule from models.diffusion.respace import space_timesteps, SpacedDiffusion from trainer.injectors.audio_injectors import denormalize_mel, TorchMelSpectrogramInjector, pixel_shuffle_1d, \ - normalize_mel + normalize_mel, KmeansQuantizerInjector from utils.music_utils import get_music_codegen, get_mel2wav_model, get_cheater_decoder, get_cheater_encoder, \ - get_mel2wav_v3_model + get_mel2wav_v3_model, get_ar_prior from utils.util import opt_get, load_model_from_config @@ -84,7 +84,13 @@ class MusicDiffusionFid(evaluator.Evaluator): conditioning_free=True, conditioning_free_k=1) self.spec_decoder = get_mel2wav_v3_model() # The only reason the other functions don't use v3 is because earlier models were trained with v1 and I want to keep metrics consistent. self.local_modules['spec_decoder'] = self.spec_decoder - + elif 'from_ar_prior' == mode: + self.diffusion_fn = self.perform_diffusion_from_codes_ar_prior + self.local_modules['cheater_encoder'] = get_cheater_encoder() + self.kmeans_inj = KmeansQuantizerInjector({'centroids': '../experiments/music_k_means_centroids.pth', 'in': 'in', 'out': 'out'}, {}) + self.local_modules['ar_prior'] = get_ar_prior() + self.spec_decoder = get_mel2wav_v3_model() + self.local_modules['spec_decoder'] = self.spec_decoder if not hasattr(self, 'spec_decoder'): self.spec_decoder = get_mel2wav_model() self.local_modules['spec_decoder'] = self.spec_decoder @@ -235,6 +241,30 @@ class MusicDiffusionFid(evaluator.Evaluator): return gen_wav, real_wav.squeeze(0), gen_mel, mel_norm, sample_rate + def perform_diffusion_from_codes_ar_prior(self, audio, sample_rate=22050): + assert self.ddim, "DDIM mode expected for reconstructing cheater gen. Do you like to waste resources??" + audio = audio.unsqueeze(0) + + mel = self.spec_fn({'in': audio})['out'] + mel_norm = normalize_mel(mel) + cheater = self.local_modules['cheater_encoder'].to(audio.device)(mel_norm) + cheater_codes = self.kmeans_inj({'in': cheater})['out'] + ar_latent = self.local_modules['ar_prior'].to(audio.device)(cheater_codes, cheater, return_latent=True) + + gen_mel = self.diffuser.ddim_sample_loop(self.model, mel_norm.shape, model_kwargs={'codes': ar_latent}, progress=True) + + gen_mel_denorm = denormalize_mel(gen_mel) + output_shape = (1,16,audio.shape[-1]//16) + self.spec_decoder = self.spec_decoder.to(audio.device) + gen_wav = self.spectral_diffuser.ddim_sample_loop(self.spec_decoder, output_shape, + model_kwargs={'codes': gen_mel_denorm}) + gen_wav = pixel_shuffle_1d(gen_wav, 16) + + real_wav = self.spectral_diffuser.ddim_sample_loop(self.spec_decoder, output_shape, + model_kwargs={'codes': mel}) + real_wav = pixel_shuffle_1d(real_wav, 16) + + return gen_wav, real_wav.squeeze(0), gen_mel, mel_norm, sample_rate def project(self, sample, sample_rate): sample = torchaudio.functional.resample(sample, sample_rate, 22050) @@ -304,17 +334,17 @@ class MusicDiffusionFid(evaluator.Evaluator): if __name__ == '__main__': - diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_cheater_gen_r8.yml', 'generator', + diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_tfd12_finetune_ar_outputs.yml', 'generator', also_load_savepoint=False, - load_path='X:\\dlas\\experiments\\train_music_cheater_gen_v5\\models\\71000_generator_ema.pth' + load_path='X:\\dlas\\experiments\\train_music_diffusion_tfd12_finetune_from_cheater_ar\\models\\7500_generator.pth' ).cuda() - opt_eval = {#'path': 'Y:\\split\\yt-music-eval', # eval music, mostly electronica. :) - 'path': 'E:\\music_eval', # this is music from the training dataset, including a lot more variety. - 'diffusion_steps': 128, - 'conditioning_free': True, 'conditioning_free_k': 2, 'clip_audio': False, 'use_ddim': True, - 'diffusion_schedule': 'linear', 'diffusion_type': 'cheater_gen', + opt_eval = {'path': 'Y:\\split\\yt-music-eval', # eval music, mostly electronica. :) + #'path': 'E:\\music_eval', # this is music from the training dataset, including a lot more variety. + 'diffusion_steps': 32, + 'conditioning_free': True, 'conditioning_free_k': 1, 'use_ddim': True, # 'clip_audio': False, + 'diffusion_schedule': 'linear', 'diffusion_type': 'from_ar_prior', #'partial_low': 128, 'partial_high': 192 } - env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 225, 'device': 'cuda', 'opt': {}} + env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 230, 'device': 'cuda', 'opt': {}} eval = MusicDiffusionFid(diffusion, opt_eval, env) print(eval.perform_eval()) diff --git a/codes/trainer/injectors/audio_injectors.py b/codes/trainer/injectors/audio_injectors.py index 2e656d03..53cb1815 100644 --- a/codes/trainer/injectors/audio_injectors.py +++ b/codes/trainer/injectors/audio_injectors.py @@ -432,7 +432,7 @@ class KmeansQuantizerInjector(Injector): class MusicCheaterArInjector(Injector): def __init__(self, opt, env): super().__init__(opt, env) - self.cheater_ar = ConditioningAR(1024, layers=24, dropout=0, cond_free_percent=0) + self.cheater_ar = ConditioningAR(1024, layers=24, dropout=0, cond_free_percent=0).eval() self.cheater_ar.load_state_dict(torch.load('../experiments/music_cheater_ar.pth', map_location=torch.device('cpu'))) self.cond_key = opt['cheater_latent_key'] self.needs_move = True diff --git a/codes/utils/music_utils.py b/codes/utils/music_utils.py index 5f358428..9161c386 100644 --- a/codes/utils/music_utils.py +++ b/codes/utils/music_utils.py @@ -50,3 +50,11 @@ def get_cheater_decoder(): model.load_state_dict(torch.load(f'../experiments/music_cheater_decoder.pth', map_location=torch.device('cpu'))) model = model.eval() return model + + +def get_ar_prior(): + from models.audio.music.cheater_gen_ar import ConditioningAR + cheater_ar = ConditioningAR(1024, layers=24, dropout=0, cond_free_percent=0) + cheater_ar.load_state_dict(torch.load('../experiments/music_cheater_ar.pth', map_location=torch.device('cpu'))) + #cheater_ar = cheater_ar.eval() + return cheater_ar \ No newline at end of file