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

Update mdf to be compatible with cheater_gen

Browse Source
This commit is contained in:
James Betker 2022-06-27 10:11:23 -06:00
parent 69b614e08a
commit 0278576f37
3 changed files with 83 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
codes/models/audio/music/tfdpc_v4.py
View File

@ -270,7 +270,7 @@ def test_cheater_model():
def inference_tfdpc4_with_cheater():
with torch.no_grad():
os.makedirs('results/tfdpc_v3', exist_ok=True)
os.makedirs('results/tfdpc_v4', exist_ok=True)
#length = 40 * 22050 // 256 // 16
samples = {'electronica1': load_audio('Y:\\split\\yt-music-eval\\00001.wav', 22050),
@ -290,7 +290,7 @@ def inference_tfdpc4_with_cheater():
model = TransformerDiffusionWithPointConditioning(in_channels=256, out_channels=512, model_channels=1024,
contraction_dim=512, num_heads=8, num_layers=12, dropout=0,
use_fp16=False, unconditioned_percentage=0).eval().cuda()
model.load_state_dict(torch.load('x:/dlas/experiments/train_music_cheater_gen_v3/models/59000_generator_ema.pth'))
model.load_state_dict(torch.load('x:/dlas/experiments/train_music_cheater_gen_v4/models/28000_generator_ema.pth'))
from trainer.injectors.audio_injectors import TorchMelSpectrogramInjector
spec_fn = TorchMelSpectrogramInjector({'n_mel_channels': 256, 'mel_fmax': 11000, 'filter_length': 16000, 'true_normalization': True,
@ -308,7 +308,7 @@ def inference_tfdpc4_with_cheater():
conditioning_free=True, conditioning_free_k=1)
# Conventional decoding method:
gen_cheater = diffuser.ddim_sample_loop(model, (1,256,length), progress=True, model_kwargs={'true_cheater': ref_cheater})
gen_cheater = diffuser.ddim_sample_loop(model, (1,256,length), progress=True, model_kwargs={'conditioning_input': ref_cheater})
# Guidance decoding method:
#mask = torch.ones_like(ref_cheater)
@ -330,7 +330,7 @@ def inference_tfdpc4_with_cheater():
cheater_to_mel = wrap.diff
gen_mel = diffuser.ddim_sample_loop(cheater_to_mel, (1,256,gen_cheater.shape[-1]*16), progress=True,
model_kwargs={'codes': gen_cheater.permute(0,2,1)})
torchvision.utils.save_image((gen_mel + 1)/2, f'results/tfdpc_v3/{k}.png')
torchvision.utils.save_image((gen_mel + 1)/2, f'results/tfdpc_v4/{k}.png')
from utils.music_utils import get_mel2wav_v3_model
m2w = get_mel2wav_v3_model().cuda()
@ -344,9 +344,9 @@ def inference_tfdpc4_with_cheater():
from trainer.injectors.audio_injectors import pixel_shuffle_1d
gen_wav = pixel_shuffle_1d(gen_wav, 16)
torchaudio.save(f'results/tfdpc_v3/{k}.wav', gen_wav.squeeze(1).cpu(), 22050)
torchaudio.save(f'results/tfdpc_v3/{k}_ref.wav', sample.unsqueeze(0).cpu(), 22050)
torchaudio.save(f'results/tfdpc_v4/{k}.wav', gen_wav.squeeze(1).cpu(), 22050)
torchaudio.save(f'results/tfdpc_v4/{k}_ref.wav', sample.unsqueeze(0).cpu(), 22050)
if __name__ == '__main__':
test_cheater_model()
#inference_tfdpc4_with_cheater()
#test_cheater_model()
inference_tfdpc4_with_cheater()

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

@ -22,7 +22,8 @@ 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
from utils.music_utils import get_music_codegen, get_mel2wav_model
from utils.music_utils import get_music_codegen, get_mel2wav_model, get_cheater_decoder, get_cheater_encoder, \
get_mel2wav_v3_model
from utils.util import opt_get, load_model_from_config
@ -34,6 +35,8 @@ class MusicDiffusionFid(evaluator.Evaluator):
super().__init__(model, opt_eval, env, uses_all_ddp=True)
self.real_path = opt_eval['path']
self.data = self.load_data(self.real_path)
self.clip = opt_get(opt_eval, ['clip_audio'], True) # Recommend setting true for more efficient eval passes.
self.ddim = opt_get(opt_eval, ['use_ddim'], False)
if distributed.is_initialized() and distributed.get_world_size() > 1:
self.skip = distributed.get_world_size() # One batch element per GPU.
else:
@ -48,17 +51,16 @@ class MusicDiffusionFid(evaluator.Evaluator):
self.diffuser = SpacedDiffusion(use_timesteps=space_timesteps(4000, [diffusion_steps]), model_mean_type='epsilon',
model_var_type='learned_range', loss_type='mse', betas=get_named_beta_schedule(diffusion_schedule, 4000),
conditioning_free=conditioning_free_diffusion_enabled, conditioning_free_k=conditioning_free_k)
self.spectral_diffuser = SpacedDiffusion(use_timesteps=space_timesteps(4000, [100]), model_mean_type='epsilon',
self.spectral_diffuser = SpacedDiffusion(use_timesteps=space_timesteps(4000, [16 if self.ddim else 100]), model_mean_type='epsilon',
model_var_type='learned_range', loss_type='mse', betas=get_named_beta_schedule('linear', 4000),
conditioning_free=False, conditioning_free_k=1)
self.dev = self.env['device']
mode = opt_get(opt_eval, ['diffusion_type'], 'spec_decode')
self.spec_decoder = get_mel2wav_model()
self.projector = ContrastiveAudio(model_dim=512, transformer_heads=8, dropout=0, encoder_depth=8, mel_channels=256)
self.projector.load_state_dict(torch.load('../experiments/music_eval_projector.pth', map_location=torch.device('cpu')))
self.local_modules = {'spec_decoder': self.spec_decoder, 'projector': self.projector}
self.local_modules = {'projector': self.projector}
if mode == 'spec_decode':
self.diffusion_fn = self.perform_diffusion_spec_decode
self.squeeze_ratio = opt_eval['squeeze_ratio']
@ -73,6 +75,16 @@ class MusicDiffusionFid(evaluator.Evaluator):
partial_high=opt_eval['partial_high'])
elif 'from_codes_quant_gradual_decode' == mode:
self.diffusion_fn = self.perform_diffusion_from_codes_quant_gradual_decode
elif 'cheater_gen' == mode:
self.diffusion_fn = self.perform_reconstruction_from_cheater_gen
self.local_modules['cheater_encoder'] = get_cheater_encoder()
self.local_modules['cheater_decoder'] = get_cheater_decoder()
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
if not hasattr(self, 'spec_decoder'):
self.spec_decoder = get_mel2wav_model()
self.local_modules['spec_decoder'] = self.spec_decoder
self.spec_fn = TorchMelSpectrogramInjector({'n_mel_channels': 256, 'mel_fmax': 11000, 'filter_length': 16000,
'normalize': True, 'in': 'in', 'out': 'out'}, {})
@ -191,6 +203,35 @@ class MusicDiffusionFid(evaluator.Evaluator):
return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate
def perform_reconstruction_from_cheater_gen(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)
# 1. Generate the cheater latent using the input as a reference.
gen_cheater = self.diffuser.ddim_sample_loop(self.model, cheater.shape, progress=True, model_kwargs={'conditioning_input': cheater})
# 2. Decode the cheater into a MEL
gen_mel = self.diffuser.ddim_sample_loop(self.local_modules['cheater_decoder'].diff.to(audio.device), (1,256,gen_cheater.shape[-1]*16), progress=True,
model_kwargs={'codes': gen_cheater.permute(0,2,1)})
# 3. And then the MEL back into a spectrogram
output_shape = (1,16,audio.shape[-1]//16)
self.spec_decoder = self.spec_decoder.to(audio.device)
gen_mel_denorm = denormalize_mel(gen_mel)
gen_wav = self.spectral_diffuser.p_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.p_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)
@ -229,7 +270,8 @@ class MusicDiffusionFid(evaluator.Evaluator):
for i in tqdm(list(range(0, len(self.data), self.skip))):
path = self.data[(i + self.env['rank']) % len(self.data)]
audio = load_audio(path, 22050).to(self.dev)
audio = audio[:, :100000]
if self.clip:
audio = audio[:, :100000]
sample, ref, sample_mel, ref_mel, sample_rate = self.diffusion_fn(audio)
gen_projections.append(self.project(sample, sample_rate).cpu()) # Store on CPU to avoid wasting GPU memory.
@ -259,17 +301,17 @@ class MusicDiffusionFid(evaluator.Evaluator):
if __name__ == '__main__':
diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_waveform_gen.yml', 'generator',
diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_cheater_gen.yml', 'generator',
also_load_savepoint=False,
load_path='X:\\dlas\\experiments\\train_music_waveform_gen_retry\\models\\22000_generator_ema.pth'
load_path='X:\\dlas\\experiments\\train_music_cheater_gen_v4\\models\\28000_generator_ema.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': 100,
'conditioning_free': False, 'conditioning_free_k': 1,
'diffusion_schedule': 'linear', 'diffusion_type': 'spec_decode',
'diffusion_steps': 32,
'conditioning_free': False, 'conditioning_free_k': 1, 'clip_audio': False, 'use_ddim': True,
'diffusion_schedule': 'linear', 'diffusion_type': 'cheater_gen',
#'partial_low': 128, 'partial_high': 192
}
env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 100, 'device': 'cuda', 'opt': {}}
env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 200, 'device': 'cuda', 'opt': {}}
eval = MusicDiffusionFid(diffusion, opt_eval, env)
print(eval.perform_eval())

22
codes/utils/music_utils.py
View File

@ -10,6 +10,7 @@ def get_mel2wav_model():
model.eval()
return model
def get_mel2wav_v3_model():
from models.audio.music.unet_diffusion_waveform_gen3 import DiffusionWaveformGen
model = DiffusionWaveformGen(model_channels=256, in_channels=16, in_mel_channels=256, out_channels=32, channel_mult=[1,1.5,2,4],
@ -19,6 +20,7 @@ def get_mel2wav_v3_model():
model.eval()
return model
def get_music_codegen():
from models.audio.mel2vec import ContrastiveTrainingWrapper
model = ContrastiveTrainingWrapper(mel_input_channels=256, inner_dim=1024, layers=24, dropout=0,
@ -28,3 +30,23 @@ def get_music_codegen():
model.load_state_dict(torch.load(f"../experiments/m2v_music.pth", map_location=torch.device('cpu')))
model = model.eval()
return model
def get_cheater_encoder():
from models.audio.music.gpt_music2 import UpperEncoder
encoder = UpperEncoder(256, 1024, 256)
encoder.load_state_dict(
torch.load('../experiments/music_cheater_encoder_256.pth', map_location=torch.device('cpu')))
encoder = encoder.eval()
return encoder
def get_cheater_decoder():
from models.audio.music.transformer_diffusion12 import TransformerDiffusionWithCheaterLatent
model = TransformerDiffusionWithCheaterLatent(in_channels=256, out_channels=512, model_channels=1024,
contraction_dim=512, prenet_channels=1024, input_vec_dim=256,
prenet_layers=6, num_heads=8, num_layers=16, new_code_expansion=True,
dropout=0, unconditioned_percentage=0)
model.load_state_dict(torch.load(f'../experiments/music_cheater_decoder.pth', map_location=torch.device('cpu')))
model = model.eval()
return model