xinference 1.9.1__py3-none-any.whl → 1.10.1__py3-none-any.whl

xinference/thirdparty/indextts/utils/maskgct/models/codec/amphion_codec/codec.py

@@ -0,0 +1,427 @@
+# 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 math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+from indextts.utils.maskgct.models.codec.amphion_codec.quantize import (
+    ResidualVQ,
+    VectorQuantize,
+    FactorizedVectorQuantize,
+    LookupFreeQuantize,
+)
+
+from indextts.utils.maskgct.models.codec.amphion_codec.vocos import Vocos
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+    if isinstance(m, nn.Linear):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            WNConv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class CodecEncoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        up_ratios: list = [4, 5, 5, 6],
+        out_channels: int = 256,
+        use_tanh: bool = False,
+        cfg=None,
+    ):
+        super().__init__()
+
+        d_model = cfg.d_model if cfg is not None else d_model
+        up_ratios = cfg.up_ratios if cfg is not None else up_ratios
+        out_channels = cfg.out_channels if cfg is not None else out_channels
+        use_tanh = cfg.use_tanh if cfg is not None else use_tanh
+
+        # Create first convolution
+        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in up_ratios:
+            d_model *= 2
+            self.block += [EncoderBlock(d_model, stride=stride)]
+
+        # Create last convolution
+        self.block += [
+            Snake1d(d_model),
+            WNConv1d(d_model, out_channels, kernel_size=3, padding=1),
+        ]
+
+        if use_tanh:
+            self.block += [nn.Tanh()]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        return self.block(x)
+
+    def reset_parameters(self):
+        self.apply(init_weights)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim),
+            WNConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=stride // 2 + stride % 2,
+                output_padding=stride % 2,
+            ),
+            ResidualUnit(output_dim, dilation=1),
+            ResidualUnit(output_dim, dilation=3),
+            ResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class CodecDecoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 256,
+        upsample_initial_channel: int = 1536,
+        up_ratios: list = [5, 5, 4, 2],
+        num_quantizers: int = 8,
+        codebook_size: int = 1024,
+        codebook_dim: int = 256,
+        quantizer_type: str = "vq",
+        quantizer_dropout: float = 0.5,
+        commitment: float = 0.25,
+        codebook_loss_weight: float = 1.0,
+        use_l2_normlize: bool = False,
+        codebook_type: str = "euclidean",
+        kmeans_init: bool = False,
+        kmeans_iters: int = 10,
+        decay: float = 0.8,
+        eps: float = 1e-5,
+        threshold_ema_dead_code: int = 2,
+        weight_init: bool = False,
+        use_vocos: bool = False,
+        vocos_dim: int = 384,
+        vocos_intermediate_dim: int = 1152,
+        vocos_num_layers: int = 8,
+        n_fft: int = 800,
+        hop_size: int = 200,
+        padding: str = "same",
+        cfg=None,
+    ):
+        super().__init__()
+
+        in_channels = (
+            cfg.in_channels
+            if cfg is not None and hasattr(cfg, "in_channels")
+            else in_channels
+        )
+        upsample_initial_channel = (
+            cfg.upsample_initial_channel
+            if cfg is not None and hasattr(cfg, "upsample_initial_channel")
+            else upsample_initial_channel
+        )
+        up_ratios = (
+            cfg.up_ratios
+            if cfg is not None and hasattr(cfg, "up_ratios")
+            else up_ratios
+        )
+        num_quantizers = (
+            cfg.num_quantizers
+            if cfg is not None and hasattr(cfg, "num_quantizers")
+            else num_quantizers
+        )
+        codebook_size = (
+            cfg.codebook_size
+            if cfg is not None and hasattr(cfg, "codebook_size")
+            else codebook_size
+        )
+        codebook_dim = (
+            cfg.codebook_dim
+            if cfg is not None and hasattr(cfg, "codebook_dim")
+            else codebook_dim
+        )
+        quantizer_type = (
+            cfg.quantizer_type
+            if cfg is not None and hasattr(cfg, "quantizer_type")
+            else quantizer_type
+        )
+        quantizer_dropout = (
+            cfg.quantizer_dropout
+            if cfg is not None and hasattr(cfg, "quantizer_dropout")
+            else quantizer_dropout
+        )
+        commitment = (
+            cfg.commitment
+            if cfg is not None and hasattr(cfg, "commitment")
+            else commitment
+        )
+        codebook_loss_weight = (
+            cfg.codebook_loss_weight
+            if cfg is not None and hasattr(cfg, "codebook_loss_weight")
+            else codebook_loss_weight
+        )
+        use_l2_normlize = (
+            cfg.use_l2_normlize
+            if cfg is not None and hasattr(cfg, "use_l2_normlize")
+            else use_l2_normlize
+        )
+        codebook_type = (
+            cfg.codebook_type
+            if cfg is not None and hasattr(cfg, "codebook_type")
+            else codebook_type
+        )
+        kmeans_init = (
+            cfg.kmeans_init
+            if cfg is not None and hasattr(cfg, "kmeans_init")
+            else kmeans_init
+        )
+        kmeans_iters = (
+            cfg.kmeans_iters
+            if cfg is not None and hasattr(cfg, "kmeans_iters")
+            else kmeans_iters
+        )
+        decay = cfg.decay if cfg is not None and hasattr(cfg, "decay") else decay
+        eps = cfg.eps if cfg is not None and hasattr(cfg, "eps") else eps
+        threshold_ema_dead_code = (
+            cfg.threshold_ema_dead_code
+            if cfg is not None and hasattr(cfg, "threshold_ema_dead_code")
+            else threshold_ema_dead_code
+        )
+        weight_init = (
+            cfg.weight_init
+            if cfg is not None and hasattr(cfg, "weight_init")
+            else weight_init
+        )
+        use_vocos = (
+            cfg.use_vocos
+            if cfg is not None and hasattr(cfg, "use_vocos")
+            else use_vocos
+        )
+        vocos_dim = (
+            cfg.vocos_dim
+            if cfg is not None and hasattr(cfg, "vocos_dim")
+            else vocos_dim
+        )
+        vocos_intermediate_dim = (
+            cfg.vocos_intermediate_dim
+            if cfg is not None and hasattr(cfg, "vocos_intermediate_dim")
+            else vocos_intermediate_dim
+        )
+        vocos_num_layers = (
+            cfg.vocos_num_layers
+            if cfg is not None and hasattr(cfg, "vocos_num_layers")
+            else vocos_num_layers
+        )
+        n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft
+        hop_size = (
+            cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size
+        )
+        padding = (
+            cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding
+        )
+
+        if quantizer_type == "vq":
+            self.quantizer = ResidualVQ(
+                input_dim=in_channels,
+                num_quantizers=num_quantizers,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_type=quantizer_type,
+                quantizer_dropout=quantizer_dropout,
+                commitment=commitment,
+                codebook_loss_weight=codebook_loss_weight,
+                use_l2_normlize=use_l2_normlize,
+                codebook_type=codebook_type,
+                kmeans_init=kmeans_init,
+                kmeans_iters=kmeans_iters,
+                decay=decay,
+                eps=eps,
+                threshold_ema_dead_code=threshold_ema_dead_code,
+                weight_init=weight_init,
+            )
+        elif quantizer_type == "fvq":
+            self.quantizer = ResidualVQ(
+                input_dim=in_channels,
+                num_quantizers=num_quantizers,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_type=quantizer_type,
+                quantizer_dropout=quantizer_dropout,
+                commitment=commitment,
+                codebook_loss_weight=codebook_loss_weight,
+                use_l2_normlize=use_l2_normlize,
+            )
+        elif quantizer_type == "lfq":
+            self.quantizer = ResidualVQ(
+                input_dim=in_channels,
+                num_quantizers=num_quantizers,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_type=quantizer_type,
+            )
+        else:
+            raise ValueError(f"Unknown quantizer type {quantizer_type}")
+
+        if not use_vocos:
+            # Add first conv layer
+            channels = upsample_initial_channel
+            layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)]
+
+            # Add upsampling + MRF blocks
+            for i, stride in enumerate(up_ratios):
+                input_dim = channels // 2**i
+                output_dim = channels // 2 ** (i + 1)
+                layers += [DecoderBlock(input_dim, output_dim, stride)]
+
+            # Add final conv layer
+            layers += [
+                Snake1d(output_dim),
+                WNConv1d(output_dim, 1, kernel_size=7, padding=3),
+                nn.Tanh(),
+            ]
+
+            self.model = nn.Sequential(*layers)
+
+        if use_vocos:
+            self.model = Vocos(
+                input_channels=in_channels,
+                dim=vocos_dim,
+                intermediate_dim=vocos_intermediate_dim,
+                num_layers=vocos_num_layers,
+                adanorm_num_embeddings=None,
+                n_fft=n_fft,
+                hop_size=hop_size,
+                padding=padding,
+            )
+
+        self.reset_parameters()
+
+    def forward(self, x=None, vq=False, eval_vq=False, n_quantizers=None):
+        """
+        if vq is True, x = encoder output, then return quantized output;
+        else, x = quantized output, then return decoder output
+        """
+        if vq is True:
+            if eval_vq:
+                self.quantizer.eval()
+            (
+                quantized_out,
+                all_indices,
+                all_commit_losses,
+                all_codebook_losses,
+                all_quantized,
+            ) = self.quantizer(x, n_quantizers=n_quantizers)
+            return (
+                quantized_out,
+                all_indices,
+                all_commit_losses,
+                all_codebook_losses,
+                all_quantized,
+            )
+
+        return self.model(x)
+
+    def quantize(self, x, n_quantizers=None):
+        self.quantizer.eval()
+        quantized_out, vq, _, _, _ = self.quantizer(x, n_quantizers=n_quantizers)
+        return quantized_out, vq
+
+    # TODO: check consistency of vq2emb and quantize
+    def vq2emb(self, vq, n_quantizers=None):
+        return self.quantizer.vq2emb(vq, n_quantizers=n_quantizers)
+
+    def decode(self, x):
+        return self.model(x)
+
+    def latent2dist(self, x, n_quantizers=None):
+        return self.quantizer.latent2dist(x, n_quantizers=n_quantizers)
+
+    def reset_parameters(self):
+        self.apply(init_weights)

