models/codec/amphion_codec/quantize/vector_quantize.py

# Copyright (c) 2024 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from torch.nn.utils import weight_norm


def WNConv1d(*args, **kwargs):
    return weight_norm(nn.Conv1d(*args, **kwargs))


def WNConvTranspose1d(*args, **kwargs):
    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))


def l2norm(t):
    return F.normalize(t, p=2, dim=-1)


def ema_inplace(moving_avg, new, decay):
    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))


def laplace_smoothing(x, n_categories, eps=1e-5):
    return (x + eps) / (x.sum() + n_categories * eps)


def sample_vectors(samples, num):
    num_samples, device = samples.shape[0], samples.device

    if num_samples >= num:
        indices = torch.randperm(num_samples, device=device)[:num]
    else:
        indices = torch.randint(0, num_samples, (num,), device=device)

    return samples[indices]


def kmeans(samples, num_clusters, num_iters=10, use_cosine_sim=False):
    dim, dtype, device = samples.shape[-1], samples.dtype, samples.device

    means = sample_vectors(samples, num_clusters)

    for _ in range(num_iters):
        if use_cosine_sim:
            dists = samples @ means.t()
        else:
            diffs = rearrange(samples, "n d -> n () d") - rearrange(
                means, "c d -> () c d"
            )
            dists = -(diffs**2).sum(dim=-1)

        buckets = dists.max(dim=-1).indices
        bins = torch.bincount(buckets, minlength=num_clusters)
        zero_mask = bins == 0
        bins_min_clamped = bins.masked_fill(zero_mask, 1)

        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
        new_means = new_means / bins_min_clamped[..., None]

        if use_cosine_sim:
            new_means = l2norm(new_means)

        means = torch.where(zero_mask[..., None], means, new_means)

    return means, bins


class EuclideanCodebook(nn.Module):
    def __init__(
        self,
        dim,
        codebook_size,
        kmeans_init=False,
        kmeans_iters=10,
        decay=0.8,
        eps=1e-5,
        threshold_ema_dead_code=2,
        weight_init=False,
    ):
        super().__init__()

        self.decay = decay
        init_fn = torch.randn if not weight_init else torch.zeros
        embed = init_fn(codebook_size, dim)

        if weight_init:
            nn.init.uniform_(embed, -1 / codebook_size, 1 / codebook_size)

        self.codebook_size = codebook_size
        self.kmeans_iters = kmeans_iters
        self.eps = eps
        self.threshold_ema_dead_code = threshold_ema_dead_code

        self.register_buffer(
            "initted", torch.Tensor([not kmeans_init])
        )  # if kmeans_init is True, then initted is False; otherwise, initted is True
        self.register_buffer("cluster_size", torch.zeros(codebook_size))
        self.register_buffer("embed", embed)
        self.register_buffer("embed_avg", embed.clone())

    def init_embed_(self, data):
        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
        self.embed.data.copy_(embed)
        self.embed_avg.data.copy_(embed)
        self.cluster_size.data.copy_(cluster_size)
        self.initted.data.copy_(torch.Tensor([True]))

    def replace(self, samples, mask):
        modified_codebook = torch.where(
            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
        )
        self.embed.data.copy_(modified_codebook)

    def expire_codes_(self, batch_samples):
        if self.threshold_ema_dead_code == 0:
            return

        expired_codes = self.cluster_size < self.threshold_ema_dead_code
        if not torch.any(expired_codes):
            return
        batch_samples = rearrange(batch_samples, "... d -> (...) d")
        self.replace(batch_samples, mask=expired_codes)

    def forward(self, x):
        shape, dtype = x.shape, x.dtype
        flatten = rearrange(x, "... d -> (...) d")
        embed = self.embed.t()  # (codebook_size, dim) -> (dim, codebook_size)

        if not self.initted:
            self.init_embed_(flatten)

        dist = -(
            flatten.pow(2).sum(1, keepdim=True)
            - 2 * flatten @ embed
            + embed.pow(2).sum(0, keepdim=True)
        )

        embed_ind = dist.max(dim=-1).indices
        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
        embed_ind = embed_ind.view(*shape[:-1])
        quantize = F.embedding(embed_ind, self.embed)

        if self.training:
            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
            embed_sum = (
                flatten.t() @ embed_onehot
            )  # (dim, ...) @ (..., codebook_size) -> (dim, codebook_size)
            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
            cluster_size = (
                laplace_smoothing(self.cluster_size, self.codebook_size, self.eps)
                * self.cluster_size.sum()
            )
            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
            self.embed.data.copy_(embed_normalized)
            self.expire_codes_(x)

        return quantize, embed_ind

    def vq2emb(self, vq):
        quantize = F.embedding(vq, self.embed)
        return quantize

    def latent2dist(self, x):
        shape, dtype = x.shape, x.dtype
        flatten = rearrange(x, "... d -> (...) d")
        embed = self.embed.t()  # (codebook_size, dim) -> (dim, codebook_size)

        if not self.initted:
            self.init_embed_(flatten)

        dist = -(
            flatten.pow(2).sum(1, keepdim=True)
            - 2 * flatten @ embed
            + embed.pow(2).sum(0, keepdim=True)
        )

        embed_ind = dist.max(dim=-1).indices
        embed_ind = embed_ind.view(*shape[:-1])
        quantize = F.embedding(embed_ind, self.embed)

        dist = dist.view(*shape[:-1], -1)

        return dist, embed_ind, quantize


class SimpleCodebook(nn.Module):
    def __init__(
        self,
        dim,
        codebook_size,
        use_l2_normlize=False,
    ):
        super().__init__()

        self.dim = dim
        self.codebook_size = codebook_size
        self.use_l2_normlize = use_l2_normlize

        self.embed = nn.Embedding(self.codebook_size, self.dim)

    def forward(self, x):
        shape, dtype = x.shape, x.dtype
        flatten = rearrange(x, "... d -> (...) d")
        embed = self.embed.weight.t()  # (codebook_size, dim) -> (dim, codebook_size)

        if self.use_l2_normlize:
            flatten = F.normalize(flatten)
            embed = F.normalize(embed)

        dist = -(
            flatten.pow(2).sum(1, keepdim=True)
            - 2 * flatten @ embed
            + embed.pow(2).sum(0, keepdim=True)
        )

        embed_ind = dist.max(dim=-1).indices
        embed_ind = embed_ind.view(*shape[:-1])
        quantize = F.embedding(embed_ind, self.embed)

        return quantize, embed_ind

    def vq2emb(self, vq):
        quantize = F.embedding(vq, self.embed.weight)
        return quantize

    def latent2dist(self, x):
        shape, dtype = x.shape, x.dtype
        flatten = rearrange(x, "... d -> (...) d")
        embed = self.embed.weight.t()  # (codebook_size, dim) -> (dim, codebook_size)

        if self.use_l2_normlize:
            flatten = F.normalize(flatten)
            embed = F.normalize(embed)

        dist = -(
            flatten.pow(2).sum(1, keepdim=True)
            - 2 * flatten @ embed
            + embed.pow(2).sum(0, keepdim=True)
        )

        embed_ind = dist.max(dim=-1).indices
        embed_ind = embed_ind.view(*shape[:-1])
        quantize = F.embedding(embed_ind, self.embed)

        dist = dist.view(*shape[:-1], -1)

        return dist, embed_ind, quantize


class VectorQuantize(nn.Module):
    """Vector quantization and factorized vecotor quantization implementation
    Args:
        input_dim (int): Dimension of input.
        codebook_size (int): Codebook size.
        codebook_dim (int): Codebook dimension. We suggest use codebook_dim = input_dim
            if use codebook_type == "euclidean", otherwise, if you want to use
            factorized vector quantization, use codebook_dim as small number (e.g. 8 or 32).
        commitment (float): Weight for commitment loss.
        use_l2_normlize (bool): Whether to use l2 normlized codes for factorized vecotor quantization,
            we suggest use it as True if you want to use factorized vector quantization
        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
        kmeans_iters (int): Number of iterations used for kmeans initialization.
        decay (float): Decay for exponential moving average over the codebooks.
        epsilon (float): Epsilon value for numerical stability.
        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
            that have an exponential moving average cluster size less than the specified threshold with
            randomly selected vector from the current batch.
    """

    def __init__(
        self,
        input_dim,
        codebook_size,
        codebook_dim,
        commitment=0.005,
        codebook_loss_weight=1.0,
        use_l2_normlize=False,
        codebook_type="euclidean",  # "euclidean" or "simple"
        kmeans_init=False,
        kmeans_iters=10,
        decay=0.8,
        eps=1e-5,
        threshold_ema_dead_code=2,
        weight_init=False,
    ):
        super().__init__()
        self.input_dim = input_dim
        self.codebook_size = codebook_size
        self.codebook_dim = codebook_dim
        self.commitment = commitment
        self.codebook_loss_weight = codebook_loss_weight
        self.use_l2_normlize = use_l2_normlize
        self.codebook_type = codebook_type
        self.kmeans_init = kmeans_init
        self.kmeans_iters = kmeans_iters
        self.decay = decay
        self.eps = eps
        self.threshold_ema_dead_code = threshold_ema_dead_code
        self.weight_init = weight_init

        if self.input_dim != self.codebook_dim:
            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
            self.out_project = WNConv1d(
                self.codebook_dim, self.input_dim, kernel_size=1
            )

        else:
            self.in_project = nn.Identity()
            self.out_project = nn.Identity()

        if self.codebook_type == "euclidean":
            self.codebook = EuclideanCodebook(
                self.codebook_dim,
                codebook_size=self.codebook_size,
                kmeans_init=self.kmeans_init,
                kmeans_iters=self.kmeans_iters,
                decay=self.decay,
                eps=self.eps,
                threshold_ema_dead_code=self.threshold_ema_dead_code,
                weight_init=self.weight_init,
            )
        elif self.codebook_type == "simple":
            self.codebook = SimpleCodebook(
                self.codebook_dim,
                codebook_size=self.codebook_size,
                use_l2_normlize=self.use_l2_normlize,
            )
        else:
            raise NotImplementedError(
                f"codebook_type {self.codebook_type} is not implemented!"
            )

    def forward(self, z):
        """
        Parameters
        ----------
        z: torch.Tensor[B x D x T]

        Returns
        -------
        z_q: torch.Tensor[B x D x T]
            Quantized continuous representation of input
        commit_loss: Tensor[B]
            Commitment loss to train encoder to predict vectors closer to codebook entries
        codebook_loss: Tensor[B]
            Codebook loss to update the codebook
        indices: torch.Tensor[B x T]
            Codebook indices (quantized discrete representation of input)
        z_e: torch.Tensor[B x D x T]
            Projected latents (continuous representation of input before quantization)
        """

        # Factorized codes project input into low-dimensional space if self.input_dim != self.codebook_dim
        z_e = self.in_project(z)
        z_q, indices = self.decode_latents(z_e)

        # Compute commitment loss and codebook loss
        if self.training:
            commit_loss = (
                F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
                * self.commitment
            )
            codebook_loss = (
                F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
                * self.codebook_loss_weight
            )
        else:
            commit_loss = torch.zeros(z.shape[0], device=z.device)
            codebook_loss = torch.zeros(z.shape[0], device=z.device)

        z_q = z_e + (z_q - z_e).detach()

        z_q = self.out_project(z_q)

        return z_q, commit_loss, codebook_loss, indices, z_e

    def decode_latents(self, latents):
        encodings = rearrange(latents, "b d t -> b t d")
        z_q, indices = self.codebook(encodings)
        z_q = z_q.transpose(1, 2)
        return z_q, indices

    def vq2emb(self, vq, out_proj=True):
        emb = self.codebook.vq2emb(vq)
        emb = emb.transpose(1, 2)
        if out_proj:
            emb = self.out_project(emb)
        return emb

    def latent2dist(self, latents):
        latents = rearrange(latents, "b d t -> b t d")
        dist, embed_ind, quantize = self.codebook.latent2dist(latents)
        return dist, embed_ind, quantize.transpose(1, 2)