diff --git a/api.py b/api.py
index 6c3fb1e..7c33484 100644
--- a/api.py
+++ b/api.py
@@ -117,13 +117,14 @@ def do_spectrogram_diffusion(diffusion_model, diffuser, mel_codes, conditioning_
cond_mels.append(cond_mel)
cond_mels = torch.stack(cond_mels, dim=1)
- output_shape = (mel_codes.shape[0], 100, mel_codes.shape[-1]*4)
- precomputed_embeddings = diffusion_model.timestep_independent(mel_codes, cond_mels, False)
+ output_seq_len = mel_codes.shape[-1]*4*24000//22050 # This diffusion model converts from 22kHz spectrogram codes to a 24kHz spectrogram signal.
+ output_shape = (mel_codes.shape[0], 100, output_seq_len)
+ precomputed_embeddings = diffusion_model.timestep_independent(mel_codes, cond_mels, output_seq_len, False)
noise = torch.randn(output_shape, device=mel_codes.device) * temperature
mel = diffuser.p_sample_loop(diffusion_model, output_shape, noise=noise,
model_kwargs={'precomputed_aligned_embeddings': precomputed_embeddings})
- return denormalize_tacotron_mel(mel)[:,:,:mel_codes.shape[-1]*4]
+ return denormalize_tacotron_mel(mel)[:,:,:output_seq_len]
class TextToSpeech:
diff --git a/api_new_autoregressive.py b/api_new_autoregressive.py
new file mode 100644
index 0000000..7a9d5ce
--- /dev/null
+++ b/api_new_autoregressive.py
@@ -0,0 +1,245 @@
+import argparse
+import os
+import random
+from urllib import request
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+import progressbar
+import ocotillo
+
+from models.diffusion_decoder import DiffusionTts
+from models.autoregressive import UnifiedVoice
+from tqdm import tqdm
+
+from models.arch_util import TorchMelSpectrogram
+from models.new_autoregressive import AutoregressiveCodegen
+from models.text_voice_clip import VoiceCLIP
+from models.vocoder import UnivNetGenerator
+from utils.audio import load_audio, wav_to_univnet_mel, denormalize_tacotron_mel
+from utils.diffusion import SpacedDiffusion, space_timesteps, get_named_beta_schedule
+from utils.tokenizer import VoiceBpeTokenizer, lev_distance
+
+
+pbar = None
+def download_models():
+ MODELS = {
+ 'clip.pth': 'https://huggingface.co/jbetker/tortoise-tts-clip/resolve/main/pytorch-model.bin',
+ 'diffusion.pth': 'https://huggingface.co/jbetker/tortoise-tts-diffusion-v1/resolve/main/pytorch-model.bin',
+ 'autoregressive.pth': 'https://huggingface.co/jbetker/tortoise-tts-autoregressive/resolve/main/pytorch-model.bin'
+ }
+ os.makedirs('.models', exist_ok=True)
+ def show_progress(block_num, block_size, total_size):
+ global pbar
+ if pbar is None:
+ pbar = progressbar.ProgressBar(maxval=total_size)
+ pbar.start()
+
+ downloaded = block_num * block_size
+ if downloaded < total_size:
+ pbar.update(downloaded)
+ else:
+ pbar.finish()
+ pbar = None
+ for model_name, url in MODELS.items():
+ if os.path.exists(f'.models/{model_name}'):
+ continue
+ print(f'Downloading {model_name} from {url}...')
+ request.urlretrieve(url, f'.models/{model_name}', show_progress)
+ print('Done.')
+
+
+def pad_or_truncate(t, length):
+ if t.shape[-1] == length:
+ return t
+ elif t.shape[-1] < length:
+ return F.pad(t, (0, length-t.shape[-1]))
+ else:
+ return t[..., :length]
+
+
+def load_discrete_vocoder_diffuser(trained_diffusion_steps=4000, desired_diffusion_steps=200, cond_free=True, cond_free_k=1):
+ """
+ Helper function to load a GaussianDiffusion instance configured for use as a vocoder.
+ """
+ return SpacedDiffusion(use_timesteps=space_timesteps(trained_diffusion_steps, [desired_diffusion_steps]), model_mean_type='epsilon',
+ model_var_type='learned_range', loss_type='mse', betas=get_named_beta_schedule('linear', trained_diffusion_steps),
+ conditioning_free=cond_free, conditioning_free_k=cond_free_k)
+
+
+def load_conditioning(clip, cond_length=132300):
+ gap = clip.shape[-1] - cond_length
+ if gap < 0:
+ clip = F.pad(clip, pad=(0, abs(gap)))
+ elif gap > 0:
+ rand_start = random.randint(0, gap)
+ clip = clip[:, rand_start:rand_start + cond_length]
+ mel_clip = TorchMelSpectrogram()(clip.unsqueeze(0)).squeeze(0)
+ return mel_clip.unsqueeze(0).cuda()
+
+
+def fix_autoregressive_output(codes, stop_token):
+ """
+ This function performs some padding on coded audio that fixes a mismatch issue between what the diffusion model was
+ trained on and what the autoregressive code generator creates (which has no padding or end).
+ This is highly specific to the DVAE being used, so this particular coding will not necessarily work if used with
+ a different DVAE. This can be inferred by feeding a audio clip padded with lots of zeros on the end through the DVAE
+ and copying out the last few codes.
+
+ Failing to do this padding will produce speech with a harsh end that sounds like "BLAH" or similar.
+ """
+ # Strip off the autoregressive stop token and add padding.
+ stop_token_indices = (codes == stop_token).nonzero()
+ if len(stop_token_indices) == 0:
+ print("No stop tokens found, enjoy that output of yours!")
+ return codes
+ else:
+ codes[stop_token_indices] = 83
+ stm = stop_token_indices.min().item()
+ codes[stm:] = 83
+ if stm - 3 < codes.shape[0]:
+ codes[-3] = 45
+ codes[-2] = 45
+ codes[-1] = 248
+
+ return codes
+
+
+def do_spectrogram_diffusion(diffusion_model, diffuser, mel_codes, conditioning_samples, temperature=1):
+ """
+ Uses the specified diffusion model to convert discrete codes into a spectrogram.
+ """
+ with torch.no_grad():
+ cond_mels = []
+ for sample in conditioning_samples:
+ sample = pad_or_truncate(sample, 102400)
+ cond_mel = wav_to_univnet_mel(sample.to(mel_codes.device), do_normalization=False)
+ cond_mels.append(cond_mel)
+ cond_mels = torch.stack(cond_mels, dim=1)
+
+ output_seq_len = mel_codes.shape[-1]*4*24000//22050 # This diffusion model converts from 22kHz spectrogram codes to a 24kHz spectrogram signal.
+ output_shape = (mel_codes.shape[0], 100, output_seq_len)
+ precomputed_embeddings = diffusion_model.timestep_independent(mel_codes, cond_mels, output_seq_len, False)
+
+ noise = torch.randn(output_shape, device=mel_codes.device) * temperature
+ mel = diffuser.p_sample_loop(diffusion_model, output_shape, noise=noise,
+ model_kwargs={'precomputed_aligned_embeddings': precomputed_embeddings})
+ return denormalize_tacotron_mel(mel)[:,:,:output_seq_len]
+
+
+class TextToSpeech:
+ def __init__(self, autoregressive_batch_size=32):
+ self.autoregressive_batch_size = autoregressive_batch_size
+ self.tokenizer = VoiceBpeTokenizer()
+ download_models()
+
+ self.autoregressive = AutoregressiveCodegen(512, 12).cpu().eval()
+ self.autoregressive.load_state_dict(torch.load('D:\\dlas\\experiments\\train_autoregressive_codegen\\models\\23000_codegen_ema.pth'))
+
+ self.clip = VoiceCLIP(dim_text=512, dim_speech=512, dim_latent=512, num_text_tokens=256, text_enc_depth=12,
+ text_seq_len=350, text_heads=8,
+ num_speech_tokens=8192, speech_enc_depth=12, speech_heads=8, speech_seq_len=430,
+ use_xformers=True).cpu().eval()
+ self.clip.load_state_dict(torch.load('.models/clip.pth'))
+
+ self.diffusion = DiffusionTts(model_channels=1024, num_layers=10, in_channels=100, out_channels=200,
+ in_latent_channels=1024, in_tokens=8193, dropout=0, use_fp16=False, num_heads=16,
+ layer_drop=0, unconditioned_percentage=0).cpu().eval()
+ self.diffusion.load_state_dict(torch.load('.models/diffusion.pth'))
+
+ self.vocoder = UnivNetGenerator().cpu()
+ self.vocoder.load_state_dict(torch.load('.models/vocoder.pth')['model_g'])
+ self.vocoder.eval(inference=True)
+
+ def tts(self, text, voice_samples, k=1,
+ # autoregressive generation parameters follow
+ num_autoregressive_samples=512, temperature=.5, length_penalty=2, repetition_penalty=2.0, top_p=.5,
+ typical_sampling=False, typical_mass=.9,
+ # diffusion generation parameters follow
+ diffusion_iterations=100, cond_free=True, cond_free_k=2, diffusion_temperature=.7,):
+ text = torch.IntTensor(self.tokenizer.encode(text)).unsqueeze(0).cuda()
+ text = F.pad(text, (0, 1)) # This may not be necessary.
+
+ conds = []
+ if not isinstance(voice_samples, list):
+ voice_samples = [voice_samples]
+ for vs in voice_samples:
+ conds.append(load_conditioning(vs))
+ conds = torch.stack(conds, dim=1)
+
+ diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=diffusion_iterations, cond_free=cond_free, cond_free_k=cond_free_k)
+
+ with torch.no_grad():
+ samples = []
+ num_batches = num_autoregressive_samples // self.autoregressive_batch_size
+ stop_mel_token = self.autoregressive.STOP_TOKEN
+ self.autoregressive = self.autoregressive.cuda()
+ for _ in tqdm(range(num_batches)):
+ codes = self.autoregressive.generate(conds, text,
+ do_sample=True,
+ top_p=top_p,
+ temperature=temperature,
+ num_return_sequences=self.autoregressive_batch_size,
+ length_penalty=length_penalty,
+ repetition_penalty=repetition_penalty,
+ typical_sampling=typical_sampling,
+ typical_mass=typical_mass)
+ padding_needed = 250 - codes.shape[1]
+ codes = F.pad(codes, (0, padding_needed), value=stop_mel_token)
+ samples.append(codes)
+ #self.autoregressive = self.autoregressive.cpu()
+
+ clip_results = []
+ self.clip = self.clip.cuda()
+ for batch in samples:
+ for i in range(batch.shape[0]):
+ batch[i] = fix_autoregressive_output(batch[i], stop_mel_token)
+ bad_toks = batch >= 8192
+ batch = batch * bad_toks.logical_not()
+ clip_results.append(self.clip(text.repeat(batch.shape[0], 1), batch, return_loss=False))
+ clip_results = torch.cat(clip_results, dim=0)
+ samples = torch.cat(samples, dim=0)
+ best_results = samples[torch.topk(clip_results, k=k).indices]
+ self.clip = self.clip.cpu()
+ del samples
+
+ print("Performing vocoding..")
+ wav_candidates = []
+ self.diffusion = self.diffusion.cuda()
+ self.vocoder = self.vocoder.cuda()
+ for b in range(best_results.shape[0]):
+ code = best_results[b].unsqueeze(0)
+ mel = do_spectrogram_diffusion(self.diffusion, diffuser, code, voice_samples, temperature=diffusion_temperature)
+ wav = self.vocoder.inference(mel)
+ wav_candidates.append(wav.cpu())
+ self.diffusion = self.diffusion.cpu()
+ self.vocoder = self.vocoder.cpu()
+
+ if len(wav_candidates) > 1:
+ return wav_candidates
+ return wav_candidates[0]
+
+ def refine_for_intellibility(self, wav_candidates, corresponding_codes, output_path):
+ """
+ Further refine the remaining candidates using a ASR model to pick out the ones that are the most understandable.
+ TODO: finish this function
+ :param wav_candidates:
+ :return:
+ """
+ transcriber = ocotillo.Transcriber(on_cuda=True)
+ transcriptions = transcriber.transcribe_batch(torch.cat(wav_candidates, dim=0).squeeze(1), 24000)
+ best = 99999999
+ for i, transcription in enumerate(transcriptions):
+ dist = lev_distance(transcription, args.text.lower())
+ if dist < best:
+ best = dist
+ best_codes = corresponding_codes[i].unsqueeze(0)
+ best_wav = wav_candidates[i]
+ del transcriber
+ torchaudio.save(os.path.join(output_path, f'{voice}_poor.wav'), best_wav.squeeze(0).cpu(), 24000)
+
+ # Perform diffusion again with the high-quality diffuser.
+ mel = do_spectrogram_diffusion(diffusion, final_diffuser, best_codes, cond_diffusion, mean=False)
+ wav = vocoder.inference(mel)
+ torchaudio.save(os.path.join(args.output_path, f'{voice}.wav'), wav.squeeze(0).cpu(), 24000)
\ No newline at end of file
diff --git a/do_tts.py b/do_tts.py
index af5c780..e48e9d5 100644
--- a/do_tts.py
+++ b/do_tts.py
@@ -5,7 +5,7 @@ import torch
import torch.nn.functional as F
import torchaudio
-from api import TextToSpeech, load_conditioning
+from api_new_autoregressive import TextToSpeech, load_conditioning
from utils.audio import load_audio
from utils.tokenizer import VoiceBpeTokenizer
@@ -28,7 +28,7 @@ if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-text', type=str, help='Text to speak.', default="I am a language model that has learned to speak.")
parser.add_argument('-voice', type=str, help='Use a preset conditioning voice (defined above). Overrides cond_path.', default='dotrice,harris,lescault,otto,atkins,grace,kennard,mol')
- parser.add_argument('-num_samples', type=int, help='How many total outputs the autoregressive transformer should produce.', default=512)
+ parser.add_argument('-num_samples', type=int, help='How many total outputs the autoregressive transformer should produce.', default=32)
parser.add_argument('-batch_size', type=int, help='How many samples to process at once in the autoregressive model.', default=16)
parser.add_argument('-num_diffusion_samples', type=int, help='Number of outputs that progress to the diffusion stage.', default=16)
parser.add_argument('-output_path', type=str, help='Where to store outputs.', default='results/')
diff --git a/models/diffusion_decoder.py b/models/diffusion_decoder.py
index cacdfc1..1baf809 100644
--- a/models/diffusion_decoder.py
+++ b/models/diffusion_decoder.py
@@ -212,7 +212,7 @@ class DiffusionTts(nn.Module):
}
return groups
- def timestep_independent(self, aligned_conditioning, conditioning_input, return_code_pred):
+ def timestep_independent(self, aligned_conditioning, conditioning_input, expected_seq_len, return_code_pred):
# Shuffle aligned_latent to BxCxS format
if is_latent(aligned_conditioning):
aligned_conditioning = aligned_conditioning.permute(0, 2, 1)
@@ -227,7 +227,7 @@ class DiffusionTts(nn.Module):
cond_emb = conds.mean(dim=-1)
cond_scale, cond_shift = torch.chunk(cond_emb, 2, dim=1)
if is_latent(aligned_conditioning):
- code_emb = self.latent_converter(aligned_conditioning)
+ code_emb = self.autoregressive_latent_converter(aligned_conditioning)
else:
code_emb = self.code_embedding(aligned_conditioning).permute(0, 2, 1)
code_emb = self.code_converter(code_emb)
@@ -240,7 +240,7 @@ class DiffusionTts(nn.Module):
device=code_emb.device) < self.unconditioned_percentage
code_emb = torch.where(unconditioned_batches, self.unconditioned_embedding.repeat(aligned_conditioning.shape[0], 1, 1),
code_emb)
- expanded_code_emb = F.interpolate(code_emb, size=aligned_conditioning.shape[-1]*4, mode='nearest')
+ expanded_code_emb = F.interpolate(code_emb, size=expected_seq_len, mode='nearest')
if not return_code_pred:
return expanded_code_emb
@@ -250,7 +250,6 @@ class DiffusionTts(nn.Module):
mel_pred = mel_pred * unconditioned_batches.logical_not()
return expanded_code_emb, mel_pred
-
def forward(self, x, timesteps, aligned_conditioning=None, conditioning_input=None, precomputed_aligned_embeddings=None, conditioning_free=False, return_code_pred=False):
"""
Apply the model to an input batch.
@@ -275,11 +274,12 @@ class DiffusionTts(nn.Module):
if precomputed_aligned_embeddings is not None:
code_emb = precomputed_aligned_embeddings
else:
- code_emb, mel_pred = self.timestep_independent(aligned_conditioning, conditioning_input, True)
+ code_emb, mel_pred = self.timestep_independent(aligned_conditioning, conditioning_input, x.shape[-1], True)
if is_latent(aligned_conditioning):
unused_params.extend(list(self.code_converter.parameters()) + list(self.code_embedding.parameters()))
else:
unused_params.extend(list(self.latent_converter.parameters()))
+
unused_params.append(self.unconditioned_embedding)
time_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
diff --git a/models/new_autoregressive.py b/models/new_autoregressive.py
new file mode 100644
index 0000000..a6d8dee
--- /dev/null
+++ b/models/new_autoregressive.py
@@ -0,0 +1,293 @@
+import functools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPT2PreTrainedModel, GPT2Config
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from x_transformers import TransformerWrapper, Encoder, Decoder
+
+from models.arch_util import AttentionBlock
+
+
+class InferenceModel(GPT2PreTrainedModel):
+ """
+ Implementation of GPT2PreTrainedModel from transformers, which allows us to use their generation library with
+ this transformer.
+ """
+ def __init__(self, model):
+ super().__init__(GPT2Config())
+ self.transformer = model
+ self.context = None
+
+ def parallelize(self, device_map=None):
+ # Not implemented.
+ pass
+
+ def deparallelize(self):
+ # Not implemented.
+ pass
+
+ def get_output_embeddings(self):
+ assert False, "Unsupported operation."
+
+ def set_output_embeddings(self, new_embeddings):
+ assert False, "Unsupported operation."
+
+ def store_context(self, context):
+ self.context = context
+
+ def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+ token_type_ids = kwargs.get("token_type_ids", None)
+ # only last token for inputs_ids if past is defined in kwargs
+ if past:
+ input_ids = input_ids[:, -1].unsqueeze(-1)
+ if token_type_ids is not None:
+ token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+ attention_mask = kwargs.get("attention_mask", None)
+ position_ids = kwargs.get("position_ids", None)
+
+ if attention_mask is not None and position_ids is None:
+ # create position_ids on the fly for batch generation
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ if past:
+ position_ids = position_ids[:, -1].unsqueeze(-1)
+ else:
+ position_ids = None
+ return {
+ "input_ids": input_ids,
+ "past_key_values": past,
+ "use_cache": kwargs.get("use_cache"),
+ "position_ids": position_ids,
+ "attention_mask": attention_mask,
+ "token_type_ids": token_type_ids,
+ }
+
+ def forward(
+ self,
+ input_ids=None,
+ past_key_values=None,
+ attention_mask=None,
+ token_type_ids=None,
+ position_ids=None,
+ head_mask=None,
+ inputs_embeds=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ labels=None,
+ use_cache=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ assert self.context is not None
+ assert inputs_embeds is None # Not supported by this inference model.
+ assert labels is None # Training not supported by this inference model.
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ hidden_states = self.transformer.decoder(input_ids, context=self.context, return_embeddings=True)
+ logits = self.transformer.decoder.transformer.to_logits(hidden_states)
+
+ if not return_dict:
+ return (logits, )
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=None,
+ logits=logits,
+ past_key_values=None,
+ hidden_states=hidden_states,
+ attentions=None,
+ cross_attentions=None,
+ )
+
+ @staticmethod
+ def _reorder_cache(past, beam_idx):
+ """
+ This function is used to re-order the :obj:`past_key_values` cache if
+ :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+ called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+ """
+ return tuple(
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
+ for layer_past in past
+ )
+
+
+class ResBlock(nn.Module):
+ """
+ Basic residual convolutional block that uses GroupNorm.
+ """
+ def __init__(self, chan):
+ super().__init__()
+ self.net = nn.Sequential(
+ nn.Conv1d(chan, chan, kernel_size=3, padding=1),
+ nn.GroupNorm(chan//8, chan),
+ nn.ReLU(),
+ nn.Conv1d(chan, chan, kernel_size=3, padding=1),
+ nn.GroupNorm(chan//8, chan)
+ )
+
+ def forward(self, x):
+ return F.relu(self.net(x) + x)
+
+
+class ConditioningEncoder(nn.Module):
+ def __init__(self,
+ spec_dim,
+ embedding_dim,
+ attn_blocks=6,
+ num_attn_heads=4,
+ do_checkpointing=False):
+ super().__init__()
+ attn = []
+ self.init = nn.Sequential(nn.Conv1d(spec_dim, embedding_dim//4, kernel_size=5, padding=2),
+ nn.Conv1d(embedding_dim//4, embedding_dim//2, kernel_size=3, padding=1, stride=2),
+ ResBlock(embedding_dim//2),
+ nn.Conv1d(embedding_dim//2, embedding_dim, kernel_size=3, padding=1, stride=2))
+ for a in range(attn_blocks):
+ attn.append(AttentionBlock(embedding_dim, num_attn_heads, do_checkpoint=do_checkpointing))
+ self.attn = nn.Sequential(*attn)
+ self.dim = embedding_dim
+
+ def forward(self, x):
+ h = self.init(x)
+ h = self.attn(h)
+ return h.mean(dim=2)
+
+
+class CheckpointedLayer(nn.Module):
+ """
+ Wraps a module. When forward() is called, passes kwargs that require_grad through torch.checkpoint() and bypasses
+ checkpoint for all other args.
+ """
+ def __init__(self, wrap):
+ super().__init__()
+ self.wrap = wrap
+
+ def forward(self, x, *args, **kwargs):
+ for k, v in kwargs.items():
+ assert not (isinstance(v, torch.Tensor) and v.requires_grad) # This would screw up checkpointing.
+ partial = functools.partial(self.wrap, **kwargs)
+ return torch.utils.checkpoint.checkpoint(partial, x, *args)
+
+
+class CheckpointedXTransformerWrapper(nn.Module):
+ """
+ Wraps a TransformerWrapper and applies CheckpointedLayer to each layer.
+ """
+ def __init__(self, checkpoint=True, **xtransformer_kwargs):
+ super().__init__()
+ self.transformer = TransformerWrapper(**xtransformer_kwargs)
+
+ if not checkpoint:
+ return
+ for i in range(len(self.transformer.attn_layers.layers)):
+ n, b, r = self.transformer.attn_layers.layers[i]
+ self.transformer.attn_layers.layers[i] = nn.ModuleList([n, CheckpointedLayer(b), r])
+
+ def forward(self, x, **kwargs):
+ return self.transformer(x, **kwargs)
+
+
+class AutoregressiveCodegen(nn.Module):
+ def __init__(self, model_dim, depth, num_text_tokens=256, num_mel_tokens=8194, max_text_tokens=4000,
+ max_mel_tokens=4000, dropout=.1):
+ super().__init__()
+
+ self.START_TOKEN=8192
+ self.STOP_TOKEN=8193
+ self.max_mel_tokens = max_mel_tokens
+ self.minicoder = ConditioningEncoder(80, model_dim, do_checkpointing=False)
+ self.encoder = CheckpointedXTransformerWrapper(
+ num_tokens=num_text_tokens,
+ max_seq_len=max_text_tokens,
+ attn_layers = Encoder(
+ depth=depth//2,
+ heads=model_dim//64,
+ dim=model_dim,
+ attn_dropout=dropout,
+ ff_dropout=dropout,
+ use_rmsnorm=True,
+ ff_glu=True,
+ ff_mult=1,
+ rotary_pos_emb=True,
+ rel_pos_bias=True,
+ ))
+ self.decoder = CheckpointedXTransformerWrapper(
+ num_tokens=num_mel_tokens,
+ max_seq_len=max_mel_tokens,
+ attn_layers=Decoder(
+ depth=depth,
+ heads=model_dim//64,
+ dim=model_dim,
+ attn_dropout=dropout,
+ ff_dropout=dropout,
+ use_rmsnorm=True,
+ ff_glu=True,
+ ff_mult=1,
+ rotary_pos_emb=True,
+ rel_pos_bias=True,
+ cross_attend=True,
+ ))
+
+ def get_grad_norm_parameter_groups(self):
+ return {
+ 'encoder': list(self.encoder.parameters()),
+ 'decoder': list(self.decoder.parameters()),
+ 'minicoder': list(self.minicoder.parameters()),
+ }
+
+ def forward(self, text_codes, conditioning_signal, mel_codes, wav_lengths, return_loss=True):
+ # Format mel_codes with a stop token on the end.
+ mel_lengths = wav_lengths // 1024 + 1
+ for b in range(mel_codes.shape[0]):
+ mel_codes[b, mel_lengths[b]:] = self.STOP_TOKEN
+ mel_codes = F.pad(mel_codes, (0, 1), value=self.STOP_TOKEN)
+
+ # Build the context
+ if len(conditioning_signal.shape) != 4:
+ conditioning_signal = conditioning_signal.unsqueeze(1)
+ cond_embs = []
+ for i in range(conditioning_signal.shape[1]):
+ cond_embs.append(self.minicoder(conditioning_signal[:, i]))
+ cond_emb = torch.stack(cond_embs, dim=1).mean(dim=1, keepdim=True)
+ enc_text = self.encoder(text_codes, return_embeddings=True)
+ context = torch.cat([cond_emb, enc_text], dim=1)
+
+ # Execute the decoder
+ dec_inputs = F.pad(mel_codes, (1,0), value=self.START_TOKEN)[:, :-1]
+ dec = self.decoder(dec_inputs, context=context)
+ if not return_loss:
+ return dec
+ loss_mel = F.cross_entropy(dec.permute(0,2,1), mel_codes)
+ return loss_mel
+
+ def generate(self, conditioning_signal, text_codes, **hf_generate_kwargs):
+ if not hasattr(self, 'inference_model'):
+ self.inference_model = InferenceModel(self)
+
+ if len(conditioning_signal.shape) != 4:
+ conditioning_signal = conditioning_signal.unsqueeze(1)
+ cond_embs = []
+ for i in range(conditioning_signal.shape[1]):
+ cond_embs.append(self.minicoder(conditioning_signal[:, i]))
+ cond_emb = torch.stack(cond_embs, dim=1).mean(dim=1, keepdim=True)
+ enc_text = self.encoder(text_codes, return_embeddings=True)
+ context = torch.cat([cond_emb, enc_text], dim=1)
+ self.inference_model.store_context(context)
+
+ gen = self.inference_model.generate(bos_token_id=self.START_TOKEN, pad_token_id=self.STOP_TOKEN, eos_token_id=self.STOP_TOKEN,
+ max_length=250, output_attentions=False, return_dict_in_generate=True,
+ **hf_generate_kwargs)
+ return gen.sequences
+
+
+if __name__ == '__main__':
+ codegen = AutoregressiveCodegen(1024, 20)
+ codegen.generate(torch.randn((1,80,120)), torch.randint(0,256,(1,200)))
+ codegen(torch.randint(0,256, (2,200)),
+ torch.randn(2,80,120),
+ torch.randint(0,8192, (2,350)),
+ torch.tensor([192,350]))
\ No newline at end of file