diff --git a/codes/data/audio/unsupervised_audio_dataset.py b/codes/data/audio/unsupervised_audio_dataset.py index 7808033e..ce73c440 100644 --- a/codes/data/audio/unsupervised_audio_dataset.py +++ b/codes/data/audio/unsupervised_audio_dataset.py @@ -208,6 +208,8 @@ if __name__ == '__main__': for b in tqdm(dl): for b_ in range(b['clip'].shape[0]): #pass - #torchaudio.save(f'{i}_clip_{b_}.wav', b['clip'][b_], ds.sampling_rate) - #torchaudio.save(f'{i}_alt_clip_{b_}.wav', b['alt_clips'][b_], ds.sampling_rate) + torchaudio.save(f'{i}_clip_{b_}.wav', b['clip'][b_], ds.sampling_rate) + torchaudio.save(f'{i}_alt_clip_{b_}.wav', b['alt_clips'][b_], ds.sampling_rate) i += 1 + if i > 200: + break diff --git a/codes/models/audio/music/transformer_diffusion9.py b/codes/models/audio/music/transformer_diffusion9.py deleted file mode 100644 index 6670b7f3..00000000 --- a/codes/models/audio/music/transformer_diffusion9.py +++ /dev/null @@ -1,334 +0,0 @@ -import itertools - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from models.audio.music.music_quantizer2 import MusicQuantizer2 -from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear -from models.diffusion.unet_diffusion import TimestepBlock, AttentionBlock, TimestepEmbedSequential -from models.lucidrains.x_transformers import Encoder -from trainer.networks import register_model -from utils.util import checkpoint, print_network - - -def is_latent(t): - return t.dtype == torch.float - -def is_sequence(t): - return t.dtype == torch.long - - -class MultiGroupEmbedding(nn.Module): - def __init__(self, tokens, groups, dim): - super().__init__() - self.m = nn.ModuleList([nn.Embedding(tokens, dim // groups) for _ in range(groups)]) - - def forward(self, x): - h = [embedding(x[:, :, i]) for i, embedding in enumerate(self.m)] - return torch.cat(h, dim=-1) - - -class ResBlock(TimestepBlock): - def __init__( - self, - channels, - emb_channels, - dropout, - out_channels=None, - dims=2, - kernel_size=3, - efficient_config=False, - use_scale_shift_norm=False, - ): - super().__init__() - self.channels = channels - self.emb_channels = emb_channels - self.dropout = dropout - self.out_channels = out_channels or channels - self.use_scale_shift_norm = use_scale_shift_norm - padding = {1: 0, 3: 1, 5: 2}[kernel_size] - eff_kernel = 1 if efficient_config else 3 - eff_padding = 0 if efficient_config else 1 - - self.in_layers = nn.Sequential( - normalization(channels), - nn.SiLU(), - conv_nd(dims, channels, self.out_channels, eff_kernel, padding=eff_padding), - ) - - self.emb_layers = nn.Sequential( - nn.SiLU(), - linear( - emb_channels, - 2 * self.out_channels if use_scale_shift_norm else self.out_channels, - ), - ) - self.out_layers = nn.Sequential( - normalization(self.out_channels), - nn.SiLU(), - nn.Dropout(p=dropout), - zero_module( - conv_nd(dims, self.out_channels, self.out_channels, kernel_size, padding=padding) - ), - ) - - if self.out_channels == channels: - self.skip_connection = nn.Identity() - else: - self.skip_connection = conv_nd(dims, channels, self.out_channels, eff_kernel, padding=eff_padding) - - def forward(self, x, emb): - """ - Apply the block to a Tensor, conditioned on a timestep embedding. - - :param x: an [N x C x ...] Tensor of features. - :param emb: an [N x emb_channels] Tensor of timestep embeddings. - :return: an [N x C x ...] Tensor of outputs. - """ - return checkpoint( - self._forward, x, emb - ) - - def _forward(self, x, emb): - h = self.in_layers(x) - emb_out = self.emb_layers(emb).type(h.dtype) - while len(emb_out.shape) < len(h.shape): - emb_out = emb_out[..., None] - if self.use_scale_shift_norm: - out_norm, out_rest = self.out_layers[0], self.out_layers[1:] - scale, shift = torch.chunk(emb_out, 2, dim=1) - h = out_norm(h) * (1 + scale) + shift - h = out_rest(h) - else: - h = h + emb_out - h = self.out_layers(h) - return self.skip_connection(x) + h - - -class DiffusionLayer(TimestepBlock): - def __init__(self, model_channels, dropout, num_heads): - super().__init__() - self.resblk = ResBlock(model_channels, model_channels, dropout, model_channels, dims=1, use_scale_shift_norm=True) - self.attn = AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True) - - def forward(self, x, time_emb): - y = self.resblk(x, time_emb) - return self.attn(y) - - -class TransformerDiffusion(nn.Module): - """ - A diffusion model composed entirely of stacks of transformer layers. Why would you do it any other way? - """ - def __init__( - self, - model_channels=512, - prenet_layers=3, - num_layers=8, - in_channels=256, - input_vec_dim=512, - out_channels=512, # mean and variance - dropout=0, - use_fp16=False, - # Parameters for regularization. - unconditioned_percentage=.1, # This implements a mechanism similar to what is used in classifier-free training. - ): - super().__init__() - - self.in_channels = in_channels - self.model_channels = model_channels - self.out_channels = out_channels - self.dropout = dropout - self.unconditioned_percentage = unconditioned_percentage - self.enable_fp16 = use_fp16 - - self.inp_block = conv_nd(1, in_channels, model_channels, 3, 1, 1) - - self.time_embed = nn.Sequential( - linear(model_channels, model_channels), - nn.SiLU(), - linear(model_channels, model_channels), - ) - - self.input_converter = nn.Linear(input_vec_dim, model_channels) - self.code_converter = Encoder( - dim=model_channels, - depth=prenet_layers, - heads=model_channels//64, - ff_dropout=dropout, - attn_dropout=dropout, - use_rmsnorm=True, - ff_glu=True, - rotary_pos_emb=True, - zero_init_branch_output=True, - ff_mult=1, - ) - - self.unconditioned_embedding = nn.Parameter(torch.randn(1,1,model_channels)) - self.intg = nn.Conv1d(model_channels*2, model_channels, kernel_size=1) - self.layers = TimestepEmbedSequential(*[DiffusionLayer(model_channels, dropout, model_channels // 64) for _ in range(num_layers)]) - - self.out = nn.Sequential( - normalization(model_channels), - nn.SiLU(), - zero_module(conv_nd(1, model_channels, out_channels, 3, padding=1)), - ) - - self.debug_codes = {} - - def get_grad_norm_parameter_groups(self): - groups = { - 'layers': list(self.layers.parameters()) + list(self.inp_block.parameters()), - 'code_converters': list(self.input_converter.parameters()) + list(self.code_converter.parameters()), - 'time_embed': list(self.time_embed.parameters()), - } - return groups - - def timestep_independent(self, codes, expected_seq_len): - code_emb = self.input_converter(codes) - - # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance. - if self.training and self.unconditioned_percentage > 0: - unconditioned_batches = torch.rand((code_emb.shape[0], 1, 1), - device=code_emb.device) < self.unconditioned_percentage - code_emb = torch.where(unconditioned_batches, self.unconditioned_embedding.repeat(codes.shape[0], 1, 1), - code_emb) - code_emb = self.code_converter(code_emb) - - expanded_code_emb = F.interpolate(code_emb.permute(0,2,1), size=expected_seq_len, mode='nearest').permute(0,2,1) - return expanded_code_emb - - def forward(self, x, timesteps, codes=None, conditioning_input=None, precomputed_code_embeddings=None, conditioning_free=False): - if precomputed_code_embeddings is not None: - assert codes is None and conditioning_input is None, "Do not provide precomputed embeddings and the other parameters. It is unclear what you want me to do here." - - unused_params = [] - if conditioning_free: - code_emb = self.unconditioned_embedding.repeat(x.shape[0], x.shape[-1], 1) - unused_params.extend(list(self.code_converter.parameters())) - else: - if precomputed_code_embeddings is not None: - code_emb = precomputed_code_embeddings - else: - code_emb = self.timestep_independent(codes, x.shape[-1]) - unused_params.append(self.unconditioned_embedding) - code_emb = code_emb.permute(0,2,1) - - blk_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) - x = self.inp_block(x) - x = self.intg(torch.cat([x, code_emb], dim=1)) - for layer in self.layers: - x = checkpoint(layer, x, blk_emb) - - x = x.float() - out = self.out(x) - - # Involve probabilistic or possibly unused parameters in loss so we don't get DDP errors. - extraneous_addition = 0 - for p in unused_params: - extraneous_addition = extraneous_addition + p.mean() - out = out + extraneous_addition * 0 - - return out - - -class TransformerDiffusionWithQuantizer(nn.Module): - def __init__(self, freeze_quantizer_until=20000, **kwargs): - super().__init__() - - self.internal_step = 0 - self.freeze_quantizer_until = freeze_quantizer_until - self.diff = TransformerDiffusion(**kwargs) - self.quantizer = MusicQuantizer2(inp_channels=256, inner_dim=[1024], codevector_dim=1024, codebook_size=256, - codebook_groups=2, max_gumbel_temperature=4, min_gumbel_temperature=.5) - self.quantizer.quantizer.temperature = self.quantizer.min_gumbel_temperature - del self.quantizer.up - - def update_for_step(self, step, *args): - self.internal_step = step - qstep = max(0, self.internal_step - self.freeze_quantizer_until) - self.quantizer.quantizer.temperature = max( - self.quantizer.max_gumbel_temperature * self.quantizer.gumbel_temperature_decay ** qstep, - self.quantizer.min_gumbel_temperature, - ) - - def forward(self, x, timesteps, truth_mel, conditioning_input=None, disable_diversity=False, conditioning_free=False): - quant_grad_enabled = self.internal_step > self.freeze_quantizer_until - with torch.set_grad_enabled(quant_grad_enabled): - proj, diversity_loss = self.quantizer(truth_mel, return_decoder_latent=True) - proj = proj.permute(0,2,1) - - # Make sure this does not cause issues in DDP by explicitly using the parameters for nothing. - if not quant_grad_enabled: - unused = 0 - for p in self.quantizer.parameters(): - unused = unused + p.mean() * 0 - proj = proj + unused - diversity_loss = diversity_loss * 0 - - diff = self.diff(x, timesteps, codes=proj, conditioning_input=conditioning_input, conditioning_free=conditioning_free) - if disable_diversity: - return diff - return diff, diversity_loss - - def get_debug_values(self, step, __): - if self.quantizer.total_codes > 0: - return {'histogram_codes': self.quantizer.codes[:self.quantizer.total_codes]} - else: - return {} - - def get_grad_norm_parameter_groups(self): - groups = { - 'attention_layers': list(itertools.chain.from_iterable([lyr.attn.parameters() for lyr in self.diff.layers])), - 'res_layers': list(itertools.chain.from_iterable([lyr.resblk.parameters() for lyr in self.diff.layers])), - 'quantizer_encoder': list(self.quantizer.encoder.parameters()), - 'quant_codebook': [self.quantizer.quantizer.codevectors], - 'out': list(self.diff.out.parameters()), - 'x_proj': list(self.diff.inp_block.parameters()), - 'layers': list(self.diff.layers.parameters()), - 'code_converters': list(self.diff.input_converter.parameters()) + list(self.diff.code_converter.parameters()), - 'time_embed': list(self.diff.time_embed.parameters()), - } - return groups - - -@register_model -def register_transformer_diffusion9(opt_net, opt): - return TransformerDiffusion(**opt_net['kwargs']) - - -@register_model -def register_transformer_diffusion9_with_quantizer(opt_net, opt): - return TransformerDiffusionWithQuantizer(**opt_net['kwargs']) - - -""" -# For TFD5 -if __name__ == '__main__': - clip = torch.randn(2, 256, 400) - aligned_sequence = torch.randn(2,100,512) - cond = torch.randn(2, 256, 400) - ts = torch.LongTensor([600, 600]) - model = TransformerDiffusion(model_channels=3072, model_channels=1536, model_channels=1536) - torch.save(model, 'sample.pth') - print_network(model) - o = model(clip, ts, aligned_sequence, cond) -""" - -if __name__ == '__main__': - clip = torch.randn(2, 256, 400) - cond = torch.randn(2, 256, 400) - ts = torch.LongTensor([600, 600]) - model = TransformerDiffusionWithQuantizer(model_channels=1024, input_vec_dim=1024, num_layers=16, prenet_layers=6) - model.get_grad_norm_parameter_groups() - - quant_weights = torch.load('D:\\dlas\\experiments\\train_music_quant_r4\\models\\5000_generator.pth') - #diff_weights = torch.load('X:\\dlas\\experiments\\train_music_diffusion_tfd5\\models\\48000_generator_ema.pth') - model.quantizer.load_state_dict(quant_weights, strict=False) - #model.diff.load_state_dict(diff_weights) - - torch.save(model.state_dict(), 'sample.pth') - print_network(model) - o = model(clip, ts, clip, cond) - diff --git a/codes/scripts/audio/gen/ctc_codes.py b/codes/scripts/audio/gen/ctc_codes.py index 3f0444cd..7b52254c 100644 --- a/codes/scripts/audio/gen/ctc_codes.py +++ b/codes/scripts/audio/gen/ctc_codes.py @@ -1,9 +1,15 @@ from itertools import groupby import torch +import torchaudio from transformers import Wav2Vec2CTCTokenizer +from data.audio.voice_tokenizer import VoiceBpeTokenizer from models.audio.tts.ctc_code_generator import CtcCodeGenerator +from models.audio.tts.transformer_diffusion_tts import TransformerDiffusionTTS +from scripts.audio.gen.speech_synthesis_utils import load_discrete_vocoder_diffuser, load_univnet_vocoder, load_clvp +from trainer.injectors.audio_injectors import TorchMelSpectrogramInjector, denormalize_mel +from utils.util import load_audio def get_ctc_metadata(codes): @@ -28,26 +34,61 @@ def get_ctc_metadata(codes): return rcodes, pads, repeats +def decode_ctc_metadata(rcodes, pads, repeats): + outp = [] + for s in range(rcodes.shape[-1]): + outp = outp + [0 for _ in range(pads[s])] + outp = outp + [rcodes[s].item() for _ in range(repeats[s])] + return torch.tensor(outp, device=rcodes.device) + + +def diffuse(text, codes, cond): + RATIO = 263/140 + + codes = codes.cuda(); + cond = cond.cuda() + + bpe_tokenizer = VoiceBpeTokenizer('../experiments/bpe_lowercase_asr_256.json') + clvp = load_clvp().cuda() + diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=200, schedule='linear', + enable_conditioning_free_guidance=False, + conditioning_free_k=1) + diffusion_model = TransformerDiffusionTTS(model_channels=896, num_layers=16, in_channels=100, in_latent_channels=1024, + token_count=256, out_channels=200, dropout=0, unconditioned_percentage=0) + diffusion_model.load_state_dict(torch.load('X:\\dlas\\experiments\\train_speech_diffusion_from_ctc_tfd5\\models\\26500_generator_ema.pth')) + diffusion_model = diffusion_model.cuda().eval() + with torch.no_grad(): + text_codes = torch.LongTensor(bpe_tokenizer.encode(text)).unsqueeze(0).to(codes.device) + clvp_latent = clvp.embed_text(text_codes) + cond_mel = TorchMelSpectrogramInjector({'n_mel_channels': 100, 'mel_fmax': 11000, 'filter_length': 8000, 'normalize': True, + 'true_normalization': True, 'in': 'in', 'out': 'out'}, {})({'in': cond})['out'] + gen = diffuser.p_sample_loop(diffusion_model, (1,100,int(codes.shape[-1]*RATIO)), model_kwargs={'codes': codes, + 'conditioning_input': cond_mel, + 'type': torch.tensor([0], device=codes.device), + 'clvp_input': clvp_latent}) + gen_denorm = denormalize_mel(gen) + vocoder = load_univnet_vocoder().cuda() + gen_wav = vocoder.inference(gen_denorm) + return gen_wav + + if __name__ == '__main__': model = CtcCodeGenerator(model_dim=512, layers=16, dropout=0).eval().cuda() model.load_state_dict(torch.load('../experiments/train_encoder_build_ctc_alignments_toy/models/76000_generator_ema.pth')) tokenizer = Wav2Vec2CTCTokenizer.from_pretrained('jbetker/tacotron-symbols') - text = "and now, what do you want." - seq = [0, 0, 0, 38, 51, 51, 41, 11, 11, 51, 51, 0, 0, 0, 0, 52, 0, 60, 0, 0, 0, 0, 0, 0, 6, 11, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 60, 45, 0, 38, 57, 57, 11, 0, 41, 52, 52, 11, 11, 62, 52, 52, 58, 0, 11, 11, 60, 0, 0, 0, 0, 38, 0, 0, 51, 51, 0, 0, 57, 0, 0, 7, 7, 0, 0, 0] - codes, pads, repeats = get_ctc_metadata(seq) + text = "Can I have tea and a pot of butter, please?" + #seq = [0, 0, 0, 38, 51, 51, 41, 11, 11, 51, 51, 0, 0, 0, 0, 52, 0, 60, 0, 0, 0, 0, 0, 0, 6, 11, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 60, 45, 0, 38, 57, 57, 11, 0, 41, 52, 52, 11, 11, 62, 52, 52, 58, 0, 11, 11, 60, 0, 0, 0, 0, 38, 0, 0, 51, 51, 0, 0, 57, 0, 0, 7, 7, 0, 0, 0] + #codes, pads, repeats = get_ctc_metadata(seq) + codes = tokenizer.encode(text) with torch.no_grad(): - codes = codes.cuda().unsqueeze(0) - pads = pads.cuda().unsqueeze(0) - repeats = repeats.cuda().unsqueeze(0) - - ppads = pads.clone() - prepeats = repeats.clone() - mask = torch.zeros_like(pads) - conf_str = tokenizer.decode(codes[0].tolist()) + codes = torch.tensor(codes).cuda().unsqueeze(0) + ppads = torch.zeros_like(codes) + prepeats = torch.zeros_like(codes) + mask = torch.zeros_like(codes) for s in range(codes.shape[-1]): - logits, confidences = model.inference(codes, pads * mask, repeats * mask) + logits, confidences = model.inference(codes, ppads * mask, prepeats * mask) confidences = confidences * mask.logical_not() # prevent prediction of tokens that have already been predicted. i = confidences.argmax(dim=-1) @@ -59,5 +100,9 @@ if __name__ == '__main__': prepeats[0,i] = pred_repeats mask[0,i] = 1 - conf_str = conf_str[:i] + conf_str[i].upper() + conf_str[i+1:] - print(f"conf: {conf_str} pads={pred_pads}:{pads[0,i].item()} repeats={pred_repeats}:{repeats[0,i].item()}") \ No newline at end of file + #print(f"conf: {conf_str} pads={pred_pads}:{pads[0,i].item()} repeats={pred_repeats}:{repeats[0,i].item()}") + + decoded_codes = decode_ctc_metadata(codes[0], ppads[0], prepeats[0]).unsqueeze(0) + cond = load_audio('D:\\tortoise-tts\\tortoise\\voices\\train_dotrice\\1.wav', 22050).unsqueeze(0).cuda() + decoded_wav = diffuse(text, decoded_codes, cond) + torchaudio.save('output.wav', decoded_wav.cpu()[0], 24000) \ No newline at end of file diff --git a/codes/trainer/eval/music_diffusion_fid.py b/codes/trainer/eval/music_diffusion_fid.py index d8ea1e0e..35abc7b6 100644 --- a/codes/trainer/eval/music_diffusion_fid.py +++ b/codes/trainer/eval/music_diffusion_fid.py @@ -62,6 +62,8 @@ class MusicDiffusionFid(evaluator.Evaluator): self.local_modules['codegen'] = get_music_codegen() elif 'from_codes_quant' == mode: self.diffusion_fn = self.perform_diffusion_from_codes_quant + elif 'from_codes_quant_gradual_decode' == mode: + self.diffusion_fn = self.perform_diffusion_from_codes_quant_gradual_decode self.spec_fn = TorchMelSpectrogramInjector({'n_mel_channels': 256, 'mel_fmax': 11000, 'filter_length': 16000, 'normalize': True, 'in': 'in', 'out': 'out'}, {}) @@ -133,6 +135,38 @@ class MusicDiffusionFid(evaluator.Evaluator): return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate + 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 + audio = audio.unsqueeze(0) + + mel = self.spec_fn({'in': audio})['out'] + mel_norm = normalize_mel(mel) + guidance = torch.zeros_like(mel_norm) + mask = torch.zeros_like(mel_norm) + GRADS = 4 + for k in range(GRADS): + gen_mel = self.diffuser.p_sample_loop_with_guidance(self.model, + guidance_input=guidance, mask=mask, + model_kwargs={'truth_mel': mel, + 'conditioning_input': torch.zeros_like(mel_norm[:,:,:390]), + 'disable_diversity': True}) + pk = int(k*(mel_norm.shape[1]/GRADS)) + ek = int((k+1)*(mel_norm.shape[1]/GRADS)) + guidance[:, pk:ek] = gen_mel[:, pk:ek] + mask[:, :ek] = 1 + + 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.diffuser.p_sample_loop(self.spec_decoder, output_shape, + model_kwargs={'aligned_conditioning': gen_mel_denorm}) + gen_wav = pixel_shuffle_1d(gen_wav, 16) + + return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate + def project(self, sample, sample_rate): sample = torchaudio.functional.resample(sample, sample_rate, 22050) @@ -201,13 +235,15 @@ class MusicDiffusionFid(evaluator.Evaluator): if __name__ == '__main__': - diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_diffusion_tfd_quant7.yml', 'generator', + diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_diffusion_ar_prior.yml', 'generator', also_load_savepoint=False, - load_path='X:\\dlas\\experiments\\train_music_diffusion_unet_music\\models\\46500_generator_ema.pth' + load_path='X:\\dlas\\experiments\\train_music_diffusion_ar_prior\\models\\22000_generator_ema.pth' ).cuda() - opt_eval = {'path': 'Y:\\split\\yt-music-eval', 'diffusion_steps': 100, - 'conditioning_free': True, 'conditioning_free_k': 1, + opt_eval = {#'path': 'Y:\\split\\yt-music-eval', + 'path': 'E:\\music_eval', + 'diffusion_steps': 100, + 'conditioning_free': False, 'conditioning_free_k': 1, 'diffusion_schedule': 'linear', 'diffusion_type': 'from_codes_quant'} - env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 561, 'device': 'cuda', 'opt': {}} + env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 502, 'device': 'cuda', 'opt': {}} eval = MusicDiffusionFid(diffusion, opt_eval, env) print(eval.perform_eval())