first commit

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+test.py
+*.json
+*.pth
+dist/
+*.egg-info/

dataset.py

@@ -0,0 +1,17 @@
+import json
+from torch.utils.data import Dataset
+
+
+class AudioDataset(Dataset):
+    def __init__(self, path):
+        with open(path, encoding='utf-8') as f:
+            self.dataset = json.load(f)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        text, audio = self.dataset[idx]
+        # Format: [STOP][TEXT]<text prompt>[START]<audio tokens>[STOP]
+        res = f"[STOP][TEXT]{text}[START]{''.join(list(map(lambda x: f'[{x}]', audio)))}[STOP]"
+        return res

encoder/codec.py

@@ -0,0 +1,206 @@
+"""
+Adapted from https://github.com/gemelo-ai/vocos
+"""
+from typing import Optional
+
+import torchaudio
+import torch
+from torch import nn
+
+from .quantizer import FSQSTE
+
+def safe_log(x: torch.Tensor, clip_val: float = 5e-3) -> torch.Tensor:
+    return torch.log(torch.clip(x, min=clip_val))
+
+
+class SimpleMLP(nn.Module):
+    def __init__(self,
+        dim,
+        intermediate_dim,
+    ):
+        super().__init__()
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+    def forward(self, x):
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        return x
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        layer_scale_init_value: float,
+        dw_kernel_size: int = 7,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=dw_kernel_size, padding=dw_kernel_size//2, groups=dim)  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.mlp = SimpleMLP(dim, intermediate_dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+        x = self.norm(x)
+        x = self.mlp(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+
+        x = residual + x
+        return x
+
+
+class VocosBackbone(nn.Module):
+    """
+    Vocos backbone module built with ConvNeXt blocks.
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        layer_scale_init_value (float, optional): Initial value for layer scaling.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        dim: int,
+        intermediate_dim: int,
+        num_layers: int,
+        input_kernel_size: int = 7,
+        dw_kernel_size: int = 7,
+        layer_scale_init_value: Optional[float] = None,
+        pad: str = 'zeros',
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.dim = dim
+        self.embed = nn.Conv1d(
+            input_channels,
+            dim,
+            kernel_size=input_kernel_size,
+            padding=input_kernel_size//2,
+            padding_mode=pad
+        )
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.convnext = nn.ModuleList([
+            ConvNeXtBlock(
+                dim=dim,
+                intermediate_dim=intermediate_dim,
+                dw_kernel_size=dw_kernel_size,
+                layer_scale_init_value=layer_scale_init_value or 1 / num_layers**0.5,
+            )
+            for _ in range(num_layers)
+        ])
+        self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
+                        C denotes output features, and L is the sequence length.
+
+        Returns:
+            Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
+                    and H denotes the model dimension.
+        """
+        x = self.embed(x) # (B, C, L)
+        x = self.norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        for conv_block in self.convnext:
+            x = conv_block(x)
+        x = self.final_layer_norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self,
+        num_input_mels=50,
+        mel_hop_length=512,
+        mel_hop_scale=0.25,
+        encoder_num_layers=8,
+        encoder_dim=768,
+        encoder_intermediate_dim=None,
+        fsq_levels=[8, 8, 5, 5, 5],
+        dw_kernel=5,
+    ):
+        super().__init__()
+        self.downsample_scale = 2048 // mel_hop_length
+        self.mel_hop_length = mel_hop_length
+        self.mel_n_fft = int(mel_hop_length/mel_hop_scale)
+        self.encoder_dim = encoder_dim
+        self.encoder_intermediate_dim = encoder_intermediate_dim if encoder_intermediate_dim else encoder_dim*3
+        self.encoder_num_layers = encoder_num_layers
+        self.encoder_initial_channels = num_input_mels
+        self.bottleneck_channels = 5
+
+        self.mel_spec = torchaudio.transforms.MelSpectrogram(
+            sample_rate=32000,
+            n_fft=self.mel_n_fft,
+            hop_length=self.mel_hop_length,
+            n_mels=num_input_mels,
+            center=True,
+            power=1,
+        )
+        self.encoder = VocosBackbone(input_channels=self.encoder_initial_channels,
+            dim=self.encoder_dim,
+            intermediate_dim=self.encoder_intermediate_dim,
+            num_layers=self.encoder_num_layers,
+            input_kernel_size=1,
+            dw_kernel_size=dw_kernel,
+            pad='zeros'
+        )
+        self.downsampler = nn.Linear(self.encoder_dim, self.bottleneck_channels)
+        self.quant = FSQSTE(levels=fsq_levels)
+
+    def encode(self, x):
+        x = self.encoder(x)
+        x = x[:, :, ::self.downsample_scale]
+        x = x.transpose(1,2)
+        x = self.downsampler(x)
+        x = self.quant(x)
+        return x
+
+    def preprocess(self, audio):
+        if audio.dim() == 2: # raw audio
+            x = self.mel_spec(audio)
+            x = safe_log(x)
+        elif audio.dim() == 3: # mel spectrogram
+            x = audio
+        return x
+
+    def forward(self, audio):
+        x = self.preprocess(audio)
+        x = self.encode(x)
+        codes = self.quant.to_codebook_index(x)
+        return codes