xinference/thirdparty/indextts/utils/maskgct/models/codec/amphion_codec/quantize/__init__.py

@@ -0,0 +1,11 @@
+# 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.
+
+from indextts.utils.maskgct.models.codec.amphion_codec.quantize.factorized_vector_quantize import (
+    FactorizedVectorQuantize,
+)
+from indextts.utils.maskgct.models.codec.amphion_codec.quantize.vector_quantize import VectorQuantize
+from indextts.utils.maskgct.models.codec.amphion_codec.quantize.lookup_free_quantize import LookupFreeQuantize
+from indextts.utils.maskgct.models.codec.amphion_codec.quantize.residual_vq import ResidualVQ

xinference/thirdparty/indextts/utils/maskgct/models/codec/amphion_codec/quantize/factorized_vector_quantize.py

@@ -0,0 +1,150 @@
+# 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
+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))
+
+
+class FactorizedVectorQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+        commitment=0.005,
+        codebook_loss_weight=1.0,
+        use_l2_normlize=True,
+    ):
+        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
+
+        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()
+
+        self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim)
+
+    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 embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight
+
+        # L2 normalize encodings and codebook
+        if self.use_l2_normlize:
+            encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return z_q, indices
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = self.decode_code(vq)
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb
+
+    def latent2dist(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight
+
+        # L2 normalize encodings and codebook
+        if self.use_l2_normlize:
+            encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )  # (b*t, k)
+
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        dist = rearrange(dist, "(b t) k -> b t k", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return -dist, indices, z_q

xinference/thirdparty/indextts/utils/maskgct/models/codec/amphion_codec/quantize/lookup_free_quantize.py

@@ -0,0 +1,77 @@
+# 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
+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))
+
+
+class LookupFreeQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        assert 2**codebook_dim == codebook_size
+
+        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()
+
+    def forward(self, z):
+        z_e = self.in_project(z)
+        z_e = F.sigmoid(z_e)
+
+        z_q = z_e + (torch.round(z_e) - z_e).detach()
+
+        z_q = self.out_project(z_q)
+
+        commit_loss = torch.zeros(z.shape[0], device=z.device)
+        codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        bits = (
+            2
+            ** torch.arange(self.codebook_dim, device=z.device)
+            .unsqueeze(0)
+            .unsqueeze(-1)
+            .long()
+        )  # (1, d, 1)
+        indices = (torch.round(z_e.clone().detach()).long() * bits).sum(1).long()
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = torch.zeros(
+            vq.shape[0], self.codebook_dim, vq.shape[-1], device=vq.device
+        )  # (B, d, T)
+        for i in range(self.codebook_dim):
+            emb[:, i, :] = (vq % 2).float()
+            vq = vq // 2
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb