Add codec modules.

codec/codec_decoder/decoder.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+from encoder.codec import VocosBackbone
+
+
+class SimpleDecoder(nn.Module):
+    def __init__(
+        self,
+        n_mels=50,
+        encoder_dim=768,
+        bottleneck_channels=5,
+        num_layers=8,
+        intermediate_dim=None,
+        upsample_scale=4,
+        dw_kernel=5,
+    ):
+        super().__init__()
+
+        intermediate_dim = intermediate_dim or encoder_dim * 3
+        self.upsample_scale = upsample_scale
+
+        # project FSQ channels back to model dim
+        self.in_proj = nn.Linear(bottleneck_channels, encoder_dim)
+
+        # ConvNeXt backbone
+        self.backbone = VocosBackbone(
+            input_channels=encoder_dim,
+            dim=encoder_dim,
+            intermediate_dim=intermediate_dim,
+            num_layers=num_layers,
+            input_kernel_size=1,
+            dw_kernel_size=dw_kernel,
+        )
+
+        # output mel projection
+        self.out_proj = nn.Conv1d(encoder_dim, n_mels, kernel_size=1)
+
+    def forward(self, z):
+        """
+        z: (B, T_latent, bottleneck_channels)
+        """
+        z = self.in_proj(z)              # (B, T_latent, D)
+        z = z.transpose(1, 2)            # (B, D, T_latent)
+
+        # naive upsampling (good enough for now)
+        z = z.repeat_interleave(self.upsample_scale, dim=2)
+
+        z = self.backbone(z)             # (B, D, T_mel)
+        mel_hat = self.out_proj(z)       # (B, n_mels, T_mel)
+        return mel_hat

codec/encoder/codec.py

@@ -0,0 +1,213 @@
+"""
+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)
+
+        # import pdb;pdb.set_trace()
+
+        x = x[:, :, ::self.downsample_scale] # What the heck is this? Brute force downsampling from mel -> tokens. 
+        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)
+
+        # print("done preprocessing: ",x.shape)
+        # import pdb;pdb.set_trace()
+
+        x = self.encode(x)
+        codes = self.quant.to_codebook_index(x)
+        return codes

codec/encoder/quantizer.py

@@ -0,0 +1,64 @@
+"""
+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)
+        else:
+            self._levels = levels
+
+        _levels = self._levels
+        self.register_buffer("levels", _levels, persistent=False)
+
+        _basis = torch.cumprod(torch.tensor([1] + levels[:-1]),
+                                    dim=0,
+                                    dtype=torch.int32)
+        self.register_buffer("_basis", _basis, persistent=False)
+
+    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:
+
+            half_levels = (self.levels - 1) * (1 - 1e-3) / 2
+            offset = 0.5 - 0.5 * (self.levels % 2)
+            shift = torch.tan(offset / half_levels)
+
+            x = torch.tanh(x + shift) * half_levels - offset
+            x = x + (x.round() - x).detach()
+            x = x / (self.levels // 2)
+        return x

codec_dataset.py

@@ -0,0 +1,53 @@
+import torch
+import torchaudio
+from torch.utils.data import Dataset
+import os
+import json
+
+
+class LJSpeechDataset(Dataset):
+    def __init__(self, root, sample_rate=32000,  mode='train'):
+        """
+        root: path to LJSpeech-1.1 directory
+        """
+        self.root = root
+        self.sample_rate = sample_rate
+
+        # Write code here to handle modes. Take wave files only from mode json. 
+        self.mode = mode
+        mode_json = os.path.join(root, f"{mode}.json")
+        with open(mode_json, 'r') as f:
+            self.dataset = json.load(f)
+
+        # self.wav_dir = os.path.join(root, "wavs")
+        # self.wav_files = sorted(
+        #     [f for f in os.listdir(self.wav_dir) if f.endswith(".wav")]
+        # )
+
+        # assert len(self.wav_files) > 0, "No wav files found!"
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        # wav_path = os.path.join(self.wav_dir, self.wav_files[idx])
+        
+        # Using train and val json files
+        item = self.dataset[idx]
+        text, audio_tokens, wav_path = item
+
+        wav, sr = torchaudio.load(wav_path)
+
+        # mono
+        if wav.shape[0] > 1:
+            wav = wav.mean(dim=0, keepdim=True)
+
+        # resample if needed
+        if sr != self.sample_rate:
+            wav = torchaudio.functional.resample(
+                wav, orig_freq=sr, new_freq=self.sample_rate
+            )
+
+        # print("wav shape is: ", wav.shape, idx)
+
+        return wav

codec_model.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn
+from codec.encoder.codec import Encoder
+from codec.codec_decoder.decoder import SimpleDecoder
+
+
+class FSQAutoEncoder(nn.Module):
+    def __init__(self, encoder_cfg, decoder_cfg):
+        super().__init__()
+        self.encoder = Encoder(**encoder_cfg)
+        self.decoder = SimpleDecoder(**decoder_cfg)
+
+    def forward(self, audio):
+        mel = self.encoder.preprocess(audio)
+        z = self.encoder.encode(mel)      # FSQ STE output
+        mel_hat = self.decoder(z)
+        return mel_hat, mel

codec_train.py

@@ -0,0 +1,247 @@
+import torch
+import torchaudio
+from torch.utils.data import DataLoader
+from torch.optim import AdamW, Adam
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import os
+import wandb
+from tqdm import tqdm
+
+
+from codec_model import FSQAutoEncoder
+from codec_dataset import LJSpeechDataset
+from codec.codec_decoder.decoder import SimpleDecoder
+
+
+wandb.init(project="soprano-codec")
+
+
+def pad_collate(batch):
+    """
+    batch: list of tensors [(1, T1), (1, T2), ...]
+    """
+    lengths = torch.tensor([x.shape[-1] for x in batch])
+
+    max_len = lengths.max().item()
+
+    padded = [
+        F.pad(x, (0, max_len - x.shape[-1]))
+        for x in batch
+    ]
+
+    audio = torch.stack(padded)  # (B, 1, T_max)
+    return audio, lengths
+
+def plot_mels():
+    pass
+
+
+
+dataset = LJSpeechDataset(
+    root="/home/ubuntu/soma/data/lj_speech/LJSpeech-1.1",
+    sample_rate=32000,
+)
+
+loader = DataLoader(
+    dataset,
+    batch_size=16,
+    shuffle=True,
+    drop_last=True,
+    num_workers=8,
+    pin_memory=True,
+    collate_fn=pad_collate
+)
+
+val_dataset = LJSpeechDataset(
+    root="/home/ubuntu/soma/data/lj_speech/LJSpeech-1.1",
+    sample_rate=32000,
+    mode='val'
+)
+
+val_loader = DataLoader(
+    val_dataset,
+    batch_size=16,
+    shuffle=False,
+    drop_last=True,
+    num_workers=8,
+    pin_memory=True,
+    collate_fn=pad_collate
+)
+val_loader_it = iter(val_loader)
+
+# ------------------
+# Config
+# ------------------
+
+encoder_cfg = dict(
+    num_input_mels=50,
+    mel_hop_length=512,
+    encoder_dim=768,
+    encoder_num_layers=8,
+    fsq_levels=[8, 8, 5, 5, 5],
+)
+
+decoder_cfg = dict(
+    n_mels=50,
+    encoder_dim=768,
+    bottleneck_channels=5,
+    num_layers=8,
+    upsample_scale=2048 // 512,
+)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 3. Token rate sanity check
+# downsample_scale = 2048 // mel_hop_length
+# total_hop = mel_hop_length * downsample_scale = 2048
+# rate = sr / 2048
+sr = 32000
+total_hop = 2048 # This is the token hop rate. The mel hop rate is 512
+print(f"Token rate: {sr / total_hop:.2f} Hz")
+
+
+# ------------------
+# Setup
+# ------------------
+
+model = FSQAutoEncoder(encoder_cfg, decoder_cfg).to(device)
+
+freeze_encoder = False
+if freeze_encoder:
+    model_ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_codec/codec_1/step_42000.pt"
+
+    if os.path.exists(model_ckpt_path):
+        print(f"Loading model from {model_ckpt_path}")
+        model.load_state_dict(torch.load(model_ckpt_path))
+
+    # fix encoder. train only the decoder. reset the decoder weights.
+    for param in model.encoder.parameters():
+        param.requires_grad = False
+
+    for name, p in model.named_parameters():
+        if "quant" in name:
+            print(name, p.requires_grad)
+
+    model.decoder = SimpleDecoder(**decoder_cfg).to(device)
+
+optimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)
+
+import pdb;pdb.set_trace()
+
+
+root_dir = "/home/ubuntu/soma/ckpt/suprano/suprano_codec"
+ckpt_dir = os.path.join(root_dir, "codec_v2")
+plot_dir = os.path.join(ckpt_dir, "plots")
+os.makedirs(ckpt_dir, exist_ok=True)
+os.makedirs(plot_dir, exist_ok=True)
+
+num_epochs = 100
+step = 0
+
+for epoch in range(num_epochs):
+
+    for epoch_step, data in tqdm(enumerate(loader), total=len(loader)):
+
+        step += 1
+
+        audio, lengths = data
+        audio = audio.squeeze().to(device)
+
+        # import pdb;pdb.set_trace()
+
+        mel_hat, mel = model(audio)
+
+        # crop to match length (upsampling can overshoot)
+        T = min(mel_hat.shape[-1], mel.shape[-1])
+        mel_hat = mel_hat[..., :T]
+        mel = mel[..., :T]
+
+        loss = torch.mean(torch.abs(mel_hat - mel))
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if step % 100 == 0:
+            print(f"step {step} | loss {loss.item():.4f}")
+            wandb.log({"train/loss": loss.item()}, step=step)
+
+        if step % 400 == 0:
+            val_loss = 0.0
+            val_steps = 10
+            model.eval()
+            with torch.no_grad():
+                for _ in range(val_steps):
+                    try:
+                        vdata = next(val_loader_it)
+                    except StopIteration:
+                        val_loader_it = iter(val_loader)
+                        vdata = next(val_loader_it)
+                    
+                    vaudio, vlengths = vdata
+                    vaudio = vaudio.squeeze().to(device)
+                    
+                    vmel_hat, vmel = model(vaudio)
+                    
+                    T_v = min(vmel_hat.shape[-1], vmel.shape[-1])
+                    vmel_hat = vmel_hat[..., :T_v]
+                    vmel = vmel[..., :T_v]
+                    
+                    val_loss += torch.mean(torch.abs(vmel_hat - vmel)).item()
+            
+            val_loss /= val_steps
+            print(f"step {step} | val_loss {val_loss:.4f}")
+            wandb.log({"val/loss": val_loss}, step=step)
+
+            # Val Plotting
+            vmel_np = vmel[0].detach().cpu().numpy()
+            vmel_hat_np = vmel_hat[0].detach().cpu().numpy()
+
+            fig, axs = plt.subplots(2, 1, figsize=(10, 6))
+
+            axs[0].imshow(vmel_np, aspect="auto", origin="lower")
+            axs[0].set_title("Val Original Mel")
+
+            axs[1].imshow(vmel_hat_np, aspect="auto", origin="lower")
+            axs[1].set_title("Val Reconstructed Mel")
+
+            plt.tight_layout()
+            plt.savefig(f"{plot_dir}/val_step_{step:05d}.png")
+            wandb.log({"val/reconstruction": wandb.Image(f"{plot_dir}/val_step_{step:05d}.png")}, step=step)
+            plt.close()
+            
+            model.train()
+
+        if step % 400 == 0:
+            with torch.no_grad():
+                # 1. FSQ bin usage
+                z = model.encoder.encode(mel)
+                indices = model.encoder.quant.to_codebook_index(z)
+                total_bins = int(torch.prod(model.encoder.quant.levels))
+                unique_bins = len(torch.unique(indices))
+                print(f"[FSQ] unique bins: {unique_bins} / {total_bins}. Lens: {indices.shape} {z.shape}")
+                wandb.log({"train/unique_bins": unique_bins}, step=step)
+
+        if step % 200 == 0:
+            mel_np = mel[0].detach().cpu().numpy()
+            mel_hat_np = mel_hat[0].detach().cpu().numpy()
+
+            fig, axs = plt.subplots(2, 1, figsize=(10, 6))
+
+            axs[0].imshow(mel_np, aspect="auto", origin="lower")
+            axs[0].set_title("Original Mel")
+
+            axs[1].imshow(mel_hat_np, aspect="auto", origin="lower")
+            axs[1].set_title("Reconstructed Mel")
+
+            plt.tight_layout()
+            plt.savefig(f"{plot_dir}/step_{step:05d}.png")
+            wandb.log({"train/reconstruction": wandb.Image(f"{plot_dir}/step_{step:05d}.png")}, step=step)
+            plt.close()
+
+
+        if step % 1000 == 0:
+            ckpt_path = os.path.join(ckpt_dir, f"step_{step:05d}.pt")
+            torch.save(model.state_dict(), ckpt_path)
+            print(f"Saved checkpoint to {ckpt_path}")