Added a bundle of 3 trained A2A codecs to enhance the sound quality (…

…NFY)

spear_tts_pytorch/a2a.py

@@ -0,0 +1,151 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/6. EnCodec quality enchancement.ipynb.
+
+# %% auto 0
+__all__ = ['load_data', 'load_datasets', 'LayerNorm', 'init_transformer', 'PureEncoder', 'AAARTransformer', 'make_model']
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 1
+import io
+import time
+import random
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 2
+import torch
+import torch.nn as nn
+from torch.profiler import profile, record_function, ProfilerActivity, schedule
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 3
+from pathlib import Path
+import json
+from fastprogress import progress_bar, master_bar
+import pandas as pd
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 4
+from spear_tts_pytorch.train import *
+from spear_tts_pytorch.modules import *
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 12
+def load_data(path):
+ atoks = []
+ for name in Path(path).glob('*.encodec'):
+ atoks.append(name)
+ return pd.DataFrame(dict(atoks=atoks))
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 14
+import torch.nn.functional as F
+
+class SADataset(torch.utils.data.Dataset):
+ def __init__(self, data, fromq=2, toq=4):
+ self.fromq = fromq
+ self.toq = toq
+ self.n_ctx = 192
+ self.data = data
+ self.samples = [(i,j) for i,name in enumerate(data['atoks']) for j in range(torch.load(name).shape[-1] // self.n_ctx)]
+
+ def __len__(self):
+ return len(self.samples)
+
+ def A_tokens(self):
+ return len(self)*self.n_ctx*4
+
+ def hours(self):
+ return len(self)*self.n_ctx/2250*30/3600
+
+ def __repr__(self):
+ return f"Dataset<{len(self)} samples, ({self.fromq}->{self.toq}), {self.A_tokens()} Atokens, {self.hours():.1f} hours>"
+
+ def __getitem__(self, idx):
+ i,j = self.samples[idx]
+ row = self.data.iloc[i]
+ jA = j * self.n_ctx
+ Atoks = torch.load(row['atoks'], map_location='cpu')[0,:self.toq,jA:jA+self.n_ctx].T
+ outtoks = Atoks.reshape(-1).clone()
+ Atoks[:,self.fromq:] = 1024 # mask token
+ intoks = Atoks.reshape(-1).clone()
+ return intoks, outtoks
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 20
+def load_datasets(
+ path:Path, # encodec files path
+ subsample:float=1, # use a fraction of the files
+ fromq:int=2, # input quantizers
+ toq:int=8 # output quantizers
+ ):
+ data = load_data(path)
+
+ val_data, train_data = data[:12], data[12:int(len(data)*subsample)]
+
+ return SADataset(train_data, fromq=fromq, toq=toq), SADataset(val_data, fromq=fromq, toq=toq)
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 23
+class LayerNorm(nn.LayerNorm):
+ def forward(self, x):
+ return super().forward(x.float()).type(x.dtype)
+
+# based on https://github.com/karpathy/minGPT/blob/master/mingpt/model.py#L163
+def init_transformer(m):
+ if isinstance(m, (nn.Linear, nn.Embedding)):
+ torch.nn.init.trunc_normal_(m.weight, std=.02)
+ if isinstance(m, nn.Linear) and m.bias is not None:
+ torch.nn.init.constant_(m.bias, 0)
+ elif isinstance(m, nn.LayerNorm):
+ torch.nn.init.constant_(m.bias, 0)
+ torch.nn.init.constant_(m.weight, 1.0)
+
+class PureEncoder(nn.Module):
+ def __init__(self, depth=6, width=384, n_head=6, length=1500, codes=1024, pos_embs=None):
+ super().__init__()
+ self.length = length
+ self.codes = codes
+
+ # embed semantic tokens
+ self.embedding = nn.Embedding(codes+1, width)
+ if pos_embs is None: pos_embs = sinusoids(length, width)
+ self.register_buffer("positional_embedding", pos_embs)
+
+ self.layers = nn.ModuleList([
+ ResidualAttentionBlock(width, n_head) for _ in range(depth)
+ ])
+ self.ln_post = LayerNorm(width)
+
+ self.apply(init_transformer)
+
+ def forward(self, Stoks):
+ Sembs = self.embedding(Stoks)
+
+ xin = (Sembs + self.positional_embedding[:Sembs.shape[1]])
+
+ x = xin
+ for l in self.layers: x = l(x, causal=False)
+
+ x = self.ln_post(x)
+
+ logits = (x @ self.embedding.weight.to(x.dtype).T).float()
+ return logits
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 24
+class AAARTransformer(nn.Module):
+ def __init__(self, width=384, depth=4, ctx_n=250, n_head=6, fromq=2, toq=4):
+ super().__init__()
+
+ pos_embs = sinusoids(ctx_n * toq, width)
+
+ self.encoder = PureEncoder(pos_embs=pos_embs, length=ctx_n*toq, width=width, n_head=n_head, depth=depth)
+
+ def forward(self, intoks, outtoks, loss=True):
+ with record_function("decoder"):
+ logits = self.encoder(intoks)
+ if loss is not None:
+ with record_function("loss"):
+ loss = F.cross_entropy(logits.reshape(-1,logits.shape[-1]), outtoks.view(-1))
+ return logits, loss
+
+# %% ../nbs/6. EnCodec quality enchancement.ipynb 25
+def make_model(size:str, dataset:torch.utils.data.Dataset=None):
+ assert(dataset is not None)
+ kwargs = dict(fromq = dataset.fromq, toq = dataset.toq, ctx_n = dataset.n_ctx)
+ if size == 'tiny':
+ return AAARTransformer(depth=4, **kwargs)
+ elif size == 'base':
+ return AAARTransformer(depth=6, width=512, n_head=8, **kwargs)
+ elif size == 'small':
+ return AAARTransformer(depth=12, width=768, n_head=12, **kwargs)