mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Make MDF compatible with ar_prior models

Browse Source
This commit is contained in:
James Betker 2022-07-04 08:38:47 -06:00
parent 455943779b
commit bac9a8b728
3 changed files with 50 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
codes/trainer/eval/music_diffusion_fid.py
View File

@ -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())

2
codes/trainer/injectors/audio_injectors.py
View File

@ -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

8
codes/utils/music_utils.py
View File

@ -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