encoder/quantizer.py

@@ -0,0 +1,58 @@
+"""
+Adapted from https://github.com/duchenzhuang/FSQ-pytorch/blob/main/quantizers/fsq.py#L41
+"""
+
+import torch
+from torch import nn
+from einops import rearrange
+
+
+class FSQSTE(nn.Module):
+    def __init__(self, levels):
+        super().__init__()
+        if levels:
+            self.dim = len(levels)
+            self.levels = torch.tensor(levels, dtype=torch.int32).view(1, 1, self.dim)
+            self.half_levels = (self.levels - 1) * (1 - 1e-3) / 2
+            self.offset = 0.5 - 0.5 * (self.levels % 2)
+            self.shift = torch.tan(self.offset / self.half_levels)
+        else:
+            self.levels = levels
+
+        self._basis = torch.cumprod(torch.tensor([1] + levels[:-1]),
+                                    dim=0,
+                                    dtype=torch.int32)
+
+    def _scale_and_shift(self, zhat_normalized):
+        half_width = self.levels // 2
+        return (zhat_normalized * half_width) + half_width
+
+    def _scale_and_shift_inverse(self, zhat):
+        half_width = self.levels // 2
+        return (zhat - half_width) / half_width
+
+    def indices_to_level_indices(self, indices):
+        """ Converts indices to indices at each level, perhaps needed for a transformer with factorized embeddings """
+        indices = rearrange(indices, '... -> ... 1')
+        codes_non_centered = (indices // self._basis) % self.levels
+        return codes_non_centered
+
+    def to_codebook_index(self, zhat):
+        """ Converts a `code` to an index in the codebook. """
+        assert zhat.shape[-1] == self.dim
+        zhat = self._scale_and_shift(zhat)
+        indices = (zhat * self._basis).sum(dim = -1).round().to(torch.int32)
+        return indices
+    
+    def from_codebook_index(self, indices):
+        """ Inverse of `codes_to_indices`. """
+        level_indices = self.indices_to_level_indices(indices)
+        codes = self._scale_and_shift_inverse(level_indices)
+        return codes
+
+    def forward(self, x):
+        if self.levels is not None:
+            x = torch.tanh(x + self.shift) * self.half_levels - self.offset
+            x = x + (x.round() - x).detach()
+            x = x / (self.levels // 2)
+        return x

example_dataset/metadata.txt

@@ -0,0 +1,2 @@
+example1|Soprano is an extremely lightweight text to speech model designed to produce highly realistic speech at unprecedented speed.
+example2|Gabagool? Ova here!

generate_dataset.py

@@ -0,0 +1,90 @@
+"""
+Converts a dataset in LJSpeech format into audio tokens that can be used to train/fine-tune Soprano.
+This script creates two JSON files for train and test splits in the provided directory.
+
+Usage:
+python generate_dataset.py --input-dir path/to/files
+
+Args:
+--input-dir: Path to directory of LJSpeech-style dataset. If none is provided this defaults to the provided example dataset.
+"""
+import argparse
+import pathlib
+import random
+import json
+
+from scipy.io import wavfile
+import torchaudio
+import torch
+from tqdm import tqdm
+from huggingface_hub import hf_hub_download
+
+from encoder.codec import Encoder
+
+
+SAMPLE_RATE = 32000
+SEED = 42
+VAL_PROP = 0.1
+VAL_MAX = 512
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input-dir",
+        required=False,
+        default="./example_dataset",
+        type=pathlib.Path
+    )
+    return parser.parse_args()
+
+def main():
+    args = get_args()
+    input_dir = args.input_dir
+
+    print("Loading model.")
+    encoder = Encoder()
+    encoder_path = hf_hub_download(repo_id='ekwek/Soprano-Encoder', filename='encoder.pth')
+    encoder.load_state_dict(torch.load(encoder_path))
+    print("Model loaded.")
+
+
+    print("Reading metadata.")
+    files = []
+    with open(f'{input_dir}/metadata.txt', encoding='utf-8') as f:
+        data = f.read().split('\n')
+        for line in data:
+            filename, transcript = line.split('|', maxsplit=1)
+            files.append((filename, transcript))
+    print(f'{len(files)} samples located in directory.')
+
+    print("Encoding audio.")
+    dataset = []
+    for sample in tqdm(files):
+        filename, transcript = sample
+        sr, audio = wavfile.read(f'{input_dir}/wavs/{filename}.wav')
+        audio = torch.from_numpy(audio)
+        if sr != SAMPLE_RATE:
+            audio = torchaudio.functional.resample(audio, sr, SAMPLE_RATE)
+        audio = audio.unsqueeze(0)
+        with torch.no_grad():
+            audio_tokens = encoder(audio)
+        dataset.append([transcript, audio_tokens.squeeze(0).tolist()])
+
+    print("Generating train/test splits.")
+    random.seed(SEED)
+    random.shuffle(dataset)
+    num_val = min(int(VAL_PROP * len(dataset)) + 1, VAL_MAX)
+    train_dataset = dataset[num_val:]
+    val_dataset = dataset[:num_val]
+    print(f'# train samples: {len(train_dataset)}')
+    print(f'# val samples: {len(val_dataset)}')
+
+    print("Saving datasets.")
+    with open(f'{input_dir}/train.json', 'w', encoding='utf-8') as f:
+        json.dump(train_dataset, f, indent=2)
+    with open(f'{input_dir}/val.json', 'w', encoding='utf-8') as f:
+        json.dump(val_dataset, f, indent=2)
+    print("Datasets saved.")
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,8 @@
+einops
+huggingface_hub
+numpy
+scipy
+torch
+torchaudio
+tqdm
+transformers

train.py

@@ -0,0 +1,216 @@
+"""
+Training script for Soprano.
+
+Usage:
+python train.py --input-dir path/to/files --save-dir path/to/weights
+
+Args:
+--input-dir: Path to directory of LJSpeech-style dataset. If none is provided this defaults to the provided example dataset.
+--save-dir: Path to directory to save weights
+
+Adapted from https://github.com/karpathy/nanoGPT
+"""
+import argparse
+import pathlib
+import random
+import time
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from dataset import AudioDataset
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input-dir",
+        required=False,
+        default="./example_dataset",
+        type=pathlib.Path
+    )
+    parser.add_argument("--save-dir",
+        required=True,
+        type=pathlib.Path
+    )
+    return parser.parse_args()
+
+args = get_args()
+
+# training hyperparameters
+device = 'cuda:0'
+seed = 1337
+max_lr = 5e-4
+warmup_ratio = 0.1
+cooldown_ratio = 0.1
+min_lr = 0.1 * max_lr
+batch_size = 16
+grad_accum_steps = 1
+seq_len = 1024
+val_freq = 250
+text_factor = 0.0 # currently does not train on text inputs, you can increase to change this
+max_steps = 10000
+betas = (0.9, 0.95)
+weight_decay = 0.1
+train_dataset_path = f'{args.input_dir}/train.json'
+val_dataset_path = f'{args.input_dir}/val.json'
+save_path = args.save_dir
+
+device_type = "cuda" if device.startswith("cuda") else "cpu"
+torch.manual_seed(seed)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed(seed)
+def worker_seed_init(_):
+    worker_seed = torch.initial_seed() % (2**32-1)
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+torch.set_float32_matmul_precision('high')
+print(f"Save Path: {save_path}")
+
+# lr schedule
+warmup_steps = int(max_steps * warmup_ratio)
+cooldown_steps = int(max_steps * cooldown_ratio)
+def get_lr(it): # WSD schedule
+    if it<warmup_steps:
+        return max_lr * (it+1) / warmup_steps
+    if it<max_steps-cooldown_steps:
+        return max_lr
+    return min_lr + (max_lr-min_lr) * ((max_steps-it) / cooldown_steps)
+
+# model
+model = AutoModelForCausalLM.from_pretrained('ekwek/Soprano-80M')
+tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
+model.to(torch.bfloat16).to(device)
+model.train()
+
+# dataset
+def collate_pack(texts):
+    tokens_batch = tokenizer(texts, padding=False, truncation=False)
+    batch = []
+    cur_sample, cur_size = [], 0
+    for i in range(len(texts)):
+        tokens = torch.tensor(tokens_batch['input_ids'][i][:-1], dtype=torch.long)
+        cur_size += tokens.size(0)
+        cur_sample.append(tokens)
+        if cur_size >= seq_len + 1:
+            batch.append(torch.cat(cur_sample)[: seq_len + 1])
+            cur_sample, cur_size = [], 0
+            if len(batch) == batch_size:
+                break
+    if cur_sample and not batch: # add partial sample if there isn't enough data
+        batch.append(torch.cat(cur_sample + [torch.zeros(seq_len, dtype=torch.long)])[: seq_len + 1])
+    if len(batch) < batch_size:
+        # pad up to batch_size for consistency
+        pad = batch[-1]
+        while len(batch) < batch_size:
+            batch.append(pad)
+    batch = torch.stack(batch)
+    x = batch[:, :-1]
+    y = batch[:, 1:]
+    return x, y
+
+dataset = AudioDataset(train_dataset_path)
+# we need batch_size * 16 to have enough tokens after packing
+dataloader = DataLoader(dataset,
+    batch_size=batch_size * 16,
+    shuffle=True,
+    num_workers=16,
+    pin_memory=True,
+    persistent_workers=True,
+    worker_init_fn=worker_seed_init,
+    collate_fn=collate_pack,
+)
+dataloader_it = iter(dataloader)
+val_dataset = AudioDataset(val_dataset_path)
+val_dataloader = DataLoader(val_dataset,
+    batch_size=batch_size * 16,
+    shuffle=False,
+    num_workers=1,
+    pin_memory=True,
+    persistent_workers=True,
+    worker_init_fn=worker_seed_init,
+    collate_fn=collate_pack,
+)
+
+# optimizer
+opt = torch.optim.AdamW(model.parameters(), max_lr, betas=betas, weight_decay=weight_decay, fused=True)
+
+def compute_loss(logits, y, num_steps):
+    pred = logits.view(-1, logits.size(-1))
+    labels = y.reshape(-1)
+    loss = torch.nn.functional.cross_entropy(pred, labels, reduction='none')
+    audio_mask = torch.logical_and(y>=3, y<=8003).view(-1)
+    audio_loss = loss[audio_mask].mean()
+    text_loss = loss[~audio_mask].mean()
+    acc = (logits.argmax(dim=-1) == y).view(-1)[audio_mask].to(torch.float32).mean()
+    audio_loss = audio_loss / num_steps
+    text_loss = text_loss / num_steps
+    acc = acc / num_steps
+    return audio_loss, text_loss, acc
+
+def evaluate(val_dataloader):
+    model.eval()
+    val_dataloader_it = iter(val_dataloader)
+    with torch.no_grad():
+        val_audio_loss_accum = torch.tensor(0.0).to(device)
+        val_text_loss_accum = torch.tensor(0.0).to(device)
+        val_acc_accum = torch.tensor(0.0).to(device)
+        val_loss_steps = 1
+        for _ in range(val_loss_steps):
+            x, y = next(val_dataloader_it)
+            x, y = x.to(device), y.to(device)
+            with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+                logits = model(x).logits
+                audio_loss, text_loss, acc = compute_loss(logits, y, val_loss_steps)
+            val_audio_loss_accum += audio_loss.detach()
+            val_text_loss_accum += text_loss.detach()
+            val_acc_accum += acc.detach()
+        print(f"validation text loss: {val_text_loss_accum.item():.4f}\tvalidation audio loss: {val_audio_loss_accum.item():.4f}\tvalidation acc: {val_acc_accum.item():.4f}")
+    model.train()
+
+pbar = tqdm(range(0, max_steps), ncols=200, dynamic_ncols=True)
+for step in pbar:
+    start = time.time()
+    if val_freq>0 and (step % val_freq == 0 or step==max_steps-1):
+        evaluate(val_dataloader)
+
+    opt.zero_grad()
+    audio_loss_accum = 0.0
+    text_loss_accum = 0.0
+    acc_accum = 0.0
+    for micro_step in range(grad_accum_steps):
+        try:
+            x, y = next(dataloader_it)
+        except:
+            dataloader_it = iter(dataloader)
+            x, y = next(dataloader_it)
+        x, y = x.to(device), y.to(device)
+
+        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+            logits = model(x).logits
+            audio_loss, text_loss, acc = compute_loss(logits, y, grad_accum_steps)
+        audio_loss_accum += audio_loss.detach()
+        text_loss_accum += text_loss.detach()
+        acc_accum += acc.detach()
+        total_loss = audio_loss + text_factor*text_loss
+        total_loss.backward()
+
+    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    lr = get_lr(step)
+    for param_group in opt.param_groups:
+        param_group['lr'] = lr
+    opt.step()
+    torch.cuda.synchronize()
+    total_tokens = step * batch_size*seq_len*grad_accum_steps
+    end = time.time()
+    dt = (end-start)*1000
+    tokens_per_second = (batch_size*seq_len*grad_accum_steps) / (end-start)
+    tqdm_log = f'text loss: {text_loss_accum.item():.3f} | audio loss: {audio_loss_accum.item():.3f} | acc: {acc_accum.item():.4f} | lr: {lr:.2e} | norm: {norm:.3f} | time: {dt:.2f} ms | {tokens_per_second:.2f} t/s'
+    pbar.set_description(tqdm_log)
+
+print(f"Training complete. Saving model at {save_path}")
+model.save_pretrained(save_path)
+tokenizer.save_pretrained(save_path)
+print("Saving done.")