mdf fixes + support for tfd-based waveform gen

This commit is contained in:
James Betker 2022-06-19 15:07:24 -06:00
parent cb7569ee5e
commit 368dca18b1

View File

@ -79,23 +79,17 @@ class MusicDiffusionFid(evaluator.Evaluator):
return list(glob(f'{path}/*.wav')) return list(glob(f'{path}/*.wav'))
def perform_diffusion_spec_decode(self, audio, sample_rate=22050): def perform_diffusion_spec_decode(self, audio, sample_rate=22050):
if sample_rate != sample_rate:
real_resampled = torchaudio.functional.resample(audio, 22050, sample_rate).unsqueeze(0)
else:
real_resampled = audio real_resampled = audio
audio = audio.unsqueeze(0) audio = audio.unsqueeze(0)
output_shape = (1, 16, audio.shape[-1] // 16) output_shape = (1, 256, audio.shape[-1] // 256)
mel = self.spec_fn({'in': audio})['out'] mel = self.spec_fn({'in': audio})['out']
gen = self.diffuser.p_sample_loop(self.model, output_shape, gen = self.diffuser.p_sample_loop(self.model, output_shape,
model_kwargs={'aligned_conditioning': mel}) model_kwargs={'codes': mel})
gen = pixel_shuffle_1d(gen, 16) gen = pixel_shuffle_1d(gen, 256)
return gen, real_resampled, normalize_mel(self.spec_fn({'in': gen})['out']), normalize_mel(mel), sample_rate return gen, real_resampled, normalize_mel(self.spec_fn({'in': gen})['out']), normalize_mel(mel), sample_rate
def perform_diffusion_from_codes(self, audio, sample_rate=22050): def perform_diffusion_from_codes(self, audio, sample_rate=22050):
if sample_rate != sample_rate:
real_resampled = torchaudio.functional.resample(audio, 22050, sample_rate).unsqueeze(0)
else:
real_resampled = audio real_resampled = audio
audio = audio.unsqueeze(0) audio = audio.unsqueeze(0)
@ -116,9 +110,6 @@ class MusicDiffusionFid(evaluator.Evaluator):
return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate
def perform_diffusion_from_codes_quant(self, audio, sample_rate=22050): def perform_diffusion_from_codes_quant(self, audio, sample_rate=22050):
if sample_rate != sample_rate:
real_resampled = torchaudio.functional.resample(audio, 22050, sample_rate).unsqueeze(0)
else:
real_resampled = audio real_resampled = audio
audio = audio.unsqueeze(0) audio = audio.unsqueeze(0)
@ -148,9 +139,6 @@ class MusicDiffusionFid(evaluator.Evaluator):
return gen_wav, real_wav.squeeze(0), gen_mel, mel_norm, sample_rate return gen_wav, real_wav.squeeze(0), gen_mel, mel_norm, sample_rate
def perform_partial_diffusion_from_codes_quant(self, audio, sample_rate=22050, partial_low=0, partial_high=256): def perform_partial_diffusion_from_codes_quant(self, audio, sample_rate=22050, partial_low=0, partial_high=256):
if sample_rate != sample_rate:
real_resampled = torchaudio.functional.resample(audio, 22050, sample_rate).unsqueeze(0)
else:
real_resampled = audio real_resampled = audio
audio = audio.unsqueeze(0) audio = audio.unsqueeze(0)
@ -174,9 +162,6 @@ class MusicDiffusionFid(evaluator.Evaluator):
return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate
def perform_diffusion_from_codes_quant_gradual_decode(self, audio, sample_rate=22050): def perform_diffusion_from_codes_quant_gradual_decode(self, audio, sample_rate=22050):
if sample_rate != sample_rate:
real_resampled = torchaudio.functional.resample(audio, 22050, sample_rate).unsqueeze(0)
else:
real_resampled = audio real_resampled = audio
audio = audio.unsqueeze(0) audio = audio.unsqueeze(0)
@ -273,17 +258,17 @@ class MusicDiffusionFid(evaluator.Evaluator):
if __name__ == '__main__': if __name__ == '__main__':
diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_diffusion_tfd_quant.yml', 'generator', diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_waveform_gen.yml', 'generator',
also_load_savepoint=False, also_load_savepoint=False,
load_path='X:\\dlas\\experiments\\train_music_diffusion_tfd12\\models\\41500_generator_ema.pth' load_path='X:\\dlas\\experiments\\train_music_waveform_gen_retry\\models\\22000_generator_ema.pth'
).cuda() ).cuda()
opt_eval = {'path': 'Y:\\split\\yt-music-eval', # eval music, mostly electronica. :) 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. #'path': 'E:\\music_eval', # this is music from the training dataset, including a lot more variety.
'diffusion_steps': 200, 'diffusion_steps': 100,
'conditioning_free': True, 'conditioning_free_k': 2, 'conditioning_free': False, 'conditioning_free_k': 1,
'diffusion_schedule': 'linear', 'diffusion_type': 'from_codes_quant', 'diffusion_schedule': 'linear', 'diffusion_type': 'spec_decode',
#'partial_low': 128, 'partial_high': 192 #'partial_low': 128, 'partial_high': 192
} }
env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 605, 'device': 'cuda', 'opt': {}} env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 100, 'device': 'cuda', 'opt': {}}
eval = MusicDiffusionFid(diffusion, opt_eval, env) eval = MusicDiffusionFid(diffusion, opt_eval, env)
print(eval.perform_eval()) print(eval.perform_eval())