Initial commit.

app.py

@@ -0,0 +1,36 @@
+import gradio as gr
+import os
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["USER"] = "imagecraft"
+
+
+import gradio as gr
+from src.model.modules.imagecraft import ImageCraft
+
+model = ImageCraft.from_pretrained("nsandiman/imagecraft-ft-co-224")
+
+
+def imagecraft_interface(image_path):
+    """Process image inputs and generate audio response."""
+    transcript, audio_buffer = model.generate(image_path, output_type="buffer")
+
+    return audio_buffer, transcript
+
+
+# Define Gradio interface
+gradio_interface = gr.Interface(
+    fn=imagecraft_interface,
+    inputs=[
+        gr.Image(type="filepath", label="Upload an image"),
+        gr.Textbox(label="Reference Text (for evaluation)"),
+    ],
+    outputs=[gr.Audio(label="Speech"), gr.Textbox(label="Transcript")],
+    title="ImageCraft",
+    description="Upload an image and get the speech responses.",
+    allow_flagging="never",
+)
+
+# Launch the Gradio app
+gradio_interface.launch(share=True)

config.json

@@ -0,0 +1,58 @@
+{
+  "_name_or_path": "nsandiman/imagecraft-ft-co-224",
+  "_vocab_size": 257216,
+  "architectures": [
+    "PaliGemmaForConditionalGeneration"
+  ],
+  "bos_token_id": 2,
+  "eos_token_id": 1,
+  "hidden_size": 2048,
+  "ignore_index": -100,
+  "image_token_index": 257152,
+  "model_type": "paligemma",
+  "pad_token_id": 0,
+  "projection_dim": 2048,
+  "text_config": {
+    "hidden_size": 2048,
+    "intermediate_size": 16384,
+    "model_type": "gemma",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 18,
+    "num_image_tokens": 256,
+    "num_key_value_heads": 1,
+    "torch_dtype": "float32",
+    "vocab_size": 257216
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.0.dev0",
+  "vision_config": {
+    "hidden_size": 1152,
+    "intermediate_size": 4304,
+    "model_type": "siglip_vision_model",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 27,
+    "num_image_tokens": 256,
+    "patch_size": 14,
+    "projection_dim": 2048,
+    "projector_hidden_act": "gelu_fast",
+    "vision_use_head": false
+  },
+  "vocab_size": 257216,
+  "voicecraft_config": {
+    "model_name": "330M_TTSEnhanced.pth",
+    "encoded": "encodec_4cb2048_giga.th",
+    "voice_audio_path": "84_121550_000074_000000.wav",
+    "voice_audio_transcript": "But when I had approached so near to them The common object, which the sense deceives, Lost not by distance any of its marks",
+    "top_k": 0,
+    "top_p": 0.9,
+    "temperature": 1,
+    "kvcache": 1,
+    "codec_sr": 50,
+    "codec_audio_sr": 16000,
+    "silence_tokens": [1388, 1898, 131],
+    "stop_repetition": 3,
+    "sample_batch_size": 2,
+    "seed": 1,
+    "cut_off_sec": 7.87
+  }
+}

config.yaml

@@ -0,0 +1,14 @@
+model_flickr: nsandiman/imagecraft-pt-fk-224
+model_coco: models/imagecraft-pt-co-224
+model_tiny: models/imagecraft-pt-ty-224
+checkpoint_dir: models/checkpoint
+pretrained_dir: models/pretrained
+model_dir: models
+data:
+  raw_dir: data/raw
+  interim_dir: data/interim
+  processed_dir: data/processed
+  log_dir: data/logs
+  wandb_dir: data/wandb
+  tensorboard_log_dir: data/tensorboard/logs
+

media/voicecraft/voices/mfa_alignments/84_121550_000074_000000.csv

@@ -0,0 +1,109 @@
+Begin,End,Label,Type,Speaker
+0.03,0.18,but,words,temp
+0.18,0.32,when,words,temp
+0.32,0.48,i,words,temp
+0.48,0.64,had,words,temp
+0.64,1.19,approached,words,temp
+1.22,1.58,so,words,temp
+1.58,1.91,near,words,temp
+1.91,2.07,to,words,temp
+2.07,2.42,them,words,temp
+2.53,2.61,the,words,temp
+2.61,3.01,common,words,temp
+3.05,3.62,object,words,temp
+3.68,3.93,which,words,temp
+3.93,4.02,the,words,temp
+4.02,4.34,sense,words,temp
+4.34,4.97,deceives,words,temp
+5.04,5.54,lost,words,temp
+5.54,6.0,not,words,temp
+6.0,6.14,by,words,temp
+6.14,6.67,distance,words,temp
+6.79,7.05,any,words,temp
+7.05,7.18,of,words,temp
+7.18,7.34,its,words,temp
+7.34,7.87,marks,words,temp
+0.03,0.06,B,phones,temp
+0.06,0.09,AH1,phones,temp
+0.09,0.18,T,phones,temp
+0.18,0.23,W,phones,temp
+0.23,0.27,EH1,phones,temp
+0.27,0.32,N,phones,temp
+0.32,0.48,AY1,phones,temp
+0.48,0.49,HH,phones,temp
+0.49,0.6,AE1,phones,temp
+0.6,0.64,D,phones,temp
+0.64,0.7,AH0,phones,temp
+0.7,0.83,P,phones,temp
+0.83,0.88,R,phones,temp
+0.88,0.99,OW1,phones,temp
+0.99,1.12,CH,phones,temp
+1.12,1.19,T,phones,temp
+1.22,1.4,S,phones,temp
+1.4,1.58,OW1,phones,temp
+1.58,1.7,N,phones,temp
+1.7,1.84,IH1,phones,temp
+1.84,1.91,R,phones,temp
+1.91,2.01,T,phones,temp
+2.01,2.07,AH0,phones,temp
+2.07,2.13,DH,phones,temp
+2.13,2.3,EH1,phones,temp
+2.3,2.42,M,phones,temp
+2.53,2.55,DH,phones,temp
+2.55,2.61,AH0,phones,temp
+2.61,2.73,K,phones,temp
+2.73,2.85,AA1,phones,temp
+2.85,2.9,M,phones,temp
+2.9,2.95,AH0,phones,temp
+2.95,3.01,N,phones,temp
+3.05,3.22,AA1,phones,temp
+3.22,3.27,B,phones,temp
+3.27,3.34,JH,phones,temp
+3.34,3.48,EH0,phones,temp
+3.48,3.54,K,phones,temp
+3.54,3.62,T,phones,temp
+3.68,3.69,HH,phones,temp
+3.69,3.76,W,phones,temp
+3.76,3.8,IH1,phones,temp
+3.8,3.93,CH,phones,temp
+3.93,3.95,DH,phones,temp
+3.95,4.02,AH0,phones,temp
+4.02,4.12,S,phones,temp
+4.12,4.21,EH1,phones,temp
+4.21,4.27,N,phones,temp
+4.27,4.34,S,phones,temp
+4.34,4.42,D,phones,temp
+4.42,4.45,IH0,phones,temp
+4.45,4.59,S,phones,temp
+4.59,4.79,IY1,phones,temp
+4.79,4.87,V,phones,temp
+4.87,4.97,Z,phones,temp
+5.04,5.12,L,phones,temp
+5.12,5.33,AO1,phones,temp
+5.33,5.42,S,phones,temp
+5.42,5.54,T,phones,temp
+5.54,5.7,N,phones,temp
+5.7,5.89,AA1,phones,temp
+5.89,6.0,T,phones,temp
+6.0,6.05,B,phones,temp
+6.05,6.14,AY1,phones,temp
+6.14,6.24,D,phones,temp
+6.24,6.3,IH1,phones,temp
+6.3,6.38,S,phones,temp
+6.38,6.45,T,phones,temp
+6.45,6.51,AH0,phones,temp
+6.51,6.57,N,phones,temp
+6.57,6.67,S,phones,temp
+6.79,6.89,EH1,phones,temp
+6.89,6.95,N,phones,temp
+6.95,7.05,IY0,phones,temp
+7.05,7.13,AH0,phones,temp
+7.13,7.18,V,phones,temp
+7.18,7.22,IH0,phones,temp
+7.22,7.29,T,phones,temp
+7.29,7.34,S,phones,temp
+7.34,7.39,M,phones,temp
+7.39,7.5,AA1,phones,temp
+7.5,7.58,R,phones,temp
+7.58,7.7,K,phones,temp
+7.7,7.87,S,phones,temp

packages.txt

@@ -0,0 +1,14 @@
+espeak-ng
+espeak
+espeak-data
+libespeak1
+libespeak-dev
+festival*
+build-essential
+flac 
+libasound2-dev
+libsndfile1-dev
+vorbis-tools
+libxml2-dev
+libxslt-dev
+zlib1g-dev

requirements.txt

@@ -0,0 +1,13 @@
+-e git+https://github.com/facebookresearch/audiocraft.git@f83babff6b5e97f75562127c4cc8122229c8f099#egg=audiocraft
+git+https://github.com/huggingface/transformers.git
+phonemizer
+spaces
+huggingface-hub
+num2words
+numpy
+pillow
+safetensors
+tokenizers
+torchaudio
+torchvision
+aeneas

setup.py

@@ -0,0 +1,16 @@
+import platform
+from setuptools import setup, find_packages
+
+if platform.python_version_tuple()[:2] != ("3", "11"):
+    raise RuntimeError("Python version 3.11 required")
+
+setup(
+    name="distilvit",
+    version="0.1",
+    packages=find_packages(),
+    entry_points={
+        "console_scripts": [
+            "train=src.model.train:main",  # "main" is a function in "train_model.py"
+        ],
+    },
+)

src/model/modules/activation.py

@@ -0,0 +1,638 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/activation.py, modified by Puyuan Peng, 2024
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+from torch.nn import Linear, Module
+from torch.nn import functional as F
+from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
+from torch.nn.modules.linear import NonDynamicallyQuantizableLinear
+from torch.nn.parameter import Parameter
+import logging
+from typing import Callable, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from torch.types import _dtype as DType
+else:
+    # The JIT doesn't understand Union, nor torch.dtype here
+    DType = int
+
+def _canonical_mask(
+        mask: Optional[Tensor],
+        mask_name: str,
+        other_type: Optional[DType],
+        other_name: str,
+        target_type: DType,
+        check_other: bool = True,
+) -> Optional[Tensor]:
+
+    if mask is not None:
+        _mask_dtype = mask.dtype
+        _mask_is_float = torch.is_floating_point(mask)
+        if _mask_dtype != torch.bool and not _mask_is_float:
+            raise AssertionError(
+                f"only bool and floating types of {mask_name} are supported")
+        if check_other and other_type is not None:
+            if _mask_dtype != other_type:
+                warnings.warn(
+                    f"Support for mismatched {mask_name} and {other_name} "
+                    "is deprecated. Use same type for both instead."
+                )
+        if not _mask_is_float:
+            mask = (
+                torch.zeros_like(mask, dtype=target_type)
+                .masked_fill_(mask, float("-inf"))
+            )
+    return mask
+
+def _in_projection_packed(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w: Tensor,
+    b: Optional[Tensor] = None,
+) -> List[Tensor]:
+    r"""
+    Performs the in-projection step of the attention operation, using packed weights.
+    Output is a triple containing projection tensors for query, key and value.
+    Args:
+        q, k, v: query, key and value tensors to be projected. For self-attention,
+            these are typically the same tensor; for encoder-decoder attention,
+            k and v are typically the same tensor. (We take advantage of these
+            identities for performance if they are present.) Regardless, q, k and v
+            must share a common embedding dimension; otherwise their shapes may vary.
+        w: projection weights for q, k and v, packed into a single tensor. Weights
+            are packed along dimension 0, in q, k, v order.
+        b: optional projection biases for q, k and v, packed into a single tensor
+            in q, k, v order.
+    Shape:
+        Inputs:
+        - q: :math:`(..., E)` where E is the embedding dimension
+        - k: :math:`(..., E)` where E is the embedding dimension
+        - v: :math:`(..., E)` where E is the embedding dimension
+        - w: :math:`(E * 3, E)` where E is the embedding dimension
+        - b: :math:`E * 3` where E is the embedding dimension
+        Output:
+        - in output list :math:`[q', k', v']`, each output tensor will have the
+            same shape as the corresponding input tensor.
+    """
+    E = q.size(-1)
+    if k is v:
+        if q is k:
+            # self-attention
+            proj = F.linear(q, w, b)
+            # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()
+            proj = proj.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return proj[0], proj[1], proj[2]
+        else:
+            # encoder-decoder attention
+            w_q, w_kv = w.split([E, E * 2])
+            if b is None:
+                b_q = b_kv = None
+            else:
+                b_q, b_kv = b.split([E, E * 2])
+            q_proj = F.linear(q, w_q, b_q)
+            kv_proj = F.linear(k, w_kv, b_kv)
+            # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()
+            kv_proj = kv_proj.unflatten(-1, (2, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return (q_proj, kv_proj[0], kv_proj[1])
+    else:
+        w_q, w_k, w_v = w.chunk(3)
+        if b is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = b.chunk(3)
+        return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)
+
+def _none_or_dtype(input: Optional[Tensor]) -> Optional[DType]:
+    if input is None:
+        return None
+    elif isinstance(input, torch.Tensor):
+        return input.dtype
+    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces as described in the paper:
+    `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.
+    Multi-Head Attention is defined as:
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+    ``forward()`` will use a special optimized implementation if all of the following
+    conditions are met:
+    - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor. This
+      restriction will be loosened in the future.)
+    - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``
+    - training is disabled (using ``.eval()``)
+    - dropout is 0
+    - ``add_bias_kv`` is ``False``
+    - ``add_zero_attn`` is ``False``
+    - ``batch_first`` is ``True`` and the input is batched
+    - ``kdim`` and ``vdim`` are equal to ``embed_dim``
+    - at most one of ``key_padding_mask`` or ``attn_mask`` is passed
+    - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``
+      nor ``attn_mask`` is passed
+    If the optimized implementation is in use, a
+    `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for
+    ``query``/``key``/``value`` to represent padding more efficiently than using a
+    padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_
+    will be returned, and an additional speedup proportional to the fraction of the input
+    that is padding can be expected.
+    Args:
+        embed_dim: Total dimension of the model.
+        num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split
+            across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).
+        dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).
+        bias: If specified, adds bias to input / output projection layers. Default: ``True``.
+        add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.
+        add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.
+            Default: ``False``.
+        kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
+        vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+    Examples::
+        >>> # xdoctest: +SKIP
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+    """
+    __constants__ = ["batch_first"]
+    bias_k: Optional[torch.Tensor]
+    bias_v: Optional[torch.Tensor]
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        kdim=None,
+        vdim=None,
+        batch_first=False,
+        linear1_cls=Linear,
+        linear2_cls=Linear,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = (
+            self.kdim == embed_dim and self.vdim == embed_dim
+        )
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.batch_first = batch_first
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+
+        if add_bias_kv:
+            self.bias_k = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+            self.bias_v = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+        else:
+            self.bias_k = self.bias_v = None
+
+        if linear1_cls == Linear:
+            if not self._qkv_same_embed_dim:
+                self.q_proj_weight = Parameter(
+                    torch.empty((embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.k_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.kdim), **factory_kwargs)
+                )
+                self.v_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.vdim), **factory_kwargs)
+                )
+                self.register_parameter("in_proj_weight", None)
+            else:
+                # go down this route with voicecraft
+                self.in_proj_weight = Parameter(
+                    torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+            if bias: # True by default
+                self.in_proj_bias = Parameter(
+                    torch.empty(3 * embed_dim, **factory_kwargs)
+                )
+            else:
+                self.register_parameter("in_proj_bias", None)
+            self.out_proj = NonDynamicallyQuantizableLinear(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            self._reset_parameters()
+        else:
+            if not self._qkv_same_embed_dim:
+                raise NotImplementedError
+            else:
+                self.in_proj_linear = linear1_cls(
+                    embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs
+                )
+                self.in_proj_weight = self.in_proj_linear.weight
+
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+                if bias:
+                    self.in_proj_bias = self.in_proj_linear.bias
+                else:
+                    self.register_parameter("in_proj_bias", None)
+
+            self.out_proj = linear2_cls(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            if self.bias_k is not None:
+                xavier_normal_(self.bias_k)
+            if self.bias_v is not None:
+                xavier_normal_(self.bias_v)
+
+        self.add_zero_attn = add_zero_attn
+
+    def _reset_parameters(self):
+        if self._qkv_same_embed_dim:
+            xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.0)
+            constant_(self.out_proj.bias, 0.0)
+
+        if self.bias_k is not None:
+            xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            xavier_normal_(self.bias_v)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if "_qkv_same_embed_dim" not in state:
+            state["_qkv_same_embed_dim"] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        key_padding_mask: Optional[Tensor] = None,
+        need_weights: bool = True,
+        attn_mask: Optional[Tensor] = None,
+        average_attn_weights: bool = True,
+        past: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        r"""
+        Args:
+            query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``
+                or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,
+                :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.
+                Queries are compared against key-value pairs to produce the output.
+                See "Attention Is All You Need" for more details.
+            key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``
+                or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,
+                :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.
+                See "Attention Is All You Need" for more details.
+            value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when
+                ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source
+                sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.
+                See "Attention Is All You Need" for more details.
+            key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``
+                to ignore for the purpose of attention (i.e. treat as "padding"). For unbatched `query`, shape should be :math:`(S)`.
+                Binary and byte masks are supported.
+                For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for
+                the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.
+            need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.
+                Default: ``True``.
+            attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
+                :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
+                :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be
+                broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
+                Binary, byte, and float masks are supported. For a binary mask, a ``True`` value indicates that the
+                corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the
+                corresponding position is not allowed to attend. For a float mask, the mask values will be added to
+                the attention weight.
+            average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
+                heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
+                effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)
+        Outputs:
+            - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,
+              :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,
+              where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the
+              embedding dimension ``embed_dim``.
+            - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,
+              returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+              :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+              :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+              head of shape :math:`(\text{num\_heads}, L, S)` when input is unbatched or :math:`(N, \text{num\_heads}, L, S)`.
+            .. note::
+                `batch_first` argument is ignored for unbatched inputs.
+        """
+        is_batched = query.dim() == 3
+        if key_padding_mask is not None:
+            _kpm_dtype = key_padding_mask.dtype
+            if _kpm_dtype != torch.bool and not torch.is_floating_point(
+                key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        why_not_fast_path = ""
+        if not is_batched:
+            why_not_fast_path = f"input not batched; expected query.dim() of 3 but got {query.dim()}"
+        elif query is not key or key is not value:
+            # When lifting this restriction, don't forget to either
+            # enforce that the dtypes all match or test cases where
+            # they don't!
+            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
+        elif (
+            self.in_proj_bias is not None
+            and query.dtype != self.in_proj_bias.dtype
+        ):
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
+        elif (
+            self.in_proj_weight is not None
+            and query.dtype != self.in_proj_weight.dtype
+        ):
+            # this case will fail anyway, but at least they'll get a useful error message.
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
+        elif self.training:
+            why_not_fast_path = "training is enabled"
+        elif not self.batch_first:
+            why_not_fast_path = "batch_first was not True"
+        elif self.bias_k is not None:
+            why_not_fast_path = "self.bias_k was not None"
+        elif self.bias_v is not None:
+            why_not_fast_path = "self.bias_v was not None"
+        elif self.dropout:
+            why_not_fast_path = f"dropout was {self.dropout}, required zero"
+        elif self.add_zero_attn:
+            why_not_fast_path = "add_zero_attn was enabled"
+        elif not self._qkv_same_embed_dim:
+            why_not_fast_path = "_qkv_same_embed_dim was not True"
+        elif attn_mask is not None:
+            why_not_fast_path = "attn_mask was not None"
+        elif query.is_nested and key_padding_mask is not None:
+            why_not_fast_path = (
+                "key_padding_mask is not supported with NestedTensor input"
+            )
+        elif self.num_heads % 2 == 1:
+            why_not_fast_path = "num_heads is odd"
+        elif torch.is_autocast_enabled():
+            why_not_fast_path = "autocast is enabled"
+
+        if not why_not_fast_path:
+            tensor_args = (
+                query,
+                key,
+                value,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.out_proj.weight,
+                self.out_proj.bias,
+            )
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_fast_path = "some Tensor argument has_torch_function"
+            elif not all(
+                [
+                    (x is None or x.is_cuda or "cpu" in str(x.device))
+                    for x in tensor_args
+                ]
+            ):
+                why_not_fast_path = (
+                    "some Tensor argument is neither CUDA nor CPU"
+                )
+            elif torch.is_grad_enabled() and any(
+                [x is not None and x.requires_grad for x in tensor_args]
+            ):
+                why_not_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+            if not why_not_fast_path:
+                return torch._native_multi_head_attention(
+                    query,
+                    key,
+                    value,
+                    self.embed_dim,
+                    self.num_heads,
+                    self.in_proj_weight,
+                    self.in_proj_bias,
+                    self.out_proj.weight,
+                    self.out_proj.bias,
+                    key_padding_mask
+                    if key_padding_mask is not None
+                    else attn_mask,
+                    need_weights,
+                    average_attn_weights,
+                    1
+                    if key_padding_mask is not None
+                    else 0
+                    if attn_mask is not None
+                    else None,
+                )
+
+        any_nested = query.is_nested or key.is_nested or value.is_nested
+        assert not any_nested, (
+            "MultiheadAttention does not support NestedTensor outside of its fast path. "
+            + f"The fast path was not hit because {why_not_fast_path}"
+        )
+
+        if self.batch_first and is_batched:
+            # make sure that the transpose op does not affect the "is" property
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = [x.transpose(1, 0) for x in (query, key)]
+                    value = key
+            else:
+                query, key, value = [
+                    x.transpose(1, 0) for x in (query, key, value)
+                ]
+
+        if not self._qkv_same_embed_dim:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight,
+                k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight,
+                average_attn_weights=average_attn_weights,
+            )
+        else:
+            # re-write the self.attention here, to get k, v cache
+            tgt_len, bsz, embed_dim = query.shape
+            src_len, _, _ = key.shape
+            num_heads = self.num_heads
+            key_padding_mask = _canonical_mask(
+                mask=key_padding_mask,
+                mask_name="key_padding_mask",
+                other_type=_none_or_dtype(attn_mask),
+                other_name="attn_mask",
+                target_type=query.dtype
+            )
+            attn_mask = _canonical_mask(
+                            mask=attn_mask,
+                            mask_name="attn_mask",
+                            other_type=None,
+                            other_name="",
+                            target_type=query.dtype,
+                            check_other=False,
+                            )
+            head_dim = self.embed_dim // self.num_heads
+            assert head_dim * self.num_heads == self.embed_dim, f"embed_dim {self.embed_dim} not divisible by num_heads {self.num_heads}"
+            assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
+            q, k, v = _in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)
+            # k_present, v_present = k, v
+            
+            #
+            # reshape q, k, v for multihead attention and make em batch first
+            #
+            
+            q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+            k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+            v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1) # (bsz * num_heads, src_len, head_dim)
+            src_len = k.size(1)
+            if past is not None and past.ndim > 2:
+                expected_src_len = src_len + past[0].shape[-2]
+            else:
+                expected_src_len = src_len
+
+
+            # ensure attn_mask's dim is 3
+            if attn_mask.dim() == 2:
+                correct_2d_size = (tgt_len, expected_src_len)
+                if attn_mask.shape != correct_2d_size:
+                    raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
+                attn_mask = attn_mask.unsqueeze(0)
+            elif attn_mask.dim() == 3:
+                correct_3d_size = (bsz * num_heads, tgt_len, expected_src_len)
+                if attn_mask.shape != correct_3d_size:
+                    raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
+            else:
+                raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
+            
+            if key_padding_mask is not None:
+                assert key_padding_mask.shape == (bsz, expected_src_len), \
+                    f"expecting key_padding_mask shape of {(bsz, expected_src_len)}, but got {key_padding_mask.shape}"
+                key_padding_mask = key_padding_mask.view(bsz, 1, 1, expected_src_len).   \
+                    expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, expected_src_len)
+                if attn_mask is None:
+                    attn_mask = key_padding_mask
+                else:
+                    attn_mask = attn_mask + key_padding_mask
+            
+            if not self.training:
+                dropout_p = 0.0
+            else:
+                dropout_p = self.dropout
+
+            if need_weights:
+                raise NotImplementedError("need_weights not implemented for voicecraft")
+                # B, Nt, E = q.shape
+                # q_scaled = q / math.sqrt(E)
+
+                # assert not (is_causal and attn_mask is None), "FIXME: is_causal not implemented for need_weights"
+
+                # if attn_mask is not None:
+                #     attn_output_weights = torch.baddbmm(attn_mask, q_scaled, k.transpose(-2, -1))
+                # else:
+                #     attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
+                # attn_output_weights = softmax(attn_output_weights, dim=-1)
+                # if dropout_p > 0.0:
+                #     attn_output_weights = dropout(attn_output_weights, p=dropout_p)
+
+                # attn_output = torch.bmm(attn_output_weights, v)
+
+                # attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
+                # attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+                # attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+
+                # # optionally average attention weights over heads
+                # attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+                # if average_attn_weights:
+                #     attn_output_weights = attn_output_weights.mean(dim=1)
+
+                # if not is_batched:
+                #     # squeeze the output if input was unbatched
+                #     attn_output = attn_output.squeeze(1)
+                #     attn_output_weights = attn_output_weights.squeeze(0)
+                # return attn_output, attn_output_weights
+            else:
+                # attn_mask can be either (L,S) or (N*num_heads, L, S)
+                # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
+                # in order to match the input for SDPA of (N, num_heads, L, S)
+                if attn_mask is not None:
+                    if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                        attn_mask = attn_mask.unsqueeze(0)
+                    else:
+                        attn_mask = attn_mask.view(bsz, num_heads, -1, expected_src_len)
+
+                q = q.view(bsz, num_heads, tgt_len, head_dim)
+                k = k.view(bsz, num_heads, src_len, head_dim)
+                v = v.view(bsz, num_heads, src_len, head_dim)
+                # logging.info(f"shape of past: {past.shape}")
+                if past is not None:
+                    present = torch.stack([k, v], dim=0) # (2, bsz, num_heads, src_len, head_dim)
+                    if past.ndim > 2: # this means we use kvcache, otherwise we just pass in a placeholder, but not actually using kvcache
+                        pk, pv = past
+                        k = torch.cat([pk, k], dim=-2)
+                        v = torch.cat([pv, v], dim=-2)
+                else:
+                    present = None
+                attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal=False)
+                attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
+
+                attn_output = F.linear(attn_output, self.out_proj.weight, self.out_proj.bias)
+                attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+                if not is_batched:
+                    # squeeze the output if input was unbatched
+                    attn_output = attn_output.squeeze(1)
+                # if self.training:
+                #     return attn_output, None
+                # else:
+                #     return (attn_output, present), None
+
+        # harded coded, the code do not support returning attn weigths yet
+        attn_output_weights=None
+        if self.batch_first and is_batched:
+            return attn_output.transpose(1, 0), present
+        else:
+            return attn_output, present
+

src/model/modules/codebooks_patterns.py

@@ -0,0 +1,538 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import namedtuple
+from dataclasses import dataclass
+from functools import lru_cache
+import logging
+import typing as tp
+
+from abc import ABC, abstractmethod
+import torch
+
+LayoutCoord = namedtuple('LayoutCoord', ['t', 'q'])  # (timestep, codebook index)
+PatternLayout = tp.List[tp.List[LayoutCoord]]  # Sequence of coordinates
+
+
+@dataclass
+class Pattern:
+    """Base implementation of a pattern over a sequence with multiple codebooks.
+
+    The codebook pattern consists in a layout, defining for each sequence step
+    the list of coordinates of each codebook timestep in the resulting interleaved sequence.
+    The first item of the pattern is always an empty list in order to properly insert a special token
+    to start with. For convenience, we also keep track of ``n_q`` the number of codebooks used for the pattern
+    and ``timesteps`` the number of timesteps corresponding to the original sequence.
+
+    The pattern provides convenient methods to build and revert interleaved sequences from it:
+    ``build_pattern_sequence`` maps a given a dense input tensor of multi-codebook sequence from [B, K, T]
+        to the interleaved sequence of shape [B, K, S] applying the pattern, with S being the batch size,
+        K being the number of codebooks, T the number of original timesteps and S the number of sequence steps
+        for the output sequence. The unfilled positions are replaced with a special token and the built sequence
+        is returned along with a mask indicating valid tokens.
+    ``revert_pattern_sequence`` maps back an interleaved sequence of shape [B, K, S] to the original alignment
+        of codebooks across timesteps to an output tensor of shape [B, K, T], using again a special token and a mask
+        to fill and specify invalid positions if needed.
+    See the dedicated methods for more details.
+    """
+    # Pattern layout, for each sequence step, we have a list of coordinates
+    # corresponding to the original codebook timestep and position.
+    # The first list is always an empty list in order to properly insert
+    # a special token to start with.
+    layout: PatternLayout
+    timesteps: int
+    n_q: int
+
+    def __post_init__(self):
+        assert len(self.layout) > 0
+        assert self.layout[0] == []
+        self._validate_layout()
+        self._build_reverted_sequence_scatter_indexes = lru_cache(100)(self._build_reverted_sequence_scatter_indexes)
+        self._build_pattern_sequence_scatter_indexes = lru_cache(100)(self._build_pattern_sequence_scatter_indexes)
+        # logging.info("New pattern, time steps: %d, sequence steps: %d", self.timesteps, len(self.layout))
+
+    def _validate_layout(self):
+        """Runs checks on the layout to ensure a valid pattern is defined.
+        A pattern is considered invalid if:
+            - Multiple timesteps for a same codebook are defined in the same sequence step
+            - The timesteps for a given codebook are not in ascending order as we advance in the sequence
+              (this would mean that we have future timesteps before past timesteps).
+        """
+        q_timesteps = {q: 0 for q in range(self.n_q)}
+        for s, seq_coords in enumerate(self.layout):
+            if len(seq_coords) > 0:
+                qs = set()
+                for coord in seq_coords:
+                    qs.add(coord.q)
+                    last_q_timestep = q_timesteps[coord.q]
+                    assert coord.t >= last_q_timestep, \
+                        f"Past timesteps are found in the sequence for codebook = {coord.q} at step {s}"
+                    q_timesteps[coord.q] = coord.t
+                # each sequence step contains at max 1 coordinate per codebook
+                assert len(qs) == len(seq_coords), \
+                    f"Multiple entries for a same codebook are found at step {s}"
+
+    @property
+    def num_sequence_steps(self):
+        return len(self.layout) - 1
+
+    @property
+    def max_delay(self):
+        max_t_in_seq_coords = 0
+        for seq_coords in self.layout[1:]:
+            for coords in seq_coords:
+                max_t_in_seq_coords = max(max_t_in_seq_coords, coords.t + 1)
+        return max_t_in_seq_coords - self.timesteps
+
+    @property
+    def valid_layout(self):
+        valid_step = len(self.layout) - self.max_delay
+        return self.layout[:valid_step]
+
+    def get_sequence_coords_with_timestep(self, t: int, q: tp.Optional[int] = None):
+        """Get codebook coordinates in the layout that corresponds to the specified timestep t
+        and optionally to the codebook q. Coordinates are returned as a tuple with the sequence step
+        and the actual codebook coordinates.
+        """
+        assert t <= self.timesteps, "provided timesteps is greater than the pattern's number of timesteps"
+        if q is not None:
+            assert q <= self.n_q, "provided number of codebooks is greater than the pattern's number of codebooks"
+        coords = []
+        for s, seq_codes in enumerate(self.layout):
+            for code in seq_codes:
+                if code.t == t and (q is None or code.q == q):
+                    coords.append((s, code))
+        return coords
+
+    def get_steps_with_timestep(self, t: int, q: tp.Optional[int] = None) -> tp.List[int]:
+        return [step for step, coords in self.get_sequence_coords_with_timestep(t, q)]
+
+    def get_first_step_with_timesteps(self, t: int, q: tp.Optional[int] = None) -> tp.Optional[int]:
+        steps_with_timesteps = self.get_steps_with_timestep(t, q)
+        return steps_with_timesteps[0] if len(steps_with_timesteps) > 0 else None
+
+    def _build_pattern_sequence_scatter_indexes(self, timesteps: int, n_q: int, keep_only_valid_steps: bool,
+                                                device: tp.Union[torch.device, str] = 'cpu'):
+        """Build scatter indexes corresponding to the pattern, up to the provided sequence_steps.
+
+        Args:
+            timesteps (int): Maximum number of timesteps steps to consider.
+            keep_only_valid_steps (bool): Restrict the pattern layout to match only valid steps.
+            device (Union[torch.device, str]): Device for created tensors.
+        Returns:
+            indexes (torch.Tensor): Indexes corresponding to the sequence, of shape [K, S].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes, of shape [K, S].
+        """
+        assert n_q == self.n_q, f"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}"
+        assert timesteps <= self.timesteps, "invalid number of timesteps used to build the sequence from the pattern"
+        # use the proper layout based on whether we limit ourselves to valid steps only or not,
+        # note that using the valid_layout will result in a truncated sequence up to the valid steps
+        ref_layout = self.valid_layout if keep_only_valid_steps else self.layout
+        # single item indexing being super slow with pytorch vs. numpy, so we use numpy here
+        indexes = torch.zeros(n_q, len(ref_layout), dtype=torch.long).numpy()
+        mask = torch.zeros(n_q, len(ref_layout), dtype=torch.bool).numpy()
+        # fill indexes with last sequence step value that will correspond to our special token
+        # the last value is n_q * timesteps as we have flattened z and append special token as the last token
+        # which will correspond to the index: n_q * timesteps
+        indexes[:] = n_q * timesteps
+        # iterate over the pattern and fill scattered indexes and mask
+        for s, sequence_coords in enumerate(ref_layout):
+            for coords in sequence_coords:
+                if coords.t < timesteps:
+                    indexes[coords.q, s] = coords.t + coords.q * timesteps
+                    mask[coords.q, s] = 1
+        indexes = torch.from_numpy(indexes).to(device)
+        mask = torch.from_numpy(mask).to(device)
+        return indexes, mask
+
+    def build_pattern_sequence(self, z: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):
+        """Build sequence corresponding to the pattern from the input tensor z.
+        The sequence is built using up to sequence_steps if specified, and non-pattern
+        coordinates are filled with the special token.
+
+        Args:
+            z (torch.Tensor): Input tensor of multi-codebooks sequence, of shape [B, K, T].
+            special_token (int): Special token used to fill non-pattern coordinates in the new sequence.
+            keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.
+                Steps that are beyond valid steps will be replaced by the special_token in that case.
+        Returns:
+            values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, S] with S
+                corresponding either to the sequence_steps if provided, otherwise to the length of the pattern.
+            indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, S].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, S].
+        """
+        B, K, T = z.shape
+        indexes, mask = self._build_pattern_sequence_scatter_indexes(
+            T, K, keep_only_valid_steps=keep_only_valid_steps, device=str(z.device)
+        )
+        z = z.view(B, -1)
+        # we append the special token as the last index of our flattened z tensor
+        z = torch.cat([z, torch.zeros_like(z[:, :1]) + special_token], dim=1)
+        values = z[:, indexes.view(-1)]
+        values = values.view(B, K, indexes.shape[-1])
+        return values, indexes, mask
+
+    def _build_reverted_sequence_scatter_indexes(self, sequence_steps: int, n_q: int,
+                                                 keep_only_valid_steps: bool = False,
+                                                 is_model_output: bool = False,
+                                                 device: tp.Union[torch.device, str] = 'cpu'):
+        """Builds scatter indexes required to retrieve the original multi-codebook sequence
+        from interleaving pattern.
+
+        Args:
+            sequence_steps (int): Sequence steps.
+            n_q (int): Number of codebooks.
+            keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.
+                Steps that are beyond valid steps will be replaced by the special_token in that case.
+            is_model_output (bool): Whether to keep the sequence item corresponding to initial special token or not.
+            device (Union[torch.device, str]): Device for created tensors.
+        Returns:
+            torch.Tensor: Indexes for reconstructing the output, of shape [K, T].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].
+        """
+        ref_layout = self.valid_layout if keep_only_valid_steps else self.layout
+        # TODO(jade): Do we want to further truncate to only valid timesteps here as well?
+        timesteps = self.timesteps
+        assert n_q == self.n_q, f"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}"
+        assert sequence_steps <= len(ref_layout), \
+            f"sequence to revert is longer than the defined pattern: {sequence_steps} > {len(ref_layout)}"
+
+        # ensure we take the appropriate indexes to keep the model output from the first special token as well
+        if is_model_output:
+            ref_layout = ref_layout[1:]
+
+        # single item indexing being super slow with pytorch vs. numpy, so we use numpy here
+        indexes = torch.zeros(n_q, timesteps, dtype=torch.long).numpy()
+        mask = torch.zeros(n_q, timesteps, dtype=torch.bool).numpy()
+        # fill indexes with last sequence step value that will correspond to our special token
+        indexes[:] = n_q * sequence_steps
+        for s, sequence_codes in enumerate(ref_layout):
+            if s < sequence_steps:
+                for code in sequence_codes:
+                    if code.t < timesteps:
+                        indexes[code.q, code.t] = s + code.q * sequence_steps
+                        mask[code.q, code.t] = 1
+        indexes = torch.from_numpy(indexes).to(device)
+        mask = torch.from_numpy(mask).to(device)
+        return indexes, mask
+
+    def revert_pattern_sequence(self, s: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):
+        """Revert a sequence built from the pattern back to the original multi-codebook sequence without interleaving.
+        The sequence is reverted using up to timesteps if specified, and non-pattern coordinates
+        are filled with the special token.
+
+        Args:
+            s (torch.Tensor): Interleaved sequence tensor obtained from the pattern, of shape [B, K, S].
+            special_token (int or float): Special token used to fill non-pattern coordinates in the new sequence.
+        Returns:
+            values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, T] with T
+                corresponding either to the timesteps if provided, or the total timesteps in pattern otherwise.
+            indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, T].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].
+        """
+        B, K, S = s.shape
+        indexes, mask = self._build_reverted_sequence_scatter_indexes(
+            S, K, keep_only_valid_steps, is_model_output=False, device=str(s.device)
+        )
+        s = s.view(B, -1)
+        # we append the special token as the last index of our flattened z tensor
+        s = torch.cat([s, torch.zeros_like(s[:, :1]) + special_token], dim=1)
+        values = s[:, indexes.view(-1)]
+        values = values.view(B, K, indexes.shape[-1])
+        return values, indexes, mask
+
+    def revert_pattern_logits(self, logits: torch.Tensor, special_token: float, keep_only_valid_steps: bool = False):
+        """Revert model logits obtained on a sequence built from the pattern
+        back to a tensor matching the original sequence.
+
+        This method is similar to ``revert_pattern_sequence`` with the following specificities:
+        1. It is designed to work with the extra cardinality dimension
+        2. We return the logits for the first sequence item that matches the special_token and
+        which matching target in the original sequence is the first item of the sequence,
+        while we skip the last logits as there is no matching target
+        """
+        B, card, K, S = logits.shape
+        indexes, mask = self._build_reverted_sequence_scatter_indexes(
+            S, K, keep_only_valid_steps, is_model_output=True, device=logits.device
+        )
+        logits = logits.reshape(B, card, -1)
+        # we append the special token as the last index of our flattened z tensor
+        logits = torch.cat([logits, torch.zeros_like(logits[:, :, :1]) + special_token], dim=-1)  # [B, card, K x S]
+        values = logits[:, :, indexes.view(-1)]
+        values = values.view(B, card, K, indexes.shape[-1])
+        return values, indexes, mask
+
+
+class CodebooksPatternProvider(ABC):
+    """Abstraction around providing pattern for interleaving codebooks.
+
+    The CodebooksPatternProvider abstraction allows to implement various strategies to
+    define interleaving pattern of sequences composed of multiple codebooks. For a given
+    number of codebooks `n_q`, the pattern provider can generate a specified pattern
+    corresponding to a sequence of `T` timesteps with `n_q` parallel codebooks. This pattern
+    can be used to construct a new sequence from the original codes respecting the specified
+    pattern. The pattern is defined as a list of list of code coordinates, code coordinate
+    being a tuple with the original timestep and codebook to build the new sequence.
+    Note that all patterns must start with an empty list that is then used to insert a first
+    sequence step of special tokens in the newly generated sequence.
+
+    Args:
+        n_q (int): number of codebooks.
+        cached (bool): if True, patterns for a given length are cached. In general
+            that should be true for efficiency reason to avoid synchronization points.
+    """
+    def __init__(self, n_q: int, cached: bool = True):
+        assert n_q > 0
+        self.n_q = n_q
+        self.get_pattern = lru_cache(100)(self.get_pattern)  # type: ignore
+
+    @abstractmethod
+    def get_pattern(self, timesteps: int) -> Pattern:
+        """Builds pattern with specific interleaving between codebooks.
+
+        Args:
+            timesteps (int): Total numer of timesteps.
+        """
+        raise NotImplementedError()
+
+
+class DelayedPatternProvider(CodebooksPatternProvider):
+    """Provider for delayed pattern across delayed codebooks.
+    Codebooks are delayed in the sequence and sequence steps will contain codebooks
+    from different timesteps.
+
+    Example:
+        Taking timesteps=4 and n_q=3, delays=None, the multi-codebook sequence:
+        [[1, 2, 3, 4],
+        [1, 2, 3, 4],
+        [1, 2, 3, 4]]
+        The resulting sequence obtained from the returned pattern is:
+        [[S, 1, 2, 3, 4],
+        [S, S, 1, 2, 3],
+        [S, S, S, 1, 2]]
+        (with S being a special token)
+
+    Args:
+        n_q (int): Number of codebooks.
+        delays (Optional[List[int]]): Delay for each of the codebooks.
+            If delays not defined, each codebook is delayed by 1 compared to the previous one.
+        flatten_first (int): Flatten the first N timesteps.
+        empty_initial (int): Prepend with N empty list of coordinates.
+    """
+    def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None,
+                 flatten_first: int = 0, empty_initial: int = 0):
+        super().__init__(n_q)
+        if delays is None:
+            delays = list(range(n_q))
+        self.delays = delays
+        self.flatten_first = flatten_first
+        self.empty_initial = empty_initial
+        assert len(self.delays) == self.n_q
+        assert sorted(self.delays) == self.delays
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        max_delay = max(self.delays)
+        if self.empty_initial:
+            out += [[] for _ in range(self.empty_initial)]
+        if self.flatten_first:
+            for t in range(min(timesteps, self.flatten_first)):
+                for q in range(self.n_q):
+                    out.append([LayoutCoord(t, q)])
+        for t in range(self.flatten_first, timesteps + max_delay):
+            v = []
+            for q, delay in enumerate(self.delays):
+                t_for_q = t - delay
+                if t_for_q >= self.flatten_first:
+                    v.append(LayoutCoord(t_for_q, q))
+            out.append(v)
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class ParallelPatternProvider(DelayedPatternProvider):
+    """Provider for parallel pattern across codebooks.
+    This pattern provider is a special case of the delayed pattern with actually no delay,
+    hence delays=repeat(0, n_q).
+
+    Args:
+        n_q (int): Number of codebooks.
+    """
+    def __init__(self, n_q: int):
+        super().__init__(n_q, [0] * n_q)
+
+
+class UnrolledPatternProvider(CodebooksPatternProvider):
+    """Provider for unrolling codebooks pattern.
+    This pattern provider enables to represent the codebook flattened completely or only to some extend
+    while also specifying a given delay between the flattened codebooks representation, allowing to
+    unroll the codebooks in the sequence.
+
+    Example:
+        1. Flattening of the codebooks.
+        By default, the pattern provider will fully flatten the codebooks such as flattening=range(n_q),
+        taking n_q = 3 and timesteps = 4:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, S, 1, S, S, 2, S, S, 3, S, S, 4],
+         [S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [1, S, S, 2, S, S, 3, S, S, 4, S, S]]
+        2. Partial flattening of the codebooks. The ``flattening`` parameter allows to specify the inner step
+        for each of the codebook, allowing to define which codebook to flatten (or keep in parallel), for example
+        taking n_q = 3, timesteps = 4 and flattening = [0, 1, 1]:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [1, S, S, 2, S, S, 3, S, S, 4, S, S]]
+        3. Flattening with delay. The ``delay`` parameter allows to further unroll the sequence of codebooks
+        allowing to specify the delay per codebook. Note that the delay between codebooks flattened to the
+        same inner timestep should be coherent. For example, taking n_q = 3, timesteps = 4, flattening = [0, 1, 1]
+        and delays = [0, 3, 3]:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, S, S, 1, S, 2, S, 3, S, 4],
+         [S, S, S, 1, S, 2, S, 3, S, 4],
+         [1, 2, 3, S, 4, S, 5, S, 6, S]]
+
+    Args:
+        n_q (int): Number of codebooks.
+        flattening (Optional[List[int]]): Flattening schema over the codebooks. If not defined,
+            the codebooks will be flattened to 1 codebook per step, meaning that the sequence will
+            have n_q extra steps for each timestep.
+        delays (Optional[List[int]]): Delay for each of the codebooks. If not defined,
+            no delay is added and therefore will default to [0] * ``n_q``.
+            Note that two codebooks that will be flattened to the same inner step
+            should have the same delay, otherwise the pattern is considered as invalid.
+    """
+    FlattenedCodebook = namedtuple('FlattenedCodebook', ['codebooks', 'delay'])
+
+    def __init__(self, n_q: int, flattening: tp.Optional[tp.List[int]] = None,
+                 delays: tp.Optional[tp.List[int]] = None):
+        super().__init__(n_q)
+        if flattening is None:
+            flattening = list(range(n_q))
+        if delays is None:
+            delays = [0] * n_q
+        assert len(flattening) == n_q
+        assert len(delays) == n_q
+        assert sorted(flattening) == flattening
+        assert sorted(delays) == delays
+        self._flattened_codebooks = self._build_flattened_codebooks(delays, flattening)
+        self.max_delay = max(delays)
+
+    def _build_flattened_codebooks(self, delays: tp.List[int], flattening: tp.List[int]):
+        """Build a flattened codebooks representation as a dictionary of inner step
+        and the actual codebook indices corresponding to the flattened codebook. For convenience, we
+        also store the delay associated to the flattened codebook to avoid maintaining an extra mapping.
+        """
+        flattened_codebooks: dict = {}
+        for q, (inner_step, delay) in enumerate(zip(flattening, delays)):
+            if inner_step not in flattened_codebooks:
+                flat_codebook = UnrolledPatternProvider.FlattenedCodebook(codebooks=[q], delay=delay)
+            else:
+                flat_codebook = flattened_codebooks[inner_step]
+                assert flat_codebook.delay == delay, (
+                    "Delay and flattening between codebooks is inconsistent: ",
+                    "two codebooks flattened to the same position should have the same delay."
+                )
+                flat_codebook.codebooks.append(q)
+            flattened_codebooks[inner_step] = flat_codebook
+        return flattened_codebooks
+
+    @property
+    def _num_inner_steps(self):
+        """Number of inner steps to unroll between timesteps in order to flatten the codebooks.
+        """
+        return max([inner_step for inner_step in self._flattened_codebooks.keys()]) + 1
+
+    def num_virtual_steps(self, timesteps: int) -> int:
+        return timesteps * self._num_inner_steps + 1
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        """Builds pattern for delay across codebooks.
+
+        Args:
+            timesteps (int): Total numer of timesteps.
+        """
+        # the PatternLayout is built as a tuple of sequence position and list of coordinates
+        # so that it can be reordered properly given the required delay between codebooks of given timesteps
+        indexed_out: list = [(-1, [])]
+        max_timesteps = timesteps + self.max_delay
+        for t in range(max_timesteps):
+            # for each timestep, we unroll the flattened codebooks,
+            # emitting the sequence step with the corresponding delay
+            for step in range(self._num_inner_steps):
+                if step in self._flattened_codebooks:
+                    # we have codebooks at this virtual step to emit
+                    step_codebooks = self._flattened_codebooks[step]
+                    t_for_q = t + step_codebooks.delay
+                    coords = [LayoutCoord(t, q) for q in step_codebooks.codebooks]
+                    if t_for_q < max_timesteps and t < max_timesteps:
+                        indexed_out.append((t_for_q, coords))
+                else:
+                    # there is no codebook in this virtual step so we emit an empty list
+                    indexed_out.append((t, []))
+        out = [coords for _, coords in sorted(indexed_out)]
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class VALLEPattern(CodebooksPatternProvider):
+    """Almost VALL-E style pattern. We futher allow some delays for the
+    codebooks other than the first one.
+
+    Args:
+        n_q (int): Number of codebooks.
+        delays (Optional[List[int]]): Delay for each of the codebooks.
+            If delays not defined, each codebook is delayed by 1 compared to the previous one.
+    """
+    def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None):
+        super().__init__(n_q)
+        if delays is None:
+            delays = [0] * (n_q - 1)
+        self.delays = delays
+        assert len(self.delays) == self.n_q - 1
+        assert sorted(self.delays) == self.delays
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        for t in range(timesteps):
+            out.append([LayoutCoord(t, 0)])
+        max_delay = max(self.delays)
+        for t in range(timesteps + max_delay):
+            v = []
+            for q, delay in enumerate(self.delays):
+                t_for_q = t - delay
+                if t_for_q >= 0:
+                    v.append(LayoutCoord(t_for_q, q + 1))
+            out.append(v)
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class MusicLMPattern(CodebooksPatternProvider):
+    """Almost MusicLM style pattern. This is equivalent to full flattening
+    but in a different order.
+
+    Args:
+        n_q (int): Number of codebooks.
+        group_by (int): Number of codebooks to group together.
+    """
+    def __init__(self, n_q: int, group_by: int = 2):
+        super().__init__(n_q)
+        self.group_by = group_by
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        for offset in range(0, self.n_q, self.group_by):
+            for t in range(timesteps):
+                for q in range(offset, offset + self.group_by):
+                    out.append([LayoutCoord(t, q)])
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)

src/model/modules/embedding.py

@@ -0,0 +1,98 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/embedding.py
+# Copyright    2023                             (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        vocab_size: int,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.dim_model = dim_model
+
+        self.dropout = torch.nn.Dropout(p=dropout)
+        self.word_embeddings = nn.Embedding(self.vocab_size, self.dim_model)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.word_embeddings.weight
+
+    def embedding(self, index: int) -> torch.Tensor:
+        return self.word_embeddings.weight[index : index + 1]
+
+    def forward(self, x: torch.Tensor):
+        X = self.word_embeddings(x)
+        X = self.dropout(X)
+
+        return X
+
+
+class SinePositionalEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        dropout: float = 0.0,
+        scale: bool = False,
+        alpha: bool = False,
+    ):
+        super().__init__()
+        self.dim_model = dim_model
+        self.x_scale = math.sqrt(dim_model) if scale else 1.0
+        self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
+        self.dropout = torch.nn.Dropout(p=dropout)
+
+        self.reverse = False
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, 4000))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            if self.pe.size(1) >= x.size(1):
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        pe = torch.zeros(x.size(1), self.dim_model)
+        if self.reverse:
+            position = torch.arange(
+                x.size(1) - 1, -1, -1.0, dtype=torch.float32
+            ).unsqueeze(1)
+        else:
+            position = torch.arange(
+                0, x.size(1), dtype=torch.float32
+            ).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.dim_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.dim_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.pe = pe.to(device=x.device, dtype=x.dtype).detach()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.extend_pe(x)
+        output = x.unsqueeze(-1) if x.ndim == 2 else x
+        output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]
+        return self.dropout(output)

src/model/modules/gemma.py

@@ -0,0 +1,423 @@
+import torch
+from torch import nn
+from typing import Optional, Tuple
+import math
+from src.model.modules.kv_cache import KVCache
+
+
+class GemmaConfig:
+
+    def __init__(
+        self,
+        vocab_size,
+        hidden_size,
+        intermediate_size,
+        num_hidden_layers,
+        num_attention_heads,
+        num_key_value_heads,
+        head_dim=256,
+        max_position_embeddings=8192,
+        rms_norm_eps=1e-6,
+        rope_theta=10000.0,
+        attention_bias=False,
+        attention_dropout=0.0,
+        pad_token_id=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.pad_token_id = pad_token_id
+
+
+class GemmaRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float())
+        # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16)
+        # See https://github.com/huggingface/transformers/pull/29402
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+
+class GemmaRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim  # it is set to the head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+
+        # Calculate the theta according to the formula theta_i = base^(2i/dim) where i = 0, 1, 2, ..., dim // 2
+        inv_freq = 1.0 / (
+            self.base
+            ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim)
+        )
+        self.register_buffer("inv_freq", tensor=inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x, position_ids, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        self.inv_freq.to(x.device)
+        # Copy the inv_freq tensor for batch in the sequence
+        # inv_freq_expanded: [Batch_Size, Head_Dim // 2, 1]
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        )
+        # position_ids_expanded: [Batch_Size, 1, Seq_Len]
+        position_ids_expanded = position_ids[:, None, :].float()
+        device_type = x.device.type
+        device_type = (
+            device_type
+            if isinstance(device_type, str) and device_type != "mps"
+            else "cpu"
+        )
+        with torch.autocast(device_type=device_type, enabled=False):
+            # Multiply each theta by the position (which is the argument of the sin and cos functions)
+            # freqs: [Batch_Size, Head_Dim // 2, 1] @ [Batch_Size, 1, Seq_Len] --> [Batch_Size, Seq_Len, Head_Dim // 2]
+            freqs = (
+                inv_freq_expanded.float() @ position_ids_expanded.float()
+            ).transpose(1, 2)
+            # emb: [Batch_Size, Seq_Len, Head_Dim]
+            emb = torch.cat((freqs, freqs), dim=-1)
+            # cos, sin: [Batch_Size, Seq_Len, Head_Dim]
+            cos = emb.cos()
+            sin = emb.sin()
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    # Build the [-x2, x1, -x4, x3, ...] tensor for the sin part of the positional encoding.
+    x1 = x[..., : x.shape[-1] // 2]  # Takes the first half of the last dimension
+    x2 = x[..., x.shape[-1] // 2 :]  # Takes the second half of the last dimension
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)  # Add the head dimension
+    sin = sin.unsqueeze(unsqueeze_dim)  # Add the head dimension
+    # Apply the formula (34) of the Rotary Positional Encoding paper.
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class GemmaMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def forward(self, x):
+        # Equivalent to:
+        # y = self.gate_proj(x) # [Batch_Size, Seq_Len, Hidden_Size] -> [Batch_Size, Seq_Len, Intermediate_Size]
+        # y = torch.gelu(y, approximate="tanh") # [Batch_Size, Seq_Len, Intermediate_Size]
+        # j = self.up_proj(x) # [Batch_Size, Seq_Len, Hidden_Size] -> [Batch_Size, Seq_Len, Intermediate_Size]
+        # z = y * j # [Batch_Size, Seq_Len, Intermediate_Size]
+        # z = self.down_proj(z) # [Batch_Size, Seq_Len, Intermediate_Size] -> [Batch_Size, Seq_Len, Hidden_Size]
+        return self.down_proj(
+            nn.functional.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x)
+        )
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class GemmaAttention(nn.Module):
+
+    def __init__(self, config: GemmaConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+
+        assert self.hidden_size % self.num_heads == 0
+
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
+        )
+        self.rotary_emb = GemmaRotaryEmbedding(
+            self.head_dim,
+            max_position_embeddings=self.max_position_embeddings,
+            base=self.rope_theta,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        kv_cache: Optional[KVCache] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()  # [Batch_Size, Seq_Len, Hidden_Size]
+        # [Batch_Size, Seq_Len, Num_Heads_Q * Head_Dim]
+        query_states = self.q_proj(hidden_states)
+        # [Batch_Size, Seq_Len, Num_Heads_KV * Head_Dim]
+        key_states = self.k_proj(hidden_states)
+        # [Batch_Size, Seq_Len, Num_Heads_KV * Head_Dim]
+        value_states = self.v_proj(hidden_states)
+        # [Batch_Size, Num_Heads_Q, Seq_Len, Head_Dim]
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        # [Batch_Size, Num_Heads_KV, Seq_Len, Head_Dim]
+        key_states = key_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        # [Batch_Size, Num_Heads_KV, Seq_Len, Head_Dim]
+        value_states = value_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+
+        # [Batch_Size, Seq_Len, Head_Dim], [Batch_Size, Seq_Len, Head_Dim]
+        cos, sin = self.rotary_emb(value_states, position_ids, seq_len=None)
+        # [Batch_Size, Num_Heads_Q, Seq_Len, Head_Dim], [Batch_Size, Num_Heads_KV, Seq_Len, Head_Dim]
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if kv_cache is not None:
+            key_states, value_states = kv_cache.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        # Repeat the key and values to match the number of heads of the query
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        # Perform the calculation as usual, Q * K^T / sqrt(head_dim). Shape: [Batch_Size, Num_Heads_Q, Seq_Len_Q, Seq_Len_KV]
+        attn_weights = torch.matmul(
+            query_states, key_states.transpose(2, 3)
+        ) / math.sqrt(self.head_dim)
+
+        assert attention_mask is not None
+        attn_weights = attn_weights + attention_mask
+
+        # Apply the softmax
+        # [Batch_Size, Num_Heads_Q, Seq_Len_Q, Seq_Len_KV]
+        attn_weights = nn.functional.softmax(
+            attn_weights, dim=-1, dtype=torch.float32
+        ).to(query_states.dtype)
+        # Apply the dropout
+        attn_weights = nn.functional.dropout(
+            attn_weights, p=self.attention_dropout, training=self.training
+        )
+        # Multiply by the values. [Batch_Size, Num_Heads_Q, Seq_Len_Q, Seq_Len_KV] x [Batch_Size, Num_Heads_KV, Seq_Len_KV, Head_Dim] -> [Batch_Size, Num_Heads_Q, Seq_Len_Q, Head_Dim]
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+        # Make sure the sequence length is the second dimension. # [Batch_Size, Num_Heads_Q, Seq_Len_Q, Head_Dim] -> [Batch_Size, Seq_Len_Q, Num_Heads_Q, Head_Dim]
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        # Concatenate all the heads together. [Batch_Size, Seq_Len_Q, Num_Heads_Q, Head_Dim] -> [Batch_Size, Seq_Len_Q, Num_Heads_Q * Head_Dim]
+        attn_output = attn_output.view(bsz, q_len, -1)
+        # Multiply by W_o. [Batch_Size, Seq_Len_Q, Hidden_Size]
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class GemmaDecoderLayer(nn.Module):
+
+    def __init__(self, config: GemmaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = GemmaAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = GemmaMLP(config)
+        self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = GemmaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        kv_cache: Optional[KVCache] = None,
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        residual = hidden_states
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        (
+            hidden_states,
+            _,
+        ) = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            kv_cache=kv_cache,
+        )
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        hidden_states = residual + hidden_states
+
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        residual = hidden_states
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        hidden_states = self.mlp(hidden_states)
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class GemmaModel(nn.Module):
+
+    def __init__(self, config: GemmaConfig):
+        super().__init__()
+        self.config = config
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                GemmaDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    # Ignore copy
+    def forward(
+        self,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        kv_cache: Optional[KVCache] = None,
+    ) -> torch.FloatTensor:
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        hidden_states = inputs_embeds
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        normalizer = torch.tensor(
+            self.config.hidden_size**0.5, dtype=hidden_states.dtype
+        )
+        hidden_states = hidden_states * normalizer
+
+        for decoder_layer in self.layers:
+            # [Batch_Size, Seq_Len, Hidden_Size]
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                kv_cache=kv_cache,
+            )
+
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        hidden_states = self.norm(hidden_states)
+
+        # [Batch_Size, Seq_Len, Hidden_Size]
+        return hidden_states
+
+
+class GemmaForCausalLM(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.model = GemmaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def tie_weights(self):
+        self.lm_head.weight = self.model.embed_tokens.weight
+
+    def forward(
+        self,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        kv_cache: Optional[KVCache] = None,
+    ) -> Tuple:
+
+        # input_embeds: [Batch_Size, Seq_Len, Hidden_Size]
+        # outputs: [Batch_Size, Seq_Len, Hidden_Size]
+        outputs = self.model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            kv_cache=kv_cache,
+        )
+
+        hidden_states = outputs
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        return_data = {
+            "logits": logits,
+        }
+
+        if kv_cache is not None:
+            # Return the updated cache
+            return_data["kv_cache"] = kv_cache
+
+        return return_data

src/model/modules/imagecraft.py

@@ -0,0 +1,490 @@
+from argparse import Namespace
+import glob
+import logging
+from pathlib import Path
+import os
+import time
+from typing import Optional, Tuple
+from PIL import Image
+from safetensors import safe_open
+import torch
+from torch import nn
+import torchaudio
+from src.model.modules import voicecraft
+from src.model.modules.gemma import GemmaForCausalLM, KVCache
+from src.model.modules.imagecraftconfig import ImageCraftConfig
+from src.model.modules.imagecraftprocessor import (
+    ImageCraftProcessor,
+)
+from src.model.modules.siglip import SiglipVisionModel
+
+from transformers import AutoTokenizer
+
+from src.model.modules.tokenizer import (
+    AudioTokenizer,
+    TextTokenizer,
+    tokenize_audio,
+    tokenize_text,
+)
+
+
+from src.utils import tools
+from src.utils.image_utils import is_valid_image
+from src.utils.model_utils import get_config, get_model_inputs
+from src.utils.util import (
+    replace_numbers_with_words,
+    sample_top_p,
+    save_to_buffer,
+    save_to_file,
+    split_line_to_sentences,
+)
+
+from huggingface_hub import HfApi
+
+logger = logging.getLogger(__name__)
+
+
+class ImageCraftMultiModalProjector(nn.Module):
+    def __init__(self, config: ImageCraftConfig):
+        super().__init__()
+        self.linear = nn.Linear(
+            config.vision_config.hidden_size,
+            config.vision_config.projection_dim,
+            bias=True,
+        )
+
+    def forward(self, image_features):
+        hidden_states = self.linear(image_features)
+        return hidden_states
+
+
+class ImageCraft(nn.Module):
+    config_class = ImageCraftConfig
+
+    def __init__(self, config: ImageCraftConfig):
+        super(ImageCraft, self).__init__()
+        self.config = config
+        self.vision_tower = SiglipVisionModel(config.vision_config)
+        self.multi_modal_projector = ImageCraftMultiModalProjector(config)
+        self.vocab_size = config.text_config.vocab_size
+
+        self.language_model = GemmaForCausalLM(config.text_config)
+
+        self.pad_token_id = (
+            self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        )
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "google/paligemma-3b-pt-224", padding_side="right"
+        )
+        assert tokenizer.padding_side == "right"
+
+        num_image_tokens = config.vision_config.num_image_tokens
+        image_size = config.vision_config.image_size
+        self.processor = ImageCraftProcessor(tokenizer, num_image_tokens, image_size)
+
+        self.text_tokenizer = None
+
+        self.voicecraft_model = None
+        self.audio_tokenizer = None
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.tie_weights with Llava->PaliGemma
+    def tie_weights(self):
+        return self.language_model.tie_weights()
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: torch.FloatTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        kv_cache: Optional[KVCache] = None,
+    ) -> Tuple:
+        # Make sure the input is right-padded
+        assert torch.all(attention_mask == 1), "The input cannot be padded"
+
+        # 1. Extra the input embeddings
+        # shape: (Batch_Size, Seq_Len, Hidden_Size)
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        # 2. Merge text and images
+        # [Batch_Size, Channels, Height, Width] -> [Batch_Size, Num_Patches, Embed_Dim]
+        selected_image_feature = self.vision_tower(pixel_values.to(inputs_embeds.dtype))
+        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Hidden_Size]
+        image_features = self.multi_modal_projector(selected_image_feature)
+
+        # Merge the embeddings of the text tokens and the image tokens
+        inputs_embeds, attention_mask, position_ids = (
+            self._merge_input_ids_with_image_features(
+                image_features, inputs_embeds, input_ids, attention_mask, kv_cache
+            )
+        )
+
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            kv_cache=kv_cache,
+        )
+
+        return outputs
+
+    def _merge_input_ids_with_image_features(
+        self,
+        image_features: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        kv_cache: Optional[KVCache] = None,
+    ):
+        _, _, embed_dim = image_features.shape
+        batch_size, sequence_length = input_ids.shape
+        dtype, device = inputs_embeds.dtype, inputs_embeds.device
+        # Shape: [Batch_Size, Seq_Len, Hidden_Size]
+        scaled_image_features = image_features / (self.config.hidden_size**0.5)
+
+        # Combine the embeddings of the image tokens, the text tokens and mask out all the padding tokens.
+        final_embedding = torch.zeros(
+            batch_size,
+            sequence_length,
+            embed_dim,
+            dtype=inputs_embeds.dtype,
+            device=inputs_embeds.device,
+        )
+        # Shape: [Batch_Size, Seq_Len]. True for text tokens
+        text_mask = (input_ids != self.config.image_token_index) & (
+            input_ids != self.pad_token_id
+        )
+        # Shape: [Batch_Size, Seq_Len]. True for image tokens
+        image_mask = input_ids == self.config.image_token_index
+        # Shape: [Batch_Size, Seq_Len]. True for padding tokens
+        pad_mask = input_ids == self.pad_token_id
+
+        # We need to expand the masks to the embedding dimension otherwise we can't use them in torch.where
+        text_mask_expanded = text_mask.unsqueeze(-1).expand(-1, -1, embed_dim)
+        pad_mask_expanded = pad_mask.unsqueeze(-1).expand(-1, -1, embed_dim)
+        image_mask_expanded = image_mask.unsqueeze(-1).expand(-1, -1, embed_dim)
+
+        # Add the text embeddings
+        final_embedding = torch.where(
+            text_mask_expanded, inputs_embeds, final_embedding
+        )
+        # Insert image embeddings. We can't use torch.where because the sequence length of scaled_image_features is not equal to the sequence length of the final embedding
+        final_embedding = final_embedding.masked_scatter(
+            image_mask_expanded, scaled_image_features
+        )
+        # Zero out padding tokens
+        final_embedding = torch.where(
+            pad_mask_expanded, torch.zeros_like(final_embedding), final_embedding
+        )
+
+        #### CREATE THE ATTENTION MASK ####
+
+        dtype, device = inputs_embeds.dtype, inputs_embeds.device
+        min_dtype = torch.finfo(dtype).min
+        q_len = inputs_embeds.shape[1]
+
+        if kv_cache is None or kv_cache.num_items() == 0:
+            # Do not mask any token, because we're in the prefill phase
+            # This only works when we have no padding
+            causal_mask = torch.full(
+                (batch_size, q_len, q_len), fill_value=0, dtype=dtype, device=device
+            )
+        else:
+            # Since we are generating tokens, the query must be one single token
+            assert q_len == 1
+            kv_len = kv_cache.num_items() + q_len
+            # Also in this case we don't need to mask anything, since each query should be able to attend all previous tokens.
+            # This only works when we have no padding
+            causal_mask = torch.full(
+                (batch_size, q_len, kv_len), fill_value=0, dtype=dtype, device=device
+            )
+
+        # Add the head dimension
+        # [Batch_Size, Q_Len, KV_Len] -> [Batch_Size, Num_Heads_Q, Q_Len, KV_Len]
+        causal_mask = causal_mask.unsqueeze(1)
+
+        if kv_cache is not None and kv_cache.num_items() > 0:
+            # The position of the query is just the last position
+            position_ids = attention_mask.cumsum(-1)[:, -1]
+            if position_ids.dim() == 1:
+                position_ids = position_ids.unsqueeze(0)
+        else:
+            # Create a position_ids based on the size of the attention_mask
+            # For masked tokens, use the number 1 as position.
+            position_ids = (
+                (attention_mask.cumsum(-1))
+                .masked_fill_((attention_mask == 0), 1)
+                .to(device)
+            )
+
+        return final_embedding, causal_mask, position_ids
+
+    def _generate_caption(self, image, max_tokens=100, do_sample=False):
+        prompt = "caption en"
+        image = (
+            image.convert("RGB")
+            if is_valid_image(image)
+            else Image.open(image).convert("RGB")
+        )
+
+        inputs = get_model_inputs(
+            processor=self.processor, prompt=prompt, image=image, device=self.device
+        )
+
+        image.close()
+
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+        pixel_values = inputs["pixel_values"]
+
+        kv_cache = KVCache()
+
+        stop_token = self.processor.tokenizer.eos_token_id
+        generated_tokens = []
+
+        for _ in range(max_tokens):
+            outputs = self(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                attention_mask=attention_mask,
+                kv_cache=kv_cache,
+            )
+            kv_cache = outputs["kv_cache"]
+            next_token_logits = outputs["logits"][:, -1, :]
+            if do_sample:
+                next_token_logits = torch.softmax(
+                    next_token_logits / self.config.temperature, dim=-1
+                )
+                next_token = sample_top_p(next_token_logits, self.config.top_p)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+            assert next_token.size() == (1, 1)
+            next_token = next_token.squeeze(0)
+            generated_tokens.append(next_token)
+            if next_token.item() == stop_token:
+                break
+            input_ids = next_token.unsqueeze(-1)
+            attention_mask = torch.cat(
+                [attention_mask, torch.ones((1, 1), device=input_ids.device)], dim=-1
+            )
+
+        generated_tokens = torch.cat(generated_tokens, dim=-1)
+        decoded_text = self.processor.tokenizer.decode(
+            generated_tokens, skip_special_tokens=True
+        )
+        decoded_text = (
+            parts[1] if len(parts := decoded_text.split("\n", 1)) > 1 else decoded_text
+        )
+
+        return decoded_text.rstrip(" .").strip().capitalize() + "."
+
+    def _generate_speech(self, text: str, output_type="file"):
+
+        sentences = split_line_to_sentences(text)
+
+        voice_audio = (
+            f"media/voicecraft/voices/{self.config.voicecraft_config.voice_audio_path}"
+        )
+        voice_transcript = self.config.voicecraft_config.voice_audio_transcript
+        cut_off_sec = self.config.voicecraft_config.cut_off_sec
+
+        decode_config = {
+            "top_k": self.config.voicecraft_config.top_k,
+            "top_p": self.config.voicecraft_config.top_p,
+            "temperature": self.config.voicecraft_config.temperature,
+            "stop_repetition": self.config.voicecraft_config.stop_repetition,
+            "kvcache": self.config.voicecraft_config.kvcache,
+            "codec_audio_sr": self.config.voicecraft_config.codec_audio_sr,
+            "codec_sr": self.config.voicecraft_config.codec_sr,
+            "silence_tokens": self.config.voicecraft_config.silence_tokens,
+            "sample_batch_size": self.config.voicecraft_config.sample_batch_size,
+        }
+
+        info = torchaudio.info(voice_audio)
+        audio_dur = info.num_frames / info.sample_rate
+        prompt_end_frame = int(min(audio_dur, cut_off_sec) * info.sample_rate)
+
+        audio_tensors = []
+        transcript = voice_transcript
+
+        for sentence in sentences:
+
+            transcript += sentence + "\n"
+            transcript = replace_numbers_with_words(transcript).replace("  ", " ")
+
+            # phonemize
+            phn2num = self.voicecraft_model.args.phn2num
+            text_tokens = [
+                phn2num[phn]
+                for phn in tokenize_text(self.text_tokenizer, text=transcript.strip())
+                if phn in phn2num
+            ]
+            text_tokens = torch.LongTensor(text_tokens).unsqueeze(0)
+            text_tokens_lens = torch.LongTensor([text_tokens.shape[-1]])
+
+            # encode audio
+            encoded_frames = tokenize_audio(
+                self.audio_tokenizer,
+                voice_audio,
+                offset=0,
+                num_frames=prompt_end_frame,
+            )
+            original_audio = encoded_frames[0][0].transpose(2, 1)  # [1,T,K]
+            model_args = vars(self.voicecraft_model.args)
+            model_args = Namespace(**model_args)
+
+            assert (
+                original_audio.ndim == 3
+                and original_audio.shape[0] == 1
+                and original_audio.shape[2] == model_args.n_codebooks
+            ), original_audio.shape
+
+            # forward
+            stime = time.time()
+            if decode_config["sample_batch_size"] <= 1:
+                _, gen_frames = self.voicecraft_model.inference_tts(
+                    text_tokens.to(self.device),
+                    text_tokens_lens.to(self.device),
+                    original_audio[..., : model_args.n_codebooks].to(
+                        self.device
+                    ),  # [1,T,8]
+                    top_k=decode_config["top_k"],
+                    top_p=decode_config["top_p"],
+                    temperature=decode_config["temperature"],
+                    stop_repetition=decode_config["stop_repetition"],
+                    kvcache=decode_config["kvcache"],
+                    silence_tokens=(
+                        eval(decode_config["silence_tokens"])
+                        if type(decode_config["silence_tokens"]) == str
+                        else decode_config["silence_tokens"]
+                    ),
+                )  # output is [1,K,T]
+            else:
+                _, gen_frames = self.voicecraft_model.inference_tts_batch(
+                    text_tokens.to(self.device),
+                    text_tokens_lens.to(self.device),
+                    original_audio[..., : model_args.n_codebooks].to(
+                        self.device
+                    ),  # [1,T,8]
+                    top_k=decode_config["top_k"],
+                    top_p=decode_config["top_p"],
+                    temperature=decode_config["temperature"],
+                    stop_repetition=decode_config["stop_repetition"],
+                    kvcache=decode_config["kvcache"],
+                    batch_size=decode_config["sample_batch_size"],
+                    silence_tokens=(
+                        eval(decode_config["silence_tokens"])
+                        if type(decode_config["silence_tokens"]) == str
+                        else decode_config["silence_tokens"]
+                    ),
+                )  # output is [1,K,T]
+            gen_sample = self.audio_tokenizer.decode([(gen_frames, None)])
+            gen_audio = gen_sample[0].cpu()
+            audio_tensors.append(gen_audio)
+
+        output = None
+
+        if output_type == "file":
+            output = save_to_file(audio_tensors, decode_config["codec_audio_sr"])
+        else:
+            output = save_to_buffer(audio_tensors, decode_config["codec_audio_sr"])
+
+        # Empty cuda cache between runs
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return output
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        image,
+        max_tokens=30,
+        do_sample=False,
+        output_type="file",
+        return_output="speech",
+    ):
+        if return_output == "speech" or return_output is None:
+            transcript = self._generate_caption(image, max_tokens, do_sample)
+            speech = self._generate_speech(transcript, output_type)
+            return transcript, speech
+        else:
+            transcript = self._generate_caption(image, max_tokens, do_sample)
+            return transcript
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path=None,
+    ):
+        api = HfApi()
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        env_config = tools.load_config()
+        pretrained_dir = env_config["pretrained_dir"]
+        imagecraft_cache_dir = f"{pretrained_dir}/imagecraft"
+        voicecraft_cache_dir = f"{pretrained_dir}/voicecraft"
+
+        state_dict = {}
+
+        if Path(model_path).is_file():
+            checkpoint = torch.load(model_path, weights_only=False)
+            state_dict = checkpoint["state_dict"]
+
+        else:
+
+            model_path = api.snapshot_download(
+                repo_id=model_path,
+                repo_type="model",
+                cache_dir=imagecraft_cache_dir,
+                local_files_only=False,
+            )
+
+            safetensors_files = glob.glob(os.path.join(model_path, "*.safetensors"))
+
+            for safetensors_file in safetensors_files:
+                with safe_open(safetensors_file, framework="pt", device="cpu") as f:
+                    for key in f.keys():
+                        state_dict[key] = f.get_tensor(key)
+
+        imagecraft_config = get_config()
+
+        model = cls(imagecraft_config).to(device)
+
+        # Load the state dict of the model
+        model.load_state_dict(state_dict, strict=False)
+
+        # Tie weights
+        model.tie_weights()
+
+        model = model.eval()
+
+        # Load voicecraft module
+
+        model.voicecraft_model = voicecraft.VoiceCraft.from_pretrained(
+            f"pyp1/VoiceCraft_{model.config.voicecraft_config.model_name.replace('.pth', '')}",
+            cache_dir=voicecraft_cache_dir,
+        )
+
+        encodec_fn = f"{voicecraft_cache_dir}/{model.config.voicecraft_config.encodec}"
+
+        if not os.path.exists(encodec_fn):
+            os.system(
+                f"wget https://huggingface.co/pyp1/VoiceCraft/resolve/main/{model.config.voicecraft_config.encodec}"
+            )
+            os.system(f"mv {model.config.voicecraft_config.encodec} {encodec_fn}")
+
+        model.audio_tokenizer = AudioTokenizer(
+            signature=encodec_fn,
+            device=device,
+        )
+
+        model.text_tokenizer = TextTokenizer(backend="espeak")
+
+        model.voicecraft_model.to(device)
+
+        return model

src/model/modules/imagecraftconfig.py

@@ -0,0 +1,47 @@
+# from src.model.modules.gemma import GemmaConfig
+# from src.model.modules.siglip import SiglipVisionConfig
+from src.model.modules.voicecraftconfig import VoiceCraftConfig
+
+from transformers import SiglipVisionConfig, GemmaConfig, PretrainedConfig
+
+
+class ImageCraftConfig(PretrainedConfig):
+
+    model_type = "imagecraft"
+
+    def __init__(
+        self,
+        vision_config=None,
+        text_config=None,
+        voicecraft_config=None,
+        ignore_index=-100,
+        image_token_index=256000,
+        vocab_size=257152,
+        projection_dim=2048,
+        hidden_size=2048,
+        pad_token_id=None,
+        **kwargs
+    ):
+        super().__init__()
+        self.ignore_index = ignore_index
+        self.image_token_index = image_token_index
+        self.vocab_size = vocab_size
+        self.projection_dim = projection_dim
+        self.hidden_size = hidden_size
+        self.is_encoder_decoder = False
+
+        self.pad_token_id = pad_token_id if pad_token_id is not None else -1
+
+        self.vision_config = SiglipVisionConfig(**vision_config)
+
+        self.text_config = GemmaConfig(**text_config, pad_token_id=pad_token_id)
+        self.vocab_size = self.text_config.vocab_size
+
+        self.text_config.num_image_tokens = (
+            self.vision_config.image_size // self.vision_config.patch_size
+        ) ** 2
+        self.vision_config.projection_dim = projection_dim
+
+        self.voicecraft_config = VoiceCraftConfig(**voicecraft_config)
+
+        super().__init__(**kwargs)

src/model/modules/imagecraftprocessor.py

@@ -0,0 +1,96 @@
+from typing import List, Optional
+from PIL import Image
+import numpy as np
+import torch
+
+
+from src.utils.util import (
+    IMAGENET_STANDARD_MEAN,
+    IMAGENET_STANDARD_STD,
+    add_image_tokens_to_prompt,
+    process_images,
+)
+
+from transformers import SiglipImageProcessor
+
+
+class ImageCraftProcessor:
+
+    IMAGE_TOKEN = "<image>"
+
+    def __init__(self, tokenizer, num_image_tokens: int, image_size: int):
+        super().__init__()
+
+        self.image_seq_length = num_image_tokens
+        self.image_size = image_size
+
+        # Tokenizer described here: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/README.md#tokenizer
+        tokens_to_add = {"additional_special_tokens": [self.IMAGE_TOKEN]}
+        tokenizer.add_special_tokens(tokens_to_add)
+        EXTRA_TOKENS = [
+            f"<loc{i:04d}>" for i in range(1024)
+        ]  # These tokens are used for object detection (bounding boxes)
+        EXTRA_TOKENS += [
+            f"<seg{i:03d}>" for i in range(128)
+        ]  # These tokens are used for object segmentation
+        tokenizer.add_tokens(EXTRA_TOKENS)
+        self.image_token_id = tokenizer.convert_tokens_to_ids(self.IMAGE_TOKEN)
+        # We will add the BOS and EOS tokens ourselves
+        tokenizer.add_bos_token = False
+        tokenizer.add_eos_token = False
+
+        self.tokenizer = tokenizer
+        # self.image_processor = SiglipImageProcessor.from_pretrained(
+        #     "google/siglip-so400m-patch14-384"
+        # )
+
+    def __call__(
+        self,
+        text: List[str],
+        images: List[Image.Image],
+        padding: str = "longest",
+        truncation: bool = True,
+    ) -> dict:
+        assert (
+            len(images) == 1 and len(text) == 1
+        ), f"Received {len(images)} images for {len(text)} prompts."
+
+        # pixel_values = self.image_processor(images=images, return_tensors="pt")[
+        #     "pixel_values"
+        # ]
+        pixel_values = process_images(
+            images,
+            size=(self.image_size, self.image_size),
+            resample=Image.Resampling.BICUBIC,
+            rescale_factor=1 / 255.0,
+            image_mean=IMAGENET_STANDARD_MEAN,
+            image_std=IMAGENET_STANDARD_STD,
+        )
+        # Convert the list of numpy arrays to a single numpy array with shape [Batch_Size, Channel, Height, Width]
+        pixel_values = np.stack(pixel_values, axis=0)
+        # Convert the numpy array to a PyTorch tensor
+        pixel_values = torch.tensor(pixel_values)
+
+        input_strings = [
+            add_image_tokens_to_prompt(
+                prefix_prompt=prompt,
+                bos_token=self.tokenizer.bos_token,
+                image_seq_length=self.image_seq_length,
+                image_token=self.IMAGE_TOKEN,
+            )
+            for prompt in text
+        ]
+
+        # max_length += self.image_seq_length
+
+        inputs = self.tokenizer(
+            input_strings,
+            return_tensors="pt",
+            padding=padding,
+            max_length=512,
+            truncation=truncation,
+        )
+
+        return_data = {"pixel_values": pixel_values, **inputs}
+
+        return return_data

src/model/modules/kv_cache.py

@@ -0,0 +1,38 @@
+from typing import List, Tuple
+import torch
+
+class KVCache:
+
+    def __init__(self) -> None:
+        self.key_cache: List[torch.Tensor] = []
+        self.value_cache: List[torch.Tensor] = []
+
+    def num_items(self) -> int:
+        if len(self.key_cache) == 0:
+            return 0
+        else:
+            # The shape of the key_cache is [Batch_Size, Num_Heads_KV, Seq_Len, Head_Dim]
+            return self.key_cache[0].shape[-2]
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if len(self.key_cache) <= layer_idx:
+            # If we never added anything to the KV-Cache of this layer, let's create it.
+            self.key_cache.append(key_states)
+            self.value_cache.append(value_states)
+        else:
+            # ... otherwise we concatenate the new keys with the existing ones.
+            # each tensor has shape: [Batch_Size, Num_Heads_KV, Seq_Len, Head_Dim]
+            self.key_cache[layer_idx] = torch.cat(
+                [self.key_cache[layer_idx], key_states], dim=-2
+            )
+            self.value_cache[layer_idx] = torch.cat(
+                [self.value_cache[layer_idx], value_states], dim=-2
+            )
+
+        # ... and then we return all the existing keys + the new ones.
+        return self.key_cache[layer_idx], self.value_cache[layer_idx]

src/model/modules/sampling.py

@@ -0,0 +1,65 @@
+# cp from https://github.com/jasonppy/VoiceCraft/blob/master/models/modules/sampling.py
+
+import torch
+import torch.nn.functional as F
+
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(
+            max(top_k, min_tokens_to_keep), logits.size(-1)
+        )  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(
+            F.softmax(sorted_logits, dim=-1), dim=-1
+        )
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+            ..., :-1
+        ].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = filter_value
+    return logits
+    
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    return token

src/model/modules/scaling.py

@@ -0,0 +1,1391 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/scaling.py
+# Copyright    2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import collections
+import logging
+import random
+import math
+from functools import reduce
+from itertools import repeat
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Embedding as ScaledEmbedding
+
+# from valle.utils import Transpose
+
+class Transpose(nn.Identity):
+    """(N, T, D) -> (N, D, T)"""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input.transpose(1, 2)
+    
+class ActivationBalancerFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        scale_factor: Tensor,
+        sign_factor: Optional[Tensor],
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        if sign_factor is None:
+            ctx.save_for_backward(xgt0, scale_factor)
+        else:
+            ctx.save_for_backward(xgt0, scale_factor, sign_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        if len(ctx.saved_tensors) == 3:
+            xgt0, scale_factor, sign_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+                sign_factor = sign_factor.unsqueeze(-1)
+            factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        else:
+            xgt0, scale_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+            factor = scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+def _compute_scale_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_abs: float,
+    max_abs: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    x_abs_mean = torch.mean(x.abs(), dim=sum_dims).to(torch.float32)
+
+    if min_abs == 0.0:
+        below_threshold = 0.0
+    else:
+        # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if
+        # x_abs)_mean , min_abs.
+        below_threshold = (
+            (min_abs - x_abs_mean) * (gain_factor / min_abs)
+        ).clamp(min=0, max=max_factor)
+
+    above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(
+        min=0, max=max_factor
+    )
+
+    return below_threshold - above_threshold
+
+
+def _compute_sign_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_positive: float,
+    max_positive: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims)
+    if min_positive == 0.0:
+        factor1 = 0.0
+    else:
+        # 0 if proportion_positive >= min_positive, else can be
+        # as large as max_factor.
+        factor1 = (
+            (min_positive - proportion_positive) * (gain_factor / min_positive)
+        ).clamp_(min=0, max=max_factor)
+
+    if max_positive == 1.0:
+        factor2 = 0.0
+    else:
+        # 0 if self.proportion_positive <= max_positive, else can be
+        # as large as -max_factor.
+        factor2 = (
+            (proportion_positive - max_positive)
+            * (gain_factor / (1.0 - max_positive))
+        ).clamp_(min=0, max=max_factor)
+    sign_factor = factor1 - factor2
+    # require min_positive != 0 or max_positive != 1:
+    assert not isinstance(sign_factor, float)
+    return sign_factor
+
+
+class ActivationScaleBalancerFunction(torch.autograd.Function):
+    """
+    This object is used in class ActivationBalancer when the user specified
+    min_positive=0, max_positive=1, so there are no constraints on the signs
+    of the activations and only the absolute value has a constraint.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        sign_factor: Tensor,
+        scale_factor: Tensor,
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        ctx.save_for_backward(xgt0, sign_factor, scale_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        xgt0, sign_factor, scale_factor = ctx.saved_tensors
+        for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+            sign_factor = sign_factor.unsqueeze(-1)
+            scale_factor = scale_factor.unsqueeze(-1)
+
+        factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+class RandomClampFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        min: Optional[float],
+        max: Optional[float],
+        prob: float,
+        reflect: float,
+    ) -> Tensor:
+        x_clamped = torch.clamp(x, min=min, max=max)
+        mask = torch.rand_like(x) < prob
+        ans = torch.where(mask, x_clamped, x)
+        if x.requires_grad:
+            ctx.save_for_backward(ans == x)
+            ctx.reflect = reflect
+        if reflect != 0.0:
+            ans = ans * (1.0 + reflect) - (x * reflect)
+        return ans
+
+    @staticmethod
+    def backward(
+        ctx, ans_grad: Tensor
+    ) -> Tuple[Tensor, None, None, None, None]:
+        (is_same,) = ctx.saved_tensors
+        x_grad = ans_grad * is_same.to(ans_grad.dtype)
+        reflect = ctx.reflect
+        if reflect != 0.0:
+            x_grad = x_grad * (1.0 + reflect) - (ans_grad * reflect)
+        return x_grad, None, None, None, None
+
+
+def random_clamp(
+    x: Tensor,
+    min: Optional[float] = None,
+    max: Optional[float] = None,
+    prob: float = 0.5,
+    reflect: float = 0.0,
+):
+    return RandomClampFunction.apply(x, min, max, prob, reflect)
+
+
+def random_cast_to_half(x: Tensor, min_abs: float = 5.0e-06) -> Tensor:
+    """
+    A randomized way of casting a floating point value to half precision.
+    """
+    if x.dtype == torch.float16:
+        return x
+    x_abs = x.abs()
+    is_too_small = x_abs < min_abs
+    # for elements where is_too_small is true, random_val will contain +-min_abs with
+    # probability (x.abs() / min_abs), and 0.0 otherwise.  [so this preserves expectations,
+    # for those elements].
+    random_val = min_abs * x.sign() * (torch.rand_like(x) * min_abs < x_abs)
+    return torch.where(is_too_small, random_val, x).to(torch.float16)
+
+
+class RandomGradFunction(torch.autograd.Function):
+    """
+    Does nothing in forward pass; in backward pass, gets rid of very small grads using
+    randomized approach that preserves expectations (intended to reduce roundoff).
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, min_abs: float) -> Tensor:
+        ctx.min_abs = min_abs
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor) -> Tuple[Tensor, None]:
+        if ans_grad.dtype == torch.float16:
+            return (
+                random_cast_to_half(
+                    ans_grad.to(torch.float32), min_abs=ctx.min_abs
+                ),
+                None,
+            )
+        else:
+            return ans_grad, None
+
+
+class RandomGrad(torch.nn.Module):
+    """
+    Gets rid of very small gradients using an expectation-preserving method, intended to increase
+    accuracy of training when using amp (automatic mixed precision)
+    """
+
+    def __init__(self, min_abs: float = 5.0e-06):
+        super(RandomGrad, self).__init__()
+        self.min_abs = min_abs
+
+    def forward(self, x: Tensor):
+        if (
+            torch.jit.is_scripting()
+            or not self.training
+            or torch.jit.is_tracing()
+        ):
+            return x
+        else:
+            return RandomGradFunction.apply(x, self.min_abs)
+
+
+class SoftmaxFunction(torch.autograd.Function):
+    """
+    Tries to handle half-precision derivatives in a randomized way that should
+    be more accurate for training than the default behavior.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, dim: int):
+        ans = x.softmax(dim=dim)
+        # if x dtype is float16, x.softmax() returns a float32 because
+        # (presumably) that op does not support float16, and autocast
+        # is enabled.
+        if torch.is_autocast_enabled():
+            ans = ans.to(torch.float16)
+        ctx.save_for_backward(ans)
+        ctx.x_dtype = x.dtype
+        ctx.dim = dim
+        return ans
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        (ans,) = ctx.saved_tensors
+        with torch.cuda.amp.autocast(enabled=False):
+            ans_grad = ans_grad.to(torch.float32)
+            ans = ans.to(torch.float32)
+            x_grad = ans_grad * ans
+            x_grad = x_grad - ans * x_grad.sum(dim=ctx.dim, keepdim=True)
+            return x_grad, None
+
+
+def softmax(x: Tensor, dim: int):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x.softmax(dim)
+
+    return SoftmaxFunction.apply(x, dim)
+
+
+class MaxEigLimiterFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        coeffs: Tensor,
+        direction: Tensor,
+        channel_dim: int,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.channel_dim = channel_dim
+        ctx.grad_scale = grad_scale
+        ctx.save_for_backward(x.detach(), coeffs.detach(), direction.detach())
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad, *args):
+        with torch.enable_grad():
+            (x_orig, coeffs, new_direction) = ctx.saved_tensors
+            x_orig.requires_grad = True
+            num_channels = x_orig.shape[ctx.channel_dim]
+            x = x_orig.transpose(ctx.channel_dim, -1).reshape(-1, num_channels)
+            new_direction.requires_grad = False
+            x = x - x.mean(dim=0)
+            x_var = (x ** 2).mean()
+            x_residual = x - coeffs * new_direction
+            x_residual_var = (x_residual ** 2).mean()
+            # `variance_proportion` is the proportion of the variance accounted for
+            # by the top eigen-direction.  This is to be minimized.
+            variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20)
+            variance_proportion.backward()
+        x_orig_grad = x_orig.grad
+        x_extra_grad = (
+            x_orig.grad
+            * ctx.grad_scale
+            * x_grad.norm()
+            / (x_orig_grad.norm() + 1.0e-20)
+        )
+        return x_grad + x_extra_grad.detach(), None, None, None, None
+
+
+class BasicNorm(torch.nn.Module):
+    """
+    This is intended to be a simpler, and hopefully cheaper, replacement for
+    LayerNorm.  The observation this is based on, is that Transformer-type
+    networks, especially with pre-norm, sometimes seem to set one of the
+    feature dimensions to a large constant value (e.g. 50), which "defeats"
+    the LayerNorm because the output magnitude is then not strongly dependent
+    on the other (useful) features.  Presumably the weight and bias of the
+    LayerNorm are required to allow it to do this.
+    So the idea is to introduce this large constant value as an explicit
+    parameter, that takes the role of the "eps" in LayerNorm, so the network
+    doesn't have to do this trick.  We make the "eps" learnable.
+    Args:
+       num_channels: the number of channels, e.g. 512.
+      channel_dim: the axis/dimension corresponding to the channel,
+        interprted as an offset from the input's ndim if negative.
+        shis is NOT the num_channels; it should typically be one of
+        {-2, -1, 0, 1, 2, 3}.
+       eps: the initial "epsilon" that we add as ballast in:
+             scale = ((input_vec**2).mean() + epsilon)**-0.5
+          Note: our epsilon is actually large, but we keep the name
+          to indicate the connection with conventional LayerNorm.
+       learn_eps: if true, we learn epsilon; if false, we keep it
+         at the initial value.
+    eps_min: float
+    eps_max: float
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int = -1,  # CAUTION: see documentation.
+        eps: float = 0.25,
+        learn_eps: bool = True,
+        eps_min: float = -3.0,
+        eps_max: float = 3.0,
+    ) -> None:
+        super(BasicNorm, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        if learn_eps:
+            self.eps = nn.Parameter(torch.tensor(eps).log().detach())
+        else:
+            self.register_buffer("eps", torch.tensor(eps).log().detach())
+        self.eps_min = eps_min
+        self.eps_max = eps_max
+
+    def forward(self, x: Tensor) -> Tensor:
+        assert x.shape[self.channel_dim] == self.num_channels
+        eps = self.eps
+        if self.training and random.random() < 0.25:
+            # with probability 0.25, in training mode, clamp eps between the min
+            # and max; this will encourage it to learn parameters within the
+            # allowed range by making parameters that are outside the allowed
+            # range noisy.
+
+            # gradients to allow the parameter to get back into the allowed region if it happens to exit it.
+            eps = eps.clamp(min=self.eps_min, max=self.eps_max)
+        scales = (
+            torch.mean(x ** 2, dim=self.channel_dim, keepdim=True) + eps.exp()
+        ) ** -0.5
+        return x * scales
+
+
+def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear:
+    """
+    Behaves like a constructor of a modified version of nn.Linear
+    that gives an easy way to set the default initial parameter scale.
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Linear(*args, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(
+                ans.bias, -0.1 * initial_scale, 0.1 * initial_scale
+            )
+    return ans
+
+
+def ScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Conv1d:
+    """
+    Behaves like a constructor of a modified version of nn.Conv1d
+    that gives an easy way to set the default initial parameter scale.
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(
+                ans.bias, -0.1 * initial_scale, 0.1 * initial_scale
+            )
+    return ans
+
+
+def TransposeScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+    )
+
+
+def ScaledConv1dTranspose(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+        Transpose(),
+    )
+
+
+def TransposeConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> Conv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def Conv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    ScaledConv1d -> Transpose
+    """
+    return nn.Sequential(
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class SRLinear(nn.Linear):
+    """https://arxiv.org/abs/2303.06296
+    Stabilizing Transformer Training by Preventing Attention Entropy Collapse
+    """
+
+    def __init__(self, in_features, out_features, bias=True, **kwargs):
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        self.register_buffer(
+            "u", nn.functional.normalize(torch.randn(in_features), dim=0)
+        )
+        with torch.no_grad():
+            sigma = self.get_sigma()
+        self.register_buffer("spectral_norm", sigma)
+        self.sigma = nn.Parameter(torch.ones(1))
+
+    def get_sigma(self):
+        with torch.no_grad():
+            u = self.u
+            v = self.weight.mv(u)
+            v = nn.functional.normalize(v, dim=0)
+            u = self.weight.T.mv(v)
+            u = nn.functional.normalize(u, dim=0)
+            self.u.data.copy_(u)
+        return torch.einsum("c,cd,d->", v, self.weight, u)
+
+    def get_weight(self):
+        sigma = self.get_sigma()
+        if self.training:
+            self.spectral_norm.data.copy_(sigma)
+        weight = (self.sigma / sigma) * self.weight
+        return weight
+
+    def forward(self, x):
+        return nn.functional.linear(x, self.get_weight(), self.bias)
+
+
+class SRConv1d(SRLinear):
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        kernel_size,
+        stride: int = 1,
+        padding: str = "same",
+        bias: bool = True,
+        **kwargs,
+    ):
+        in_features = in_features * kernel_size
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+
+    def forward(self, x):
+        in_features = self.in_features // self.kernel_size
+        weight = self.get_weight().view(
+            self.out_features, in_features, self.kernel_size
+        )
+        return nn.functional.conv1d(
+            x, weight, bias=self.bias, stride=self.stride, padding=self.padding
+        )
+
+
+def TransposeSRConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> SRConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def SRConv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    SRConv1d -> Transpose
+    """
+    return nn.Sequential(
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class ActivationBalancer(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to encourage, for
+    each channel, that it is positive at least a proportion `threshold` of the
+    time.  It does this by multiplying negative derivative values by up to
+    (1+max_factor), and positive derivative values by up to (1-max_factor),
+    interpolated from 1 at the threshold to those extremal values when none
+    of the inputs are positive.
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           min_positive: the minimum, per channel, of the proportion of the time
+               that (x > 0), below which we start to modify the derivatives.
+           max_positive: the maximum, per channel, of the proportion of the time
+               that (x > 0), above which we start to modify the derivatives.
+           max_factor: the maximum factor by which we modify the derivatives for
+              either the sign constraint or the magnitude constraint;
+              e.g. with max_factor=0.02, the the derivatives would be multiplied by
+              values in the range [0.98..1.02].
+           sign_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_positive and max_positive
+              are violated.
+           scale_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_abs and max_abs
+              are violated.
+           min_abs:  the minimum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+           max_abs:  the maximum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+          min_prob: determines the minimum probability with which we modify the
+             gradients for the {min,max}_positive and {min,max}_abs constraints,
+             on each forward().  This is done randomly to prevent all layers
+             from doing it at the same time.  Early in training we may use
+             higher probabilities than this; it will decay to this value.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        min_positive: float = 0.05,
+        max_positive: float = 0.95,
+        max_factor: float = 0.04,
+        sign_gain_factor: float = 0.01,
+        scale_gain_factor: float = 0.02,
+        min_abs: float = 0.2,
+        max_abs: float = 100.0,
+        min_prob: float = 0.1,
+    ):
+        super(ActivationBalancer, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.min_positive = min_positive
+        self.max_positive = max_positive
+        self.max_factor = max_factor
+        self.min_abs = min_abs
+        self.max_abs = max_abs
+        self.min_prob = min_prob
+        self.sign_gain_factor = sign_gain_factor
+        self.scale_gain_factor = scale_gain_factor
+
+        # count measures how many times the forward() function has been called.
+        # We occasionally sync this to a tensor called `count`, that exists to
+        # make sure it is synced to disk when we load and save the model.
+        self.cpu_count = 0
+        self.register_buffer("count", torch.tensor(0, dtype=torch.int64))
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (
+            torch.jit.is_scripting()
+            or not x.requires_grad
+            or torch.jit.is_tracing()
+        ):
+            return _no_op(x)
+
+        count = self.cpu_count
+        self.cpu_count += 1
+
+        if random.random() < 0.01:
+            # Occasionally sync self.cpu_count with self.count.
+            # count affects the decay of 'prob'.  don't do this on every iter,
+            # because syncing with the GPU is slow.
+            self.cpu_count = max(self.cpu_count, self.count.item())
+            self.count.fill_(self.cpu_count)
+
+        # the prob of doing some work exponentially decreases from 0.5 till it hits
+        # a floor at min_prob (==0.1, by default)
+        prob = max(self.min_prob, 0.5 ** (1 + (count / 4000.0)))
+
+        if random.random() < prob:
+            sign_gain_factor = 0.5
+            if self.min_positive != 0.0 or self.max_positive != 1.0:
+                sign_factor = _compute_sign_factor(
+                    x,
+                    self.channel_dim,
+                    self.min_positive,
+                    self.max_positive,
+                    gain_factor=self.sign_gain_factor / prob,
+                    max_factor=self.max_factor,
+                )
+            else:
+                sign_factor = None
+
+            scale_factor = _compute_scale_factor(
+                x.detach(),
+                self.channel_dim,
+                min_abs=self.min_abs,
+                max_abs=self.max_abs,
+                gain_factor=self.scale_gain_factor / prob,
+                max_factor=self.max_factor,
+            )
+            return ActivationBalancerFunction.apply(
+                x,
+                scale_factor,
+                sign_factor,
+                self.channel_dim,
+            )
+        else:
+            return _no_op(x)
+
+
+def penalize_abs_values_gt(x: Tensor, limit: float, penalty: float) -> Tensor:
+    """
+    Returns x unmodified, but in backprop will put a penalty for the excess of
+    the absolute values of elements of x over the limit "limit".  E.g. if
+    limit == 10.0, then if x has any values over 10 it will get a penalty.
+    Caution: the value of this penalty will be affected by grad scaling used
+    in automatic mixed precision training.  For this reasons we use this,
+    it shouldn't really matter, or may even be helpful; we just use this
+    to disallow really implausible values of scores to be given to softmax.
+    """
+    x_sign = x.sign()
+    over_limit = (x.abs() - limit) > 0
+    # The following is a memory efficient way to penalize the absolute values of
+    # x that's over the limit.  (The memory efficiency comes when you think
+    # about which items torch needs to cache for the autograd, and which ones it
+    # can throw away).  The numerical value of aux_loss as computed here will
+    # actually be larger than it should be, by limit * over_limit.sum(), but it
+    # has the same derivative as the real aux_loss which is penalty * (x.abs() -
+    # limit).relu().
+    aux_loss = penalty * ((x_sign * over_limit).to(torch.int8) * x)
+    # note: we don't do sum() here on aux)_loss, but it's as if we had done
+    # sum() due to how with_loss() works.
+    x = with_loss(x, aux_loss)
+    # you must use x for something, or this will be ineffective.
+    return x
+
+
+def _diag(x: Tensor):  # like .diag(), but works for tensors with 3 dims.
+    if x.ndim == 2:
+        return x.diag()
+    else:
+        (batch, dim, dim) = x.shape
+        x = x.reshape(batch, dim * dim)
+        x = x[:, :: dim + 1]
+        assert x.shape == (batch, dim)
+        return x
+
+
+def _whitening_metric(x: Tensor, num_groups: int):
+    """
+    Computes the "whitening metric", a value which will be 1.0 if all the eigenvalues of
+    of the centered feature covariance are the same within each group's covariance matrix
+    and also between groups.
+    Args:
+        x: a Tensor of shape (*, num_channels)
+     num_groups:  the number of groups of channels, a number >=1 that divides num_channels
+    Returns:
+        Returns a scalar Tensor that will be 1.0 if the data is "perfectly white" and
+    greater than 1.0 otherwise.
+    """
+    assert x.dtype != torch.float16
+    x = x.reshape(-1, x.shape[-1])
+    (num_frames, num_channels) = x.shape
+    assert num_channels % num_groups == 0
+    channels_per_group = num_channels // num_groups
+    x = x.reshape(num_frames, num_groups, channels_per_group).transpose(0, 1)
+    # x now has shape (num_groups, num_frames, channels_per_group)
+    # subtract the mean so we use the centered, not uncentered, covariance.
+    # My experience has been that when we "mess with the gradients" like this,
+    # it's better not do anything that tries to move the mean around, because
+    # that can easily cause instability.
+    x = x - x.mean(dim=1, keepdim=True)
+    # x_covar: (num_groups, channels_per_group, channels_per_group)
+    x_covar = torch.matmul(x.transpose(1, 2), x)
+    x_covar_mean_diag = _diag(x_covar).mean()
+    # the following expression is what we'd get if we took the matrix product
+    # of each covariance and measured the mean of its trace, i.e.
+    # the same as _diag(torch.matmul(x_covar, x_covar)).mean().
+    x_covarsq_mean_diag = (x_covar ** 2).sum() / (
+        num_groups * channels_per_group
+    )
+    # this metric will be >= 1.0; the larger it is, the less 'white' the data was.
+    metric = x_covarsq_mean_diag / (x_covar_mean_diag ** 2 + 1.0e-20)
+    return metric
+
+
+class WhiteningPenaltyFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        num_groups: int,
+        whitening_limit: float,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.save_for_backward(x)
+        ctx.num_groups = num_groups
+        ctx.whitening_limit = whitening_limit
+        ctx.grad_scale = grad_scale
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor):
+        (x_orig,) = ctx.saved_tensors
+        with torch.enable_grad():
+            with torch.cuda.amp.autocast(enabled=False):
+                x_detached = x_orig.to(torch.float32).detach()
+                x_detached.requires_grad = True
+
+                metric = _whitening_metric(x_detached, ctx.num_groups)
+
+                if random.random() < 0.005 or __name__ == "__main__":
+                    logging.info(
+                        f"Whitening: num_groups={ctx.num_groups}, num_channels={x_orig.shape[-1]}, "
+                        f"metric={metric.item():.2f} vs. limit={ctx.whitening_limit}"
+                    )
+
+                (metric - ctx.whitening_limit).relu().backward()
+                penalty_grad = x_detached.grad
+                scale = ctx.grad_scale * (
+                    x_grad.to(torch.float32).norm()
+                    / (penalty_grad.norm() + 1.0e-20)
+                )
+                penalty_grad = penalty_grad * scale
+        return x_grad + penalty_grad.to(x_grad.dtype), None, None, None
+
+
+class Whiten(nn.Module):
+    def __init__(
+        self,
+        num_groups: int,
+        whitening_limit: float,
+        prob: Union[float, Tuple[float, float]],
+        grad_scale: float,
+    ):
+        """
+        Args:
+          num_groups: the number of groups to divide the channel dim into before
+            whitening.  We will attempt to make the feature covariance
+            within each group, after mean subtraction, as "white" as possible,
+            while having the same trace across all groups.
+         whitening_limit: a value greater than 1.0, that dictates how much
+           freedom we have to violate the constraints.  1.0 would mean perfectly
+           white, with exactly the same trace across groups; larger values
+           give more freedom.  E.g. 2.0.
+         prob: the probability with which we apply the gradient modification
+           (also affects the grad scale).  May be supplied as a float,
+           or as a pair (min_prob, max_prob)
+          grad_scale: determines the scale on the gradient term from this object,
+            relative to the rest of the gradient on the attention weights.
+            E.g. 0.02 (you may want to use smaller values than this if prob is large)
+        """
+        super(Whiten, self).__init__()
+        assert num_groups >= 1
+        assert whitening_limit >= 1
+        assert grad_scale >= 0
+        self.num_groups = num_groups
+        self.whitening_limit = whitening_limit
+        if isinstance(prob, float):
+            assert 0 < prob <= 1
+            self.prob = prob
+        else:
+            (self.min_prob, self.max_prob) = prob
+            assert 0 < self.min_prob < self.max_prob <= 1
+            self.prob = self.max_prob
+
+        self.grad_scale = grad_scale
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        In the forward pass, this function just returns the input unmodified.
+        In the backward pass, it will modify the gradients to ensure that the
+        distribution in each group has close to (lambda times I) as the covariance
+        after mean subtraction, with the same lambda across groups.
+        For whitening_limit > 1, there will be more freedom to violate this
+        constraint.
+        Args:
+           x: the input of shape (*, num_channels)
+        Returns:
+            x, unmodified.   You should make sure
+        you use the returned value, or the graph will be freed
+        and nothing will happen in backprop.
+        """
+        if (
+            not x.requires_grad
+            or random.random() > self.prob
+            or self.grad_scale == 0
+        ):
+            return _no_op(x)
+        else:
+            if hasattr(self, "min_prob") and random.random() < 0.25:
+                # occasionally switch between min_prob and max_prob, based on whether
+                # we are above or below the threshold.
+                if (
+                    _whitening_metric(x.to(torch.float32), self.num_groups)
+                    > self.whitening_limit
+                ):
+                    # there would be a change to the grad.
+                    self.prob = self.max_prob
+                else:
+                    self.prob = self.min_prob
+
+            return WhiteningPenaltyFunction.apply(
+                x, self.num_groups, self.whitening_limit, self.grad_scale
+            )
+
+
+class WithLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: Tensor, y: Tensor):
+        ctx.y_shape = y.shape
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        return ans_grad, torch.ones(
+            ctx.y_shape, dtype=ans_grad.dtype, device=ans_grad.device
+        )
+
+
+def with_loss(x, y):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    # returns x but adds y.sum() to the loss function.
+    return WithLoss.apply(x, y)
+
+
+def _no_op(x: Tensor) -> Tensor:
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    else:
+        # a no-op function that will have a node in the autograd graph,
+        # to avoid certain bugs relating to backward hooks
+        return x.chunk(1, dim=-1)[0]
+
+
+class Identity(torch.nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return _no_op(x)
+
+
+class MaxEig(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to discourage
+    that any given direction in activation space accounts for more than
+    a specified proportion of the covariance (e.g. 0.2).
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           max_var_per_eig:  the maximum proportion of the variance of the
+               features/channels, after mean subtraction, that can come from
+               any given eigenvalue.
+           min_prob: the minimum probability with which we apply this during any invocation
+               of forward(), assuming last time we applied the constraint it was
+               not active; supplied for speed.
+           scale: determines the scale with which we modify the gradients, relative
+               to the existing / unmodified gradients
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        max_var_per_eig: float = 0.2,
+        min_prob: float = 0.01,
+        scale: float = 0.01,
+    ):
+        super(MaxEig, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.scale = scale
+        assert max_var_per_eig == 0.0 or max_var_per_eig > 1.0 / num_channels
+        self.max_var_per_eig = max_var_per_eig
+
+        # we figure out the dominant direction using the power method: starting with
+        # a random vector, keep multiplying by the covariance and renormalizing.
+        with torch.no_grad():
+            # arbitrary.. would use randn() but want to leave the rest of the model's
+            # random parameters unchanged for comparison
+            direction = torch.arange(num_channels).to(torch.float)
+            direction = direction / direction.norm()
+            self.register_buffer("max_eig_direction", direction)
+
+        self.min_prob = min_prob
+        # cur_prob is the current probability we'll use to apply the ActivationBalancer.
+        # We'll regress this towards prob, each tiem we try to apply it and it is not
+        # active.
+        self.cur_prob = 1.0
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (
+            torch.jit.is_scripting()
+            or self.max_var_per_eig <= 0
+            or random.random() > self.cur_prob
+            or torch.jit.is_tracing()
+        ):
+            return _no_op(x)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            eps = 1.0e-20
+            orig_x = x
+            x = x.to(torch.float32)
+            with torch.no_grad():
+                x = x.transpose(self.channel_dim, -1).reshape(
+                    -1, self.num_channels
+                )
+                x = x - x.mean(dim=0)
+                new_direction, coeffs = self._find_direction_coeffs(
+                    x, self.max_eig_direction
+                )
+                x_var = (x ** 2).mean()
+                x_residual = x - coeffs * new_direction
+                x_residual_var = (x_residual ** 2).mean()
+
+                # `variance_proportion` is the proportion of the variance accounted for
+                # by the top eigen-direction.
+                variance_proportion = (x_var - x_residual_var) / (
+                    x_var + 1.0e-20
+                )
+
+                # ensure new direction is nonzero even if x == 0, by including `direction`.
+                self._set_direction(
+                    0.1 * self.max_eig_direction + new_direction
+                )
+
+            if random.random() < 0.01 or __name__ == "__main__":
+                logging.info(
+                    f"variance_proportion = {variance_proportion.item()}, shape={tuple(orig_x.shape)}, cur_prob={self.cur_prob}"
+                )
+
+            if variance_proportion >= self.max_var_per_eig:
+                # The constraint is active.  Note, we should quite rarely
+                # reach here, only near the beginning of training if we are
+                # starting to diverge, should this constraint be active.
+                cur_prob = self.cur_prob
+                self.cur_prob = (
+                    1.0  # next time, do the update with probability 1.0.
+                )
+                return MaxEigLimiterFunction.apply(
+                    orig_x, coeffs, new_direction, self.channel_dim, self.scale
+                )
+            else:
+                # let self.cur_prob exponentially approach self.min_prob, as
+                # long as the constraint is inactive.
+                self.cur_prob = 0.75 * self.cur_prob + 0.25 * self.min_prob
+                return orig_x
+
+    def _set_direction(self, direction: Tensor):
+        """
+        Sets self.max_eig_direction to a normalized version of `direction`
+        """
+        direction = direction.detach()
+        direction = direction / direction.norm()
+        direction_sum = direction.sum().item()
+        if direction_sum - direction_sum == 0:  # no inf/nan
+            self.max_eig_direction[:] = direction
+        else:
+            logging.info(
+                f"Warning: sum of direction in MaxEig is {direction_sum}, "
+                "num_channels={self.num_channels}, channel_dim={self.channel_dim}"
+            )
+
+    def _find_direction_coeffs(
+        self, x: Tensor, prev_direction: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+            Figure out (an approximation to) the proportion of the variance of a set of
+            feature vectors that can be attributed to the top eigen-direction.
+            Args:
+             x: a Tensor of shape (num_frames, num_channels), with num_frames > 1.
+          prev_direction:  a Tensor of shape (num_channels,), that is our previous estimate
+                   of the top eigen-direction, or a random direction if this is the first
+                   iteration.  Does not have to be normalized, but should be nonzero.
+        Returns: (cur_direction, coeffs), where:
+             cur_direction: a Tensor of shape (num_channels,) that is the current
+                estimate of the top eigen-direction.
+             coeffs: a Tensor of shape (num_frames, 1) that minimizes, or
+                approximately minimizes, (x - coeffs * cur_direction).norm()
+        """
+        (num_frames, num_channels) = x.shape
+        assert num_channels > 1 and num_frames > 1
+        assert prev_direction.shape == (num_channels,)
+        # `coeffs` are the coefficients of `prev_direction` in x.
+        # actually represent the coeffs up to a constant positive factor.
+        coeffs = (x * prev_direction).sum(dim=1, keepdim=True) + 1.0e-10
+        cur_direction = (x * coeffs).sum(dim=0) / (
+            (coeffs ** 2).sum() + 1.0e-20
+        )
+        return cur_direction, coeffs
+
+
+class DoubleSwishFunction(torch.autograd.Function):
+    """
+      double_swish(x) = x * torch.sigmoid(x-1)
+    This is a definition, originally motivated by its close numerical
+    similarity to swish(swish(x)), where swish(x) =  x * sigmoid(x).
+    Memory-efficient derivative computation:
+     double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)
+     double_swish'(x) = d/dx double_swish(x) =  x * s'(x) + x' * s(x) = x * s'(x) + s(x).
+     Now, s'(x) = s(x) * (1-s(x)).
+     double_swish'(x) =  x * s'(x) + s(x).
+                      =  x * s(x) * (1-s(x)) + s(x).
+                     = double_swish(x) * (1-s(x)) + s(x)
+     ... so we just need to remember s(x) but not x itself.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor) -> Tensor:
+        requires_grad = x.requires_grad
+        x_dtype = x.dtype
+        if x.dtype == torch.float16:
+            x = x.to(torch.float32)
+
+        s = torch.sigmoid(x - 1.0)
+        y = x * s
+
+        if requires_grad:
+            deriv = y * (1 - s) + s
+            # notes on derivative of x * sigmoid(x - 1):
+            # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29
+            # min \simeq -0.043638.  Take floor as -0.043637 so it's a lower bund
+            # max \simeq 1.1990.   Take ceil to be 1.2 so it's an upper bound.
+            # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which
+            # floors), should be expectation-preserving.
+            floor = -0.043637
+            ceil = 1.2
+            d_scaled = (deriv - floor) * (
+                255.0 / (ceil - floor)
+            ) + torch.rand_like(deriv)
+            if __name__ == "__main__":
+                # for self-testing only.
+                assert d_scaled.min() >= 0.0
+                assert d_scaled.max() < 256.0
+            d_int = d_scaled.to(torch.uint8)
+            ctx.save_for_backward(d_int)
+        if x.dtype == torch.float16 or torch.is_autocast_enabled():
+            y = y.to(torch.float16)
+        return y
+
+    @staticmethod
+    def backward(ctx, y_grad: Tensor) -> Tensor:
+        (d,) = ctx.saved_tensors
+        # the same constants as used in forward pass.
+        floor = -0.043637
+        ceil = 1.2
+        d = d * ((ceil - floor) / 255.0) + floor
+        return y_grad * d
+
+
+class DoubleSwish(torch.nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        """Return double-swish activation function which is an approximation to Swish(Swish(x)),
+        that we approximate closely with x * sigmoid(x-1).
+        """
+        if torch.jit.is_scripting() or torch.jit.is_tracing():
+            return x * torch.sigmoid(x - 1.0)
+        return DoubleSwishFunction.apply(x)
+
+
+def BalancedDoubleSwish(
+    d_model, channel_dim=-1, max_abs=10.0, min_prob=0.25
+) -> nn.Sequential:
+    """
+    ActivationBalancer -> DoubleSwish
+    """
+    balancer = ActivationBalancer(
+        d_model, channel_dim=channel_dim, max_abs=max_abs, min_prob=min_prob
+    )
+    return nn.Sequential(
+        balancer,
+        DoubleSwish(),
+    )
+
+
+def _test_max_eig():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = MaxEig(
+            num_channels, 1, 0.5, scale=0.1  # channel_dim  # max_var_per_eig
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad, atol=1.0e-02)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_whiten():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"_test_whiten(): proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = Whiten(
+            1, 5.0, prob=1.0, grad_scale=0.1  # num_groups  # whitening_limit,
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_activation_balancer_sign():
+    probs = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = 1.0 * (
+        (2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0
+    )
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        probs.numel(),
+        channel_dim=0,
+        min_positive=0.05,
+        max_positive=0.95,
+        max_factor=0.2,
+        min_abs=0.0,
+    )
+
+    y_grad = torch.sign(torch.randn(probs.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_sign: x = ", x)
+    print("_test_activation_balancer_sign: y grad = ", y_grad)
+    print("_test_activation_balancer_sign: x grad = ", x.grad)
+
+
+def _test_activation_balancer_magnitude():
+    magnitudes = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze(
+        -1
+    )
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        magnitudes.numel(),
+        channel_dim=0,
+        min_positive=0.0,
+        max_positive=1.0,
+        max_factor=0.2,
+        min_abs=0.2,
+        max_abs=0.8,
+        min_prob=1.0,
+    )
+
+    y_grad = torch.sign(torch.randn(magnitudes.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_magnitude: x = ", x)
+    print("_test_activation_balancer_magnitude: y grad = ", y_grad)
+    print("_test_activation_balancer_magnitude: x grad = ", x.grad)
+
+
+def _test_basic_norm():
+    num_channels = 128
+    m = BasicNorm(num_channels=num_channels, channel_dim=1)
+
+    x = torch.randn(500, num_channels)
+
+    y = m(x)
+
+    assert y.shape == x.shape
+    x_rms = (x ** 2).mean().sqrt()
+    y_rms = (y ** 2).mean().sqrt()
+    print("x rms = ", x_rms)
+    print("y rms = ", y_rms)
+    assert y_rms < x_rms
+    assert y_rms > 0.5 * x_rms
+
+
+def _test_double_swish_deriv():
+    x = torch.randn(10, 12, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    m = DoubleSwish()
+
+    tol = (1.2 - (-0.043637)) / 255.0
+    torch.autograd.gradcheck(m, x, atol=tol)
+
+    # for self-test.
+    x = torch.randn(1000, 1000, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    y = m(x)
+
+
+def _test_softmax():
+    a = torch.randn(2, 10, dtype=torch.float64)
+    b = a.clone()
+    a.requires_grad = True
+    b.requires_grad = True
+    a.softmax(dim=1)[:, 0].sum().backward()
+    print("a grad = ", a.grad)
+    softmax(b, dim=1)[:, 0].sum().backward()
+    print("b grad = ", b.grad)
+    assert torch.allclose(a.grad, b.grad)
+
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    _test_softmax()
+    _test_whiten()
+    _test_max_eig()
+    _test_activation_balancer_sign()
+    _test_activation_balancer_magnitude()
+    _test_basic_norm()
+    _test_double_swish_deriv()

src/model/modules/siglip.py

@@ -0,0 +1,258 @@
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+
+
+class SiglipVisionConfig:
+
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        image_size=224,
+        patch_size=16,
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        num_image_tokens: int = None,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.num_image_tokens = num_image_tokens
+
+
+class SiglipVisionEmbeddings(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",  # This indicates no padding is added
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+        self.register_buffer(
+            "position_ids",
+            torch.arange(self.num_positions).expand((1, -1)),
+            persistent=False,
+        )
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        _, _, height, width = (
+            pixel_values.shape
+        )  # [Batch_Size, Channels, Height, Width]
+        # Convolve the `patch_size` kernel over the image, with no overlapping patches since the stride is equal to the kernel size
+        # The output of the convolution will have shape [Batch_Size, Embed_Dim, Num_Patches_H, Num_Patches_W]
+        # where Num_Patches_H = height // patch_size and Num_Patches_W = width // patch_size
+        patch_embeds = self.patch_embedding(pixel_values)
+        # [Batch_Size, Embed_Dim, Num_Patches_H, Num_Patches_W] -> [Batch_Size, Embed_Dim, Num_Patches]
+        # where Num_Patches = Num_Patches_H * Num_Patches_W
+        embeddings = patch_embeds.flatten(2)
+        # [Batch_Size, Embed_Dim, Num_Patches] -> [Batch_Size, Num_Patches, Embed_Dim]
+        embeddings = embeddings.transpose(1, 2)
+        # Add position embeddings to each patch. Each positional encoding is a vector of size [Embed_Dim]
+        embeddings = embeddings + self.position_embedding(self.position_ids)
+        # [Batch_Size, Num_Patches, Embed_Dim]
+        return embeddings
+
+
+class SiglipAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        self.scale = self.head_dim**-0.5  # Equivalent to 1 / sqrt(self.head_dim)
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        # hidden_states: [Batch_Size, Num_Patches, Embed_Dim]
+        batch_size, seq_len, _ = hidden_states.size()
+        # query_states: [Batch_Size, Num_Patches, Embed_Dim]
+        query_states = self.q_proj(hidden_states)
+        # key_states: [Batch_Size, Num_Patches, Embed_Dim]
+        key_states = self.k_proj(hidden_states)
+        # value_states: [Batch_Size, Num_Patches, Embed_Dim]
+        value_states = self.v_proj(hidden_states)
+        # query_states: [Batch_Size, Num_Heads, Num_Patches, Head_Dim]
+        query_states = query_states.view(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+
+        key_states = key_states.view(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+
+        value_states = value_states.view(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        # Calculate the attention using the formula Q * K^T / sqrt(d_k). attn_weights: [Batch_Size, Num_Heads, Num_Patches, Num_Patches]
+        attn_weights = (
+            torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
+        )
+
+        if attn_weights.size() != (batch_size, self.num_heads, seq_len, seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(batch_size, self.num_heads, seq_len, seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        # Apply the softmax row-wise. attn_weights: [Batch_Size, Num_Heads, Num_Patches, Num_Patches]
+        attn_weights = nn.functional.softmax(
+            attn_weights, dim=-1, dtype=torch.float32
+        ).to(query_states.dtype)
+        # Apply dropout only during training
+        attn_weights = nn.functional.dropout(
+            attn_weights, p=self.dropout, training=self.training
+        )
+        # Multiply the attention weights by the value states. attn_output: [Batch_Size, Num_Heads, Num_Patches, Head_Dim]
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (batch_size, self.num_heads, seq_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(batch_size, self.num_heads, seq_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+        # [Batch_Size, Num_Heads, Num_Patches, Head_Dim] -> [Batch_Size, Num_Patches, Num_Heads, Head_Dim]
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        # [Batch_Size, Num_Patches, Num_Heads, Head_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
+        attn_output = attn_output.reshape(batch_size, seq_len, self.embed_dim)
+        # [Batch_Size, Num_Patches, Embed_Dim]
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class SiglipMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Intermediate_Size]
+        hidden_states = self.fc1(hidden_states)
+        # hidden_states: [Batch_Size, Num_Patches, Intermediate_Size]
+        hidden_states = nn.functional.gelu(hidden_states, approximate="tanh")
+        # [Batch_Size, Num_Patches, Intermediate_Size] -> [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = self.fc2(hidden_states)
+
+        return hidden_states
+
+
+class SiglipEncoderLayer(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = SiglipAttention(config)
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+    # Ignore copy
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # residual: [Batch_Size, Num_Patches, Embed_Dim]
+        residual = hidden_states
+        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = self.layer_norm1(hidden_states)
+        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
+        # [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = residual + hidden_states
+        # residual: [Batch_Size, Num_Patches, Embed_Dim]
+        residual = hidden_states
+        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = self.layer_norm2(hidden_states)
+        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = self.mlp(hidden_states)
+        # [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class SiglipEncoder(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList(
+            [SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+
+    # Ignore copy
+    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
+        # inputs_embeds: [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = inputs_embeds
+
+        for encoder_layer in self.layers:
+            # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
+            hidden_states = encoder_layer(hidden_states)
+
+        return hidden_states
+
+
+class SiglipVisionTransformer(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = SiglipVisionEmbeddings(config)
+        self.encoder = SiglipEncoder(config)
+        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        # pixel_values: [Batch_Size, Channels, Height, Width] -> [Batch_Size, Num_Patches, Embed_Dim]
+        hidden_states = self.embeddings(pixel_values)
+
+        last_hidden_state = self.encoder(inputs_embeds=hidden_states)
+
+        last_hidden_state = self.post_layernorm(last_hidden_state)
+
+        return last_hidden_state
+
+
+class SiglipVisionModel(nn.Module):
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.vision_model = SiglipVisionTransformer(config)
+
+    def forward(self, pixel_values) -> Tuple:
+        # [Batch_Size, Channels, Height, Width] -> [Batch_Size, Num_Patches, Embed_Dim]
+        return self.vision_model(pixel_values=pixel_values)

src/model/modules/tokenizer.py

@@ -0,0 +1,149 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/data/tokenizer.py
+# Copyright    2023                            (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Pattern, Union
+
+import numpy as np
+import torch
+import torchaudio
+# from lhotse.features import FeatureExtractor
+# from lhotse.utils import Seconds, compute_num_frames
+from phonemizer.backend import EspeakBackend
+from phonemizer.backend.espeak.language_switch import LanguageSwitch
+from phonemizer.backend.espeak.words_mismatch import WordMismatch
+from phonemizer.punctuation import Punctuation
+from phonemizer.separator import Separator
+
+
+
+class TextTokenizer:
+    """Phonemize Text."""
+
+    def __init__(
+        self,
+        language="en-us",
+        backend="espeak",
+        separator=Separator(word="_", syllable="-", phone="|"),
+        preserve_punctuation=True,
+        punctuation_marks: Union[str, Pattern] = Punctuation.default_marks(),
+        with_stress: bool = False,
+        tie: Union[bool, str] = False,
+        language_switch: LanguageSwitch = "keep-flags",
+        words_mismatch: WordMismatch = "ignore",
+    ) -> None:
+        phonemizer = EspeakBackend(
+            language,
+            punctuation_marks=punctuation_marks,
+            preserve_punctuation=preserve_punctuation,
+            with_stress=with_stress,
+            tie=tie,
+            language_switch=language_switch,
+            words_mismatch=words_mismatch,
+        )
+        
+        self.backend = phonemizer
+        self.separator = separator
+
+    def to_list(self, phonemized: str) -> List[str]:
+        fields = []
+        for word in phonemized.split(self.separator.word):
+            # "ɐ    m|iː|n?"    ɹ|ɪ|z|ɜː|v; h|ɪ|z.
+            pp = re.findall(r"\w+|[^\w\s]", word, re.UNICODE)
+            fields.extend(
+                [p for p in pp if p != self.separator.phone]
+                + [self.separator.word]
+            )
+        assert len("".join(fields[:-1])) == len(phonemized) - phonemized.count(
+            self.separator.phone
+        )
+        return fields[:-1]
+
+    def __call__(self, text, strip=True) -> List[List[str]]:
+        if isinstance(text, str):
+            text = [text]
+
+        phonemized = self.backend.phonemize(
+            text, separator=self.separator, strip=strip, njobs=1
+        )
+        return [self.to_list(p) for p in phonemized]
+
+
+def tokenize_text(tokenizer: TextTokenizer, text: str) -> List[str]:
+    phonemes = tokenizer([text.strip()])
+    return phonemes[0]  # k2symbols
+
+def convert_audio(wav: torch.Tensor, sr: int, target_sr: int, target_channels: int):
+    assert wav.shape[0] in [1, 2], "Audio must be mono or stereo."
+    if target_channels == 1:
+        wav = wav.mean(0, keepdim=True)
+    elif target_channels == 2:
+        *shape, _, length = wav.shape
+        wav = wav.expand(*shape, target_channels, length)
+    elif wav.shape[0] == 1:
+        wav = wav.expand(target_channels, -1)
+    wav = torchaudio.transforms.Resample(sr, target_sr)(wav)
+    return wav
+
+class AudioTokenizer:
+    """EnCodec audio."""
+
+    def __init__(
+        self,
+        device: Any = None,
+        signature = None
+    ) -> None:
+        from audiocraft.solvers import CompressionSolver
+        model = CompressionSolver.model_from_checkpoint(signature)
+        self.sample_rate = model.sample_rate
+        self.channels = model.channels
+        
+        if not device:
+            device = torch.device("cpu")
+            if torch.cuda.is_available():
+                device = torch.device("cuda:0")
+
+        self._device = device
+
+        self.codec = model.to(device)
+
+    @property
+    def device(self):
+        return self._device
+
+    def encode(self, wav: torch.Tensor) -> torch.Tensor:
+        codes = self.codec.encode(wav.to(self.device))
+        return [(codes[0], None)]
+
+    def decode(self, frames: torch.Tensor) -> torch.Tensor:
+        frames = frames[0][0] # [1,4,T]
+        return self.codec.decode(frames)
+    
+
+
+def tokenize_audio(tokenizer: AudioTokenizer, audio_path: str, offset = -1, num_frames=-1):
+    # Load and pre-process the audio waveform
+    if offset != -1 and num_frames!=-1:
+        wav, sr = torchaudio.load(audio_path, frame_offset=offset, num_frames=num_frames)
+    else:
+        wav, sr = torchaudio.load(audio_path)
+    wav = convert_audio(wav, sr, tokenizer.sample_rate, tokenizer.channels)
+    wav = wav.unsqueeze(0)
+
+    # Extract discrete codes from EnCodec
+    with torch.no_grad():
+        encoded_frames = tokenizer.encode(wav)
+    return encoded_frames

src/model/modules/transformer.py

@@ -0,0 +1,690 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py, modified by Puyuan Peng 2024
+import copy
+import numbers
+from functools import partial
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .activation import MultiheadAttention
+from .scaling import ActivationBalancer, BalancedDoubleSwish
+from .scaling import BasicNorm as _BasicNorm
+
+_shape_t = Union[int, List[int], torch.Size]
+
+
+class LayerNorm(nn.Module):
+    __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
+    normalized_shape: Tuple[int, ...]
+    eps: float
+    elementwise_affine: bool
+
+    def __init__(
+        self,
+        normalized_shape: _shape_t,
+        eps: float = 1e-5,
+        elementwise_affine: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+            self.bias = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if self.elementwise_affine:
+            nn.init.ones_(self.weight)
+            nn.init.zeros_(self.bias)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                F.layer_norm(
+                    input,
+                    self.normalized_shape,
+                    self.weight,
+                    self.bias,
+                    self.eps,
+                ),
+                embedding,
+            )
+
+        assert embedding is None
+        return F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+
+    def extra_repr(self) -> str:
+        return (
+            "{normalized_shape}, eps={eps}, "
+            "elementwise_affine={elementwise_affine}".format(**self.__dict__)
+        )
+
+
+class AdaptiveLayerNorm(nn.Module):
+    r"""Adaptive Layer Normalization"""
+
+    def __init__(self, d_model, norm) -> None:
+        super(AdaptiveLayerNorm, self).__init__()
+        self.project_layer = nn.Linear(d_model, 2 * d_model)
+        self.norm = norm
+        self.d_model = d_model
+        self.eps = self.norm.eps
+
+    def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            weight, bias = torch.split(
+                self.project_layer(embedding),
+                split_size_or_sections=self.d_model,
+                dim=-1,
+            )
+            return (weight * self.norm(input) + bias, embedding)
+
+        weight, bias = torch.split(
+            self.project_layer(embedding),
+            split_size_or_sections=self.d_model,
+            dim=-1,
+        )
+        return weight * self.norm(input) + bias
+
+
+class BasicNorm(_BasicNorm):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BasicNorm, self).__init__(d_model, eps=eps)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                super(BasicNorm, self).forward(input),
+                embedding,
+            )
+
+        assert embedding is None
+        return super(BasicNorm, self).forward(input)
+
+
+class BalancedBasicNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BalancedBasicNorm, self).__init__()
+        self.balancer = ActivationBalancer(
+            d_model,
+            channel_dim=-1,
+            min_positive=0.45,
+            max_positive=0.55,
+            max_abs=6.0,
+        )
+        self.norm = BasicNorm(d_model, eps, device=device, dtype=dtype)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return self.norm((self.balancer(input), embedding))
+
+        assert embedding is None
+        return self.norm(self.balancer(input))
+
+
+class IdentityNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ) -> None:
+        super(IdentityNorm, self).__init__()
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            return input
+
+        assert embedding is None
+        return input
+
+
+class TransformerEncoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            activation = BalancedDoubleSwish(d_model)
+
+        # # We can't test self.activation in forward() in TorchScript,
+        # # so stash some information about it instead.
+        # if activation is F.relu or isinstance(activation, torch.nn.ReLU):
+        #     self.activation_relu_or_gelu = 1
+        # elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
+        #     self.activation_relu_or_gelu = 2
+        # else:
+        #     self.activation_relu_or_gelu = 0
+        self.activation = activation
+
+        norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)
+        if layer_norm_cls == IdentityNorm:
+            norm2 = BalancedBasicNorm(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+        else:
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+        if adaptive_layer_norm:
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+        else:
+            self.norm1 = norm1
+            self.norm2 = norm2
+
+    def __setstate__(self, state):
+        super(TransformerEncoderLayer, self).__setstate__(state)
+        if not hasattr(self, "activation"):
+            self.activation = F.relu
+
+    def forward(
+        self,
+        src: Tensor,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        need_weights: Optional[bool] = False,
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layer.
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+        Shape:
+            see the docs in Transformer class.
+        """
+        x, stage_embedding = src, None
+        is_src_tuple = False
+        if isinstance(src, tuple):
+            x, stage_embedding = src
+            is_src_tuple = True
+
+        if src_key_padding_mask is not None:
+            _skpm_dtype = src_key_padding_mask.dtype
+            if _skpm_dtype != torch.bool and not torch.is_floating_point(
+                src_key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        if need_weights:
+            if self.norm_first:
+                out, attn = self._sa_block_attn(
+                    self.norm1(x, stage_embedding),
+                    src_mask,
+                    src_key_padding_mask,
+                    past
+                )
+                out, present = out # present is the kvcache of the present timestep
+                x = x + out
+                x = x + self._ff_block(self.norm2(x, stage_embedding))
+            else:
+                out, attn = self._sa_block_attn(x, src_mask, src_key_padding_mask, past)
+                out, present = out # present is the kvcache of the present timestep
+                x = self.norm1(
+                    x + out,
+                    stage_embedding,
+                )
+                x = self.norm2(x + self._ff_block(x), stage_embedding)
+            assert not is_src_tuple
+                # return (x, stage_embedding)
+            return (x, attn)
+        else:
+            if self.norm_first:
+                out = self._sa_block(
+                    self.norm1(x, stage_embedding),
+                    src_mask,
+                    src_key_padding_mask, past
+                )
+                out, present = out # present is the kvcache of the present timestep
+                x = x + out
+                x = x + self._ff_block(self.norm2(x, stage_embedding))
+            else:
+                out = self._sa_block(x, src_mask, src_key_padding_mask)
+                out, present = out # present is the kvcache of the present timestep
+                x = self.norm1(
+                    x + out,
+                    stage_embedding, past
+                )
+                x = self.norm2(x + self._ff_block(x), stage_embedding)
+            
+            if is_src_tuple:
+                x = (x, stage_embedding)
+            if present != None:
+                x = [x, present]
+            return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+            past=past
+        )
+        x, present = x
+        return self.dropout1(x), present
+
+    # self-attention block, also return attention weights
+    def _sa_block_attn(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        x, attn = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=True,
+            past=past
+        )
+        x, present = x
+        return (self.dropout1(x), present), attn
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout2(x)
+
+
+class TransformerEncoder(nn.Module):
+    r"""TransformerEncoder is a stack of N encoder layers. Users can build the
+    BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+        enable_nested_tensor: if True, input will automatically convert to nested tensor
+            (and convert back on output). This will improve the overall performance of
+            TransformerEncoder when padding rate is high. Default: ``True`` (enabled).
+    Examples::
+        >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+    __constants__ = ["norm"]
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(
+        self,
+        src: Tensor,
+        mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        return_layer_states: bool = False,
+        need_weights:Optional[bool] = False,
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+            return_layer_states: return layers' state (optional).
+        Shape:
+            see the docs in Transformer class.
+        """
+        if return_layer_states:
+            assert not need_weights
+            layer_states = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    past=past
+                )
+                layer_states.append(output[0])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_states, output
+        if need_weights:
+            assert not return_layer_states
+            layer_attn = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    need_weights=True,
+                    past=past
+                )
+                layer_attn.append(output[1])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_attn, output
+        
+        output = src
+        all_present = []
+        for n_layer, mod in enumerate(self.layers):
+            output = mod(
+                output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, past=None if past is None else past[n_layer]
+            )
+            if isinstance(output, list):
+                output, present = output
+                all_present.append(present)
+
+        if self.norm is not None:
+            output = self.norm(output)
+        if all_present != []:
+            all_present = torch.stack(all_present, dim=0) # (num_layers, 2, batch_size, num_heads, seq_len, head_dim)
+            output = [output, all_present]
+        return output
+
+
+class TransformerDecoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerDecoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        self.multihead_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            self.activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            self.activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            self.activation = BalancedDoubleSwish(d_model)
+        else:
+            self.activation = activation
+
+        if adaptive_layer_norm:
+            norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm3 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+            self.norm3 = AdaptiveLayerNorm(d_model, norm3)
+        else:
+            self.norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            self.norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            if layer_norm_cls == IdentityNorm:
+                self.norm3 = BalancedBasicNorm(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+            else:
+                self.norm3 = layer_norm_cls(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+
+    def forward(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the inputs (and mask) through the decoder layer.
+        Args:
+            tgt: the sequence to the decoder layer (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+        Shape:
+            see the docs in Transformer class.
+        """
+        tgt_is_tuple = False
+        if isinstance(tgt, tuple):
+            x, stage_embedding = tgt
+            tgt_is_tuple = True
+        else:
+            x, stage_embedding = tgt, None
+
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x, stage_embedding), tgt_mask, tgt_key_padding_mask
+            )
+            x = x + self._mha_block(
+                self.norm2(x, stage_embedding),
+                memory,
+                memory_mask,
+                memory_key_padding_mask,
+            )
+            x = x + self._ff_block(self.norm3(x, stage_embedding))
+        else:
+            x = self.norm1(
+                x + self._sa_block(x, tgt_mask, tgt_key_padding_mask),
+                stage_embedding,
+            )
+            x = self.norm2(
+                x
+                + self._mha_block(
+                    x, memory, memory_mask, memory_key_padding_mask
+                ),
+                stage_embedding,
+            )
+            x = self.norm3(x + self._ff_block(x), stage_embedding)
+
+        if tgt_is_tuple:
+            return (x, stage_embedding)
+        return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout1(x)
+
+    # multihead attention block
+    def _mha_block(
+        self,
+        x: Tensor,
+        mem: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+    ) -> Tensor:
+        x = self.multihead_attn(
+            x,
+            mem,
+            mem,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout2(x)
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout3(x)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError(
+        "activation should be relu/gelu, not {}".format(activation)
+    )

src/model/modules/voicecraft.py

@@ -0,0 +1,1999 @@
+# cp from https://github.com/jasonppy/VoiceCraft/blob/master/models/voicecraft.py
+
+import random
+
+import numpy as np
+import logging
+import argparse, copy
+from typing import Dict, Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchmetrics.classification import MulticlassAccuracy
+
+from .codebooks_patterns import DelayedPatternProvider
+
+from ...utils.util import make_pad_mask
+
+from .embedding import SinePositionalEmbedding, TokenEmbedding
+from .transformer import (
+    LayerNorm,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+
+from argparse import Namespace
+from huggingface_hub import PyTorchModelHubMixin
+
+
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = filter_value
+    return logits
+
+
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    return token
+
+
+class VoiceCraft(
+    nn.Module,
+    PyTorchModelHubMixin,
+    library_name="voicecraft",
+    repo_url="https://github.com/jasonppy/VoiceCraft",
+    tags=["text-to-speech"],
+):
+    def __new__(
+        cls, args: Optional[Namespace] = None, config: Optional[Dict] = None, **kwargs
+    ) -> "VoiceCraft":
+        # If initialized from Namespace args => convert to dict config for 'PyTorchModelHubMixin' to serialize it as config.json
+        # Won't affect instance initialization
+        if args is not None:
+            if config is not None:
+                raise ValueError("Cannot provide both `args` and `config`.")
+            config = vars(args)
+        return super().__new__(cls, args=args, config=config, **kwargs)
+
+    def __init__(self, args: Optional[Namespace] = None, config: Optional[Dict] = None):
+        super().__init__()
+
+        # If loaded from HF Hub => convert config.json to Namespace args before initializing
+        if args is None:
+            if config is None:
+                raise ValueError("Either `args` or `config` must be provided.")
+            args = Namespace(**config)
+
+        self.args = copy.copy(args)
+        self.pattern = DelayedPatternProvider(n_q=self.args.n_codebooks)
+        if not getattr(self.args, "special_first", False):
+            self.args.special_first = 0
+        if not getattr(self.args, "n_special", False):
+            self.args.n_special = 3
+        self.args.eos = getattr(self.args, "eos", -1)
+        self.eog = nn.Parameter(
+            torch.full((self.args.n_codebooks, 1), self.args.eog, dtype=torch.long),
+            requires_grad=False,
+        )  # [K 1]
+        if self.args.eos > 0:
+            assert (
+                self.args.eos != self.args.audio_pad_token
+                and self.args.eos != self.args.empty_token
+            ), self.args.eos
+            self.eos = nn.Parameter(
+                torch.full((self.args.n_codebooks, 1), self.args.eos, dtype=torch.long),
+                requires_grad=False,
+            )  # [K 1]
+        if isinstance(self.args.audio_vocab_size, str):
+            self.args.audio_vocab_size = eval(self.args.audio_vocab_size)
+
+        self.n_text_tokens = self.args.text_vocab_size + 1
+        assert (
+            self.args.text_pad_token == self.args.text_vocab_size
+        ), f"self.args.text_vocab_size: {self.args.text_vocab_size}, self.args.text_pad_token: {self.args.text_pad_token}"
+
+        self.n_audio_tokens = [
+            self.args.audio_vocab_size + self.args.n_special
+        ] * self.args.n_codebooks  # special tokens: empty token, EOG token, audio pad token
+        assert (
+            self.args.audio_vocab_size == self.args.empty_token
+        ), self.args.empty_token
+        assert self.args.eog == self.args.audio_vocab_size + 1, self.args.eog
+        assert (
+            self.args.audio_pad_token == self.args.audio_vocab_size + 2
+        ), self.args.audio_pad_token
+
+        self.text_embedding = TokenEmbedding(
+            dim_model=self.args.d_model,
+            vocab_size=self.n_text_tokens,
+            dropout=self.args.text_embedding_dropout,
+        )
+
+        self.audio_embedding = nn.ModuleList(
+            [
+                TokenEmbedding(
+                    dim_model=self.args.audio_embedding_dim,
+                    vocab_size=self.n_audio_tokens[k],
+                    dropout=self.args.audio_embedding_dropout,
+                )
+                for k in range(self.args.n_codebooks)
+            ]
+        )
+        self.mask_embedding = nn.Parameter(
+            torch.randn(self.args.max_n_spans, self.args.d_model), requires_grad=True
+        )
+        self.text_positional_embedding = SinePositionalEmbedding(
+            self.args.d_model,
+            dropout=self.args.text_positional_embedding_dropout,
+            scale=False,
+            alpha=True,  # learnable scaler, scale the volume of positional embedding
+        )
+        self.audio_positional_embedding = SinePositionalEmbedding(
+            self.args.d_model,
+            dropout=self.args.audio_positional_embedding_dropout,
+            scale=False,
+            alpha=True,  # learnable scaler, scale the volume of positional embedding
+        )
+
+        dec_layer = TransformerEncoderLayer(
+            self.args.d_model,
+            self.args.nhead,
+            dim_feedforward=self.args.d_model * 4,
+            dropout=self.args.trm_dropout,
+            batch_first=True,
+            norm_first=True,
+            layer_norm_cls=LayerNorm,
+        )
+        self.decoder = TransformerEncoder(
+            dec_layer,
+            num_layers=self.args.num_decoder_layers,
+            norm=LayerNorm(self.args.d_model),
+        )
+
+        self.predict_layer = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(self.args.d_model, self.args.audio_vocab_size // 2),
+                    nn.GELU(),
+                    nn.Linear(self.args.audio_vocab_size // 2, self.n_audio_tokens[k]),
+                )
+                for k in range(self.args.n_codebooks)
+            ]
+        )
+
+        self.accuracy_metrics = nn.ModuleList(
+            [
+                MulticlassAccuracy(
+                    self.n_audio_tokens[k],
+                    top_k=10,
+                    average="micro",
+                    multidim_average="global",
+                    ignore_index=None,
+                )
+                for k in range(self.args.n_codebooks)
+            ]
+        )
+
+    def prepare_mask_intervals(self, y_lens):
+        mask_intervals = []
+        non_mask_intervals = []
+
+        for i, y_len in enumerate(y_lens):
+            if self.args.mask_sample_dist == "uniform":
+                n_spans = random.choice(range(1, self.args.max_n_spans + 1))
+            elif "poisson" in self.args.mask_sample_dist.lower():
+                param = float(self.args.mask_sample_dist[len("poisson") :])
+                poisson_sample = torch.poisson(torch.tensor([param]))
+                n_spans = int(poisson_sample.clamp(1, self.args.max_n_spans).item())
+
+            starts = random.sample(
+                range(1, y_len - 1 - self.args.mask_len_min), n_spans
+            )
+            starts = sorted(starts)
+
+            for j in range(len(starts) - 1, 0, -1):
+                if starts[j] - starts[j - 1] < self.args.min_gap:
+                    del starts[j]  # If elements are too close, delete the later one
+            assert (
+                len(starts) > 0
+            ), f"there is no masked span left, y_len: {y_len}, sampled n_spans: {n_spans}"
+
+            temp_starts = starts + [y_len]
+            gaps = [
+                temp_starts[j + 1] - temp_starts[j] for j in range(len(temp_starts) - 1)
+            ]
+
+            ends = []
+
+            for j, (start, gap) in enumerate(zip(starts, gaps)):
+                mask_len = random.randint(
+                    self.args.mask_len_min, self.args.mask_len_max
+                )
+                # if mask_len > gap * self.args.max_mask_portion: # make sure the masks are not overlapping with each other
+                if (
+                    mask_len > gap - 1
+                ):  # make sure the masks are not overlapping with each other
+                    # temp_mask_start = int(0.6*gap*self.args.max_mask_portion)
+                    # temp_mask_end = int(gap*self.args.max_mask_portion)
+                    temp_mask_start = 1
+                    temp_mask_end = gap - 1
+                    mask_len = random.randint(temp_mask_start, temp_mask_end)
+                ends.append(start + mask_len)
+
+            mask_intervals.append([(s, e) for s, e in zip(starts, ends)])
+            non_mask_intervals.append(
+                [(ns, ne) for ns, ne in zip([0] + ends, starts + [y_len])]
+            )
+
+        return mask_intervals, non_mask_intervals
+
+    def rearrange(self, y, non_mask_intervals, mask_intervals):
+        reduced_eog = getattr(self.args, "reduced_eog", 0)
+        rearranged_y = []
+        for i in range(len(y)):
+            if self.args.eos > 0:
+                assert reduced_eog
+                cur_y = (
+                    [y[i, :, item[0] : item[1]] for item in non_mask_intervals[i][:-1]]
+                    + [
+                        torch.cat(
+                            [
+                                y[
+                                    i,
+                                    :,
+                                    non_mask_intervals[i][-1][0] : non_mask_intervals[
+                                        i
+                                    ][-1][1],
+                                ],
+                                self.eos,
+                            ],
+                            dim=-1,
+                        )
+                    ]
+                    + [
+                        torch.cat([y[i, :, item[0] : item[1]], self.eog], dim=-1)
+                        for item in mask_intervals[i]
+                    ]
+                )  # only insert eog to the last non-mask-interval, which is when the utterance actual ends
+            else:
+                if reduced_eog:
+                    cur_y = (
+                        [
+                            y[i, :, item[0] : item[1]]
+                            for item in non_mask_intervals[i][:-1]
+                        ]
+                        + [
+                            torch.cat(
+                                [
+                                    y[
+                                        i,
+                                        :,
+                                        non_mask_intervals[i][-1][
+                                            0
+                                        ] : non_mask_intervals[i][-1][1],
+                                    ],
+                                    self.eog,
+                                ],
+                                dim=-1,
+                            )
+                        ]
+                        + [
+                            torch.cat([y[i, :, item[0] : item[1]], self.eog], dim=-1)
+                            for item in mask_intervals[i]
+                        ]
+                    )  # only insert eog to the last non-mask-interval, which is when the utterance actual ends
+                else:
+                    cur_y = [
+                        torch.cat([y[i, :, item[0] : item[1]], self.eog], dim=-1)
+                        for item in non_mask_intervals[i]
+                    ] + [
+                        torch.cat([y[i, :, item[0] : item[1]], self.eog], dim=-1)
+                        for item in mask_intervals[i]
+                    ]  # eog is added to each section TODO this is not correct, I should add eog to non_mask_intervals if that segment is not the ending segment (as there is no way for the model to predict eog for those segments, and this will do harm to tts experiment, where the model randomly output eog for the first segment)
+            rearranged_y.append(cur_y)
+        return rearranged_y
+
+    def shift(self, rearranged_y):
+        shifted_y = []
+        patterns = []
+        for i in range(len(rearranged_y)):
+            cur_patterns = [
+                self.pattern.get_pattern(cur_y.shape[1]) for cur_y in rearranged_y[i]
+            ]
+            out = [
+                cur_pattern.build_pattern_sequence(
+                    z=cur_y.unsqueeze(0).contiguous(),
+                    special_token=self.args.empty_token,
+                    keep_only_valid_steps=False,
+                )
+                for cur_pattern, cur_y in zip(cur_patterns, rearranged_y[i])
+            ]
+            shifted_y.append(
+                [item[0].squeeze(0) for item in out]
+            )  # the first item is values, later two are indexes and mask
+            patterns.append(cur_patterns)
+        return shifted_y, patterns
+
+    def insert_mask(self, shifted_y):
+        inserted_y = []
+        mask_position = []
+        mask_value = []
+        for i in range(len(shifted_y)):
+            num_masks = (len(shifted_y[i]) - 1) // 2
+            assert num_masks == (len(shifted_y[i]) - 1) / 2, len(shifted_y[i])
+            emb_inds = list(range(self.args.max_n_spans))
+            if self.args.shuffle_mask_embedding:
+                random.shuffle(emb_inds)
+            emb_inds_use = emb_inds[:num_masks]
+            emb_inds_use = emb_inds_use + emb_inds_use
+            mask_value.append(emb_inds_use)
+            cur_inserted_y = []
+            cur_mask_position = []
+            for j in range(len(shifted_y[i]) - 1):
+                cur_inserted_y.append(shifted_y[i][j])
+                cur_mask_position.append(
+                    sum([item.shape[1] for item in cur_inserted_y])
+                )  # each item is of shape [K S], so take shape[1]
+                cur_inserted_y.append(
+                    self.eog
+                )  # insert mask token of shape [K, 1], BUT we are actually using the eog token as a place holder here, as the real mask will be inserted in embed_y function
+
+            cur_inserted_y.append(shifted_y[i][-1])
+
+            inserted_y.append(cur_inserted_y)
+            mask_position.append(cur_mask_position)
+        return inserted_y, mask_position, mask_value
+
+    def cat_y(self, inserted_y, mask_position, y_lens):
+        reduced_eog = getattr(self.args, "reduced_eog", 0)
+        cated_y = []
+        new_y_lens = []
+        for i in range(len(inserted_y)):
+            cur_cated_y = torch.cat(inserted_y[i], dim=1)  # [K S]
+            cur_cated_y = cur_cated_y.transpose(1, 0)  # [S K]
+            cur_cated_y_len = cur_cated_y.shape[0]
+            if reduced_eog:
+                assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (
+                    len(mask_position[i]) + 1
+                ) * self.args.n_codebooks + (
+                    len(mask_position[i]) / 2 + 1
+                ), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i])/2 + 1) ({len(mask_position[i])/2 + 1})={y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i])/2 + 1)}"
+            else:
+                assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (
+                    len(mask_position[i]) + 1
+                ) * self.args.n_codebooks + (
+                    len(mask_position[i]) + 1
+                ), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i]) + 1) ({len(mask_position[i]) + 1})"  # the last term represent the inserted eog token, originally it's inserted at the end of every token, but this is wrong
+            new_y_lens.append(cur_cated_y_len)
+            cated_y.append(cur_cated_y)
+
+        cated_y = torch.nn.utils.rnn.pad_sequence(
+            cated_y, batch_first=False, padding_value=self.args.audio_pad_token
+        )
+        assert cated_y.shape == torch.Size(
+            [max(new_y_lens), len(inserted_y), self.args.n_codebooks]
+        ), f"cated_y.shape: {cated_y.shape}, but it should be {torch.Size([max(new_y_lens,len(inserted_y), self.args.n_codebooks)])}"
+        cated_y = cated_y.permute(2, 0, 1)  # [T,B,K]->[K,T,B]
+        assert cated_y.shape[0] == self.args.n_codebooks, cated_y.shape
+        return cated_y, torch.LongTensor(new_y_lens).to(cated_y.device)
+
+    def embed_y(self, cated_y, mask_position, mask_value):
+        embedded_y = torch.stack(
+            [self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)],
+            dim=0,
+        )  # [K, T, B, D]
+        assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
+        assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
+        embedded_y = embedded_y.sum(dim=0)  # [K,T,B,D]->[T,B,D]
+        embedded_y = embedded_y.transpose(1, 0)  # [T,B,D]->[B,T,D]
+        for i in range(len(embedded_y)):
+            if len(mask_position[i]) > 0:
+                embedded_y[i, mask_position[i]] = self.mask_embedding[mask_value[i]]
+        return embedded_y
+
+    def prepare_input_target(self, y, y_lens):
+        # rearrange y
+        # assume y shape: [B T K], K is n_codebooks
+        assert y.shape[1] == self.args.n_codebooks, y.shape
+        # sample mask_intervals
+        mask_intervals, non_mask_intervals = self.prepare_mask_intervals(y_lens)
+
+        # need to have EOG in each section (SOG will be generated by the pattern class)
+        # but mask can be inserted later after we have shifted the input
+        # y could be rearranged in this way:
+        # [
+        # [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]],
+        # [tensor[4, 44], tensor[4, 56], tensor[4, 19]],
+        # ...
+        # ]
+        # for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part.
+        # NOTE #non_masked_part = #masked_part + 1
+        # NOTE *these are also the targets*
+        # added eog at the end of each segment (masked segment and unmasked segment)
+        rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals)
+        targets = rearranged_y  # each element in each sample is of shape [K T]
+        assert targets[0][0].shape[0] == self.args.n_codebooks, targets[0][0].shape
+
+        # next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token
+        #  [[5, 1, 2, 3, 4, 5, 5],
+        #  [5, 5, 1, 2, 3, 4, 5],
+        #  [5, 5, 5, 1, 2, 3, 4]]
+        shifted_y, patterns = self.shift(rearranged_y)  # each element [K S]
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape[
+            0
+        ]
+
+        # then, insert mask token at the intersection of each tensor (we want to decide the arrangement of the mask (shuffle or not)), we better have a separate nn.embedding for it
+        # we also need to record the position of the inserted mask
+        inserted_y, mask_position, mask_value = self.insert_mask(shifted_y)
+        assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][
+            0
+        ].shape[0]
+        assert inserted_y[0][1].shape == torch.Size(
+            (self.args.n_codebooks, 1)
+        ), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}"
+
+        # then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token
+        cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens)  # KTB
+        assert cated_y.shape == torch.Size(
+            (self.args.n_codebooks, cated_y.shape[1], len(inserted_y))
+        )
+
+        # embed remember to separately embed the mask tokens
+        embedded_y = self.embed_y(cated_y, mask_position, mask_value)  # BTD
+        assert embedded_y.shape[1:] == torch.Size(
+            (max(new_y_lens), self.args.d_model)
+        ), embedded_y.shape
+
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_padding_mask = make_pad_mask(new_y_lens).to(y.device)
+        y_attention_mask = (
+            torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1)
+            .bool()
+            .to(y_padding_mask.device)
+        )
+        return (
+            y_input,
+            new_y_lens,
+            targets,
+            y_padding_mask,
+            y_attention_mask,
+            mask_position,
+            patterns,
+        )
+
+    def remove_mask(self, logits, mask_position, new_y_lens):
+        # logits: [B K S card]
+        logits_use = []
+        for i in range(len(logits)):
+            non_mask_positions = [-1] + mask_position[i] + [new_y_lens[i]]
+            non_mask_intervals = [
+                [non_mask_positions[i] + 1, non_mask_positions[i + 1]]
+                for i in range(len(non_mask_positions) - 1)
+            ]
+            cur_logits_use = [logits[i, :, l:r] for l, r in non_mask_intervals]
+            logits_use.append(cur_logits_use)
+
+        return logits_use
+
+    def revert_pattern(self, patterns, logits_use):
+        logits_final = []
+        logit_masks = []
+        for i in range(len(logits_use)):
+            cur_logits = [
+                item.unsqueeze(0).permute(0, 3, 1, 2).contiguous()
+                for item in logits_use[i]
+            ]  # each item is of shape [1 K S card] [1 card K S]
+            cur_logits_final = [
+                cur_pattern.revert_pattern_logits(item, 0, keep_only_valid_steps=False)
+                for cur_pattern, item in zip(patterns[i], cur_logits)
+            ]  # if input output order doesn't match, this step will give an error
+            cur_logits_final_ret = [
+                item[0].permute(0, 2, 3, 1).squeeze(0) for item in cur_logits_final
+            ]  # each element is of shape [K,T,card]
+            logits_final.append(cur_logits_final_ret)
+            logit_masks.append([item[2] for item in cur_logits_final])
+
+        return logits_final, logit_masks
+
+    def dec_forward(
+        self,
+        x_input,
+        x_lens,
+        x_attention_mask,
+        x_padding_mask,
+        y_input,
+        new_y_lens,
+        y_attention_mask,
+        y_padding_mask,
+        past=None,
+        last_3_tokens=False,
+    ):
+        x_attn_mask = F.pad(
+            x_attention_mask,
+            (0, new_y_lens.max()),
+            value=True,
+        )  # x attn to all x, doesn't attn to any y, this follow figure 3 of the valle paper
+        y_attn_mask = F.pad(
+            y_attention_mask,
+            (x_lens.max(), 0),  # y is padded at the front
+            value=False,
+        )  # y attn to all x, for y itself use lower triangle mask to ensure autoregressive
+        xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)
+
+        # merge key padding and attention masks
+        bsz, src_len = x_input.shape[0], x_lens.max() + new_y_lens.max()
+        xy_padding_mask = torch.concat([x_padding_mask, y_padding_mask], dim=1)
+        _xy_padding_mask = (
+            xy_padding_mask.view(bsz, 1, 1, src_len)
+            .expand(-1, self.args.nhead, -1, -1)
+            .reshape(bsz * self.args.nhead, 1, src_len)
+        )
+        # Check shapes and resize+broadcast as necessary
+        if xy_attn_mask.shape != _xy_padding_mask.shape:
+            assert (
+                xy_attn_mask.ndim + 1 == _xy_padding_mask.ndim
+            ), f"xy_attn_mask.shape: {xy_attn_mask.shape}, _xy_padding_mask: {_xy_padding_mask.shape}"
+            xy_attn_mask = xy_attn_mask.unsqueeze(0).repeat(
+                _xy_padding_mask.shape[0], 1, 1
+            )  # Example approach
+        xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)
+
+        new_attn_mask = torch.zeros_like(xy_attn_mask)
+        new_attn_mask.masked_fill_(xy_attn_mask, float("-inf"))
+        xy_attn_mask = new_attn_mask
+
+        xy_input = torch.cat([x_input, y_input], dim=1)
+
+        if past == None:  # do not use kvcache
+            out, _ = self.decoder((xy_input, None), mask=xy_attn_mask)
+            return out[:, x_lens.max() :], None
+        else:  # use kvcache
+            if (
+                past.ndim > 3
+            ):  # uses kvcache, only need to pass the last tokens, this doesn't work with multi-span speech editing yet
+                if last_3_tokens:
+                    xy_input = xy_input[:, -3:]
+                    xy_attn_mask = xy_attn_mask[:, -3:]
+                else:
+                    xy_input = xy_input[:, -1:]
+                    xy_attn_mask = xy_attn_mask[:, -1:]
+
+            out, present = self.decoder((xy_input, None), mask=xy_attn_mask, past=past)
+            if isinstance(out, tuple):  # get rid of stage_embedding
+                out = out[0]
+
+            if out.shape[1] > x_lens.max():  # the first pass, not kvcache yet
+                return out[:, x_lens.max() :], present
+            else:  # used kvcache
+                return out, present
+
+    def forward(self, batch):
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (N, S).
+          x_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (N, K, T).
+            where K is the number of codebooks
+          y_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+        """
+        x, x_lens, y, y_lens = batch["x"], batch["x_lens"], batch["y"], batch["y_lens"]
+        if len(x) == 0:
+            return None
+        x = x[
+            :, : x_lens.max()
+        ]  # this deal with gradient accumulation, where x_lens.max() might not be longer than the length of the current slice of x
+        y = y[:, :, : y_lens.max()]
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3 and y.shape[1] == self.args.n_codebooks, y.shape
+        assert y_lens.ndim == 1, y_lens.shape
+        # makes attention mask and padding mask for x
+        x_padding_mask = make_pad_mask(x_lens).to(x.device)
+        x_attention_mask = (
+            torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1)
+            .bool()
+            .to(x_padding_mask.device)
+        )
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+        (
+            y_input,
+            new_y_lens,
+            targets,
+            y_padding_mask,
+            y_attention_mask,
+            mask_position,
+            patterns,
+        ) = self.prepare_input_target(y, y_lens)
+        y_out = self.dec_forward(
+            x_input,
+            x_lens,
+            x_attention_mask,
+            x_padding_mask,
+            y_input,
+            new_y_lens,
+            y_attention_mask,
+            y_padding_mask,
+        )
+        y_out = y_out[0]  # no kv-caching during training
+        assert (
+            y_out.shape == y_input.shape
+        ), f"y_out.shape: {y_out.shape}, y_input.shape: {y_input.shape}"  # [B S D]
+
+        logits = torch.stack(
+            [self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1
+        )  # [B K S card]
+        # take out the mask token (using mask_position and new_y_lens) and revert (using function provided by self.pattern)
+        assert (
+            logits.shape[1] == self.args.n_codebooks
+            and logits.shape[3] == self.n_audio_tokens[0]
+        ), logits.shape
+
+        logits_use = self.remove_mask(logits, mask_position, new_y_lens)
+
+        # revert the pattern shift for each logits section in each sample
+        logits_final, logit_masks = self.revert_pattern(patterns, logits_use)
+        assert (
+            logits_final[0][0].shape[0] == self.args.n_codebooks
+            and logits_final[0][0].shape[2] == self.n_audio_tokens[0]
+        ), f"it is: {logits_final[0][0].shape}, but should be [K, T, card]"
+        # testing
+        sample_to_test = 0
+        assert len(logits_final[sample_to_test]) == len(
+            targets[sample_to_test]
+        ), f"{len(logits_final[sample_to_test])}, {len(targets[sample_to_test])}"
+        temp = sum(
+            [
+                logits_final[sample_to_test][i].shape[:-1]
+                != targets[sample_to_test][i].shape
+                for i in range(len(targets[sample_to_test]))
+            ]
+        )
+        assert (
+            temp == 0
+        ), f"none equal positions: {temp}, total number of elements: {len(targets[sample_to_test])}"
+
+        logit_masked = sum(
+            [(item == False).any() for cur_mask in logit_masks for item in cur_mask]
+        )
+        assert logit_masked == 0, logit_masks
+
+        logits = torch.cat(
+            [torch.cat(item, dim=1) for item in logits_final], dim=1
+        )  # [K, T1+T2+T3+..., card]
+        targets = torch.cat(
+            [torch.cat(item, dim=1) for item in targets], dim=1
+        )  # [K, T1+T2+T3+...]
+        assert targets.shape[0] == logits.shape[0], f"{targets.shape}, {logits.shape}"
+        loss = []
+        ntokens = []
+        top10acc = []
+        for k, (logit, target) in enumerate(zip(logits, targets)):
+            loss.append(F.cross_entropy(logit, target, reduction="mean"))
+            top10acc.append(self.accuracy_metrics[k](logit.detach(), target))
+            ntokens.append(len(logit))
+
+        all_ntokens = sum(ntokens)
+        if self.args.codebook_weight != None:
+            codebook_weight = eval(self.args.codebook_weight)
+        else:
+            codebook_weight = [1.0] * self.args.n_codebooks
+        loss = sum([l * nt * cw for l, nt, cw in zip(loss, ntokens, codebook_weight)])
+        top10acc_by_codebook = [t10a * nt for t10a, nt in zip(top10acc, ntokens)]
+        top10acc = sum(top10acc_by_codebook)
+        ntokens = torch.tensor(all_ntokens).to(logits.device)
+
+        return {
+            "loss": loss,
+            "top10acc": top10acc,
+            "top10acc_by_codebook": top10acc_by_codebook,
+            "effective_ntoken": ntokens,
+        }
+
+    def inference(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        mask_interval: list[torch.Tensor],
+        top_k: int = -100,
+        top_p: float = 1.0,
+        temperature: float = 1.0,
+        stop_repetition: int = -1,
+        kvcache: int = 1,
+        silence_tokens: list[int] = [1388, 1898, 131],
+    ) -> torch.Tensor:
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (1, L).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, K).
+          mask_interval:
+            a list of tensors of shape (M, 2). contains M mask_start and mask_end. list length is actually 1, because we only support single sample inference for now
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          top_p: (`optional`) float
+            For Neucleus sampling
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+          eog_coef: (`optional`) float
+            if 0, no change to eog token logits, otherwise, will adjust eog token logit based on the difference between acoustic token and phn token length
+          stop_repetition (`optional`) int
+            if not -1, will set the logits of a token that repeated this many times to be -100000, to avoid generating it again. This only apply to tokens from the first codebook
+          allowed_repeat_tokens (`optional`) list of ints
+            by inspecting the validation set, get a few tokens that indeed repeat a significant amount of time, and exclude those tokens from prevent repetition
+          ultimate_stop_repetition (`optional`) int
+            no matter that token it is, stop repetition once after this number
+        """
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        if self.args.special_first:
+            y = y + int(self.args.n_special)
+        y = y.transpose(2, 1)  # [1,T,K] -> [1,K,T]
+        assert (
+            y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks
+        ), y.shape  # there is no padding
+        assert mask_interval.shape == torch.Size(
+            (1, mask_interval.shape[1], 2)
+        ), mask_interval
+
+        # make x attention mask and x_input
+        x_attention_mask = (
+            torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1)
+            .bool()
+            .to(x.device)
+        )
+        # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+
+        # make initial y_input
+
+        # make mask_interval and non_mask_interval
+        y_len = y.shape[2]
+        y_lens = torch.LongTensor([y_len]).to(y.device)
+        mask_interval = mask_interval[0]
+        starts = [item[0].item() for item in mask_interval] + [y_len]
+        ends = [0] + [item[1].item() for item in mask_interval]
+        mask_intervals = [
+            [(item[0].item(), item[1].item()) for item in mask_interval]
+        ]  # a werid name change, mask_interval is input, now is mask_intervals, with one more dimension
+        non_mask_intervals = [[(ns, ne) for ns, ne in zip(ends, starts)]]
+
+        # rearrange y
+        # will add have EOG in each section (SOG will be generated by the pattern class)
+        # but mask can be inserted later after we have shifted the input
+        # y could be rearranged in this way:
+        # [
+        # [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]],
+        # [tensor[4, 44], tensor[4, 56], tensor[4, 19]],
+        # ...
+        # ]
+        # for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part.
+        # NOTE #non_masked_part = #masked_part + 1
+        rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals)
+        assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][
+            0
+        ].shape
+
+        # shift each element of y
+        # next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token
+        #  [
+        #  [empty, 1, 2, 3, eog, empty, empty, empty],
+        #  [empty, empty, 1, 2, 3, eog, empty, empty],
+        #  [empty, empty, empty, 1, 2, 3, eog, empty],
+        #  [empty, empty, empty, empty, 1, 2, 3, eog]
+        # ]
+        shifted_y, patterns = self.shift(
+            rearranged_y
+        )  # each element [K S], patterns is not used, as we directly use the original input y
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
+
+        # insert mask token at the intersction of each tensor, but *actually inserted eog as place holder*
+        # the position of inserted mask is also recorded
+        # and the mask_value, the index of the mask emb is recorded
+        inserted_y, mask_position, mask_value = self.insert_mask(shifted_y)
+        assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][
+            0
+        ].shape[0]
+        assert inserted_y[0][1].shape == torch.Size(
+            (self.args.n_codebooks, 1)
+        ), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}"
+
+        # then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token
+        cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens)  # KTB
+        assert cated_y.shape == torch.Size(
+            (self.args.n_codebooks, cated_y.shape[1], len(inserted_y))
+        )
+        assert not (cated_y == self.args.audio_pad_token).any(), cated_y
+
+        ### NOTE this is different from forward, as we will remove the masked tokens
+        ### say there are two masked region
+        ### the cated_y should be like
+        ### [empty a a a a mask0 empty b b b mask1 empty c c mask0 empty]
+        ### which means we need to take the part after the last empty out
+        num_mask = len(mask_position[0]) // 2
+        assert num_mask == len(mask_position[0]) / 2, mask_position
+        cated_y = cated_y[:, : mask_position[0][num_mask] + 2]  # of shape [K,T,B]
+        # logging.info(f"mask_position[0][num_mask]+2: {mask_position[0][num_mask]+2}")
+        more_mask_value = mask_value[0][
+            num_mask + 1 :
+        ]  # NOTE this will be used in the generation loop for reference for inserting mask embedding
+        new_y_lens[0] = mask_position[0][num_mask] + 2
+        mask_position[0] = mask_position[0][: num_mask + 1]
+        assert (
+            mask_position[0][num_mask] + 2 == cated_y.shape[1]
+        ), f"num_mask: {num_mask}, mask_position: {mask_position}, cated_y.shape: {cated_y.shape}"
+
+        # embed: remember to separately embed the mask tokens
+        embedded_y = self.embed_y(
+            cated_y, mask_position, [mask_value[0][: num_mask + 1]]
+        )  # BTD
+        # assert embedded_y.shape == torch.Size((y.shape[0], max(new_y_lens), self.args.d_model)), embedded_y.shape
+
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_attention_mask = (
+            torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1)
+            .bool()
+            .to(y.device)
+        )
+        # y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
+
+        x_padding_mask = torch.full((1, x_lens[0]), False).to(x.device)
+        y_padding_mask = torch.full((1, new_y_lens[0]), False).to(y.device)
+
+        codebook_eog = [False] * self.args.n_codebooks
+        generated = []  # doesn't contain any empty_token, contains eog
+        cur_generated = []
+        # say 0 is empty, 4 is eog
+        # tensor([[ 1,  2,  3,  4,  0,  0],
+        #         [ 0,  1,  2,  3,  4,  0],
+        #         [ 0,  0,  1,  2,  3,  4]])
+        num_gen = []
+        cur_num_gen = 0
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        logging.info(
+            f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default"
+        )
+        consec_silence_count = 0
+        prev_token = None
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        # prepare the cache placeholder
+        # n_layers, 2, bsz, num_heads, src_len, head_dim
+        past = (
+            torch.ones(
+                [self.args.num_decoder_layers, 2, x.shape[0]],
+                device=x.device,
+                dtype=torch.float32,
+            )
+            if kvcache
+            else None
+        )
+        # handle multi-span kv-cache
+        new_masked_span = False
+
+        def sample_helper(
+            n_eog,
+            logits,
+            codebook_eog,
+            top_k,
+            top_p,
+            temperature,
+            prev_token,
+            consec_silence_count,
+            stop_repetition,
+            silence_tokens,
+            cur_num_gen,
+        ):
+            if n_eog == 0:
+                logits_adjust = logits
+                for jj in range(1, self.args.n_codebooks):
+                    logits_adjust[jj][self.args.eog] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                ##################### silence repetition handling #####################
+                if (
+                    stop_repetition > 0
+                    and prev_token in silence_tokens
+                    and consec_silence_count > stop_repetition
+                ):
+                    if logits_adjust[0, prev_token] < 0:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (
+                            consec_silence_count - (stop_repetition - 1)
+                        )
+                    else:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (
+                            consec_silence_count - (stop_repetition - 1)
+                        )
+                ##################### silence repetition handling #####################
+                if type(logits_adjust) == list:
+                    samples_list = []
+                    for logit in logits_adjust:
+                        # print(logit)
+                        # print(logit.shape)
+                        cur_sample = topk_sampling(
+                            logit.unsqueeze(0),
+                            top_k=top_k,
+                            top_p=top_p,
+                            temperature=temperature,
+                        )  # [1, 1]
+                        samples_list.append(cur_sample)
+                    samples = torch.cat(samples_list, dim=0)  # [K, 1]
+                else:
+                    samples = topk_sampling(
+                        logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                    )  # [K, 1]
+                assert samples.shape == torch.Size(
+                    (self.args.n_codebooks, 1)
+                ), f"samples.shape: {samples.shape}"
+                if cur_num_gen < self.args.n_codebooks - 1:
+                    for jj in range(1, self.args.n_codebooks - cur_num_gen):
+                        samples[-jj, 0] = self.args.empty_token
+
+                if (
+                    samples[0, 0] == self.args.eog
+                    or torch.argmax(logits[0], dim=-1) == self.args.eog
+                    or y_input.shape[1] > x_lens[0] * 10
+                ):  # last one means y is already too long, shouldn't happen, but put it here
+                    samples[0, 0] = self.args.eog
+                    codebook_eog[0] = True
+                ##################### silence repetition handling #####################
+                ##################### silence repetition handling #####################
+                if samples[0, 0] in silence_tokens and samples[0, 0] == prev_token:
+                    consec_silence_count += 1
+                else:
+                    consec_silence_count = 0
+                prev_token = samples[0, 0]
+                ##################### silence repetition handling #####################
+                ##################### silence repetition handling #####################
+                return samples, codebook_eog, prev_token, consec_silence_count
+            else:
+                assert (
+                    sum(codebook_eog[i] for i in range(n_eog)) == n_eog
+                ), f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
+                logits_adjust = logits
+                for jj in range(n_eog + 1, self.args.n_codebooks):
+                    logits_adjust[jj][self.args.eog] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                if type(logits_adjust) == list:
+                    samples_list = []
+                    for logit in logits_adjust:
+                        cur_sample = topk_sampling(
+                            logit.unsqueeze(0),
+                            top_k=top_k,
+                            top_p=top_p,
+                            temperature=temperature,
+                        )  # [1, 1]
+                        samples_list.append(cur_sample)
+                    samples = torch.cat(samples_list, dim=0)  # [K, 1]
+                else:
+                    samples = topk_sampling(
+                        logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                    )  # [K, 1]
+                for jj in range(n_eog):
+                    samples[jj, 0] = self.args.empty_token
+                samples[n_eog, 0] = self.args.eog
+                codebook_eog[n_eog] = True
+                return samples, codebook_eog, prev_token, consec_silence_count
+
+        while True:
+            y_out, present = self.dec_forward(
+                x_input,
+                x_lens,
+                x_attention_mask,
+                x_padding_mask,
+                y_input,
+                new_y_lens,
+                y_attention_mask,
+                y_padding_mask,
+                past=past,
+                last_3_tokens=new_masked_span,
+            )
+            if new_masked_span:
+                new_masked_span = False
+
+            if past != None:
+                past = (
+                    torch.cat([past, present.to(past.dtype)], dim=-2)
+                    if past.ndim > 3
+                    else present.to(past.dtype)
+                )
+
+            y_out = y_out[:, -1:]  # only take the last one
+
+            logits = torch.stack(
+                [self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)],
+                dim=1,
+            )  # [B K S card], B==S==1, so [1 K 1 card]
+            logits = logits.squeeze(0).squeeze(1)  # [K card]
+            assert logits.shape == torch.Size(
+                (self.args.n_codebooks, self.n_audio_tokens[0])
+            ), f"{logits.shape}"
+
+            n_eog = sum(codebook_eog)
+            assert n_eog < self.args.n_codebooks
+            if (
+                self.args.eos > 0
+            ):  # eos stands for end-of-sentence, which shouldn't be used as we are doing speech editing
+                for jj in range(self.args.n_codebooks):
+                    logits[jj][self.args.eos] = -10000.0
+            # need to use a helper function to hand different n_eog cases
+            samples, codebook_eog, prev_token, consec_silence_count = sample_helper(
+                n_eog,
+                logits,
+                codebook_eog,
+                top_k,
+                top_p,
+                temperature,
+                prev_token,
+                consec_silence_count,
+                stop_repetition,
+                silence_tokens,
+                cur_num_gen,
+            )
+            cur_num_gen += 1
+            cur_generated.append(samples.squeeze(-1))  # [K,1] -> [K]
+            # get samples_emb
+            samples_emb = torch.stack(
+                [
+                    self.audio_embedding[k](samples[k])
+                    for k in range(self.args.n_codebooks)
+                ],
+                dim=0,
+            )  # [K,1,D]
+            samples_emb = samples_emb.sum(dim=0, keepdim=True)  # [1,1,D]
+
+            if (
+                sum(codebook_eog) == self.args.n_codebooks
+            ):  # generation for the current span is done
+                # re-init
+                codebook_eog = [False] * self.args.n_codebooks
+                num_gen.append(cur_num_gen)
+                cur_num_gen = 0
+                generated.append(cur_generated)
+                cur_generated = []
+
+                # if the current mask span is the last span, then all done
+                # else
+                # append the next mask token and the four empty tokens to start the next generation
+                if len(more_mask_value) > 0:
+                    next_mask_ind = more_mask_value.pop(0)
+                    mask_emb = (
+                        self.mask_embedding[next_mask_ind].unsqueeze(0).unsqueeze(0)
+                    )  # [1,1,D]
+                    assert mask_emb.shape == torch.Size(
+                        (1, 1, self.args.d_model)
+                    ), mask_emb.shape
+                    empty_token = torch.LongTensor([self.args.empty_token]).to(y.device)
+                    empty_emb = torch.stack(
+                        [
+                            self.audio_embedding[k](empty_token)
+                            for k in range(self.args.n_codebooks)
+                        ],
+                        dim=0,
+                    ).sum(
+                        dim=0, keepdim=True
+                    )  # [1,1,D]
+                    assert empty_emb.shape == torch.Size(
+                        (1, 1, self.args.d_model)
+                    ), empty_emb.shape
+                    extra_emb = torch.cat([mask_emb, empty_emb], dim=1)  # [1,2,D]
+                    samples_emb = torch.cat(
+                        [samples_emb, extra_emb], dim=1
+                    )  # [1,3,D] # prev_last_token, mask_token, empty token
+                    assert samples_emb.shape == torch.Size(
+                        (1, 3, self.args.d_model)
+                    ), f"samples_emb.shape: {samples_emb.shape}"
+                    ##################### silence repetition handling #####################
+                    ##################### silence repetition handling #####################
+                    consec_silence_count = 0
+                    prev_token = None
+                    ##################### silence repetition handling #####################
+                    ##################### silence repetition handling #####################
+
+                    # handling kv-caching for multi-span editing
+                    new_masked_span = True
+                else:
+                    break
+            else:
+                assert samples_emb.shape == torch.Size(
+                    (1, 1, self.args.d_model)
+                ), f"samples_emb.shape: {samples_emb.shape}"
+
+            embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
+            # positional embedding
+            y_input = self.audio_positional_embedding(embedded_y)  # [B T D]
+            # make attention mask and padding mask
+            y_attention_mask = (
+                torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1)
+                .bool()
+                .to(y.device)
+            )
+            new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
+            y_padding_mask = torch.full((1, new_y_lens[0]), False).to(y.device)
+
+        assert (
+            len(generated) == num_mask
+        ), f"len(generated): {len(generated)}, num_mask: {num_mask}"
+
+        # # combine non_masked_span with generated spans
+        # first need to shift the generated part back
+        flatten_gen = []
+        for l, orig_span in enumerate(generated):
+            span = torch.stack(orig_span, dim=0)  # [T K]
+            span = span.transpose(1, 0)  # [K, T]
+            assert span.shape[0] == self.args.n_codebooks, span.shape
+            unshifted_span = []
+            for j, s in enumerate(span):
+                start_from = j
+                end_at = -(self.args.n_codebooks - start_from)
+                unshifted_span.append(s[start_from:end_at])
+            unshifted_span = torch.stack(unshifted_span, dim=0)
+
+            assert (
+                unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks
+            ), f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
+            flatten_gen.append(unshifted_span)
+        # logging.info(f"unshfited_span: {unshifted_span.shape}")
+        # raise
+        assert len(non_mask_intervals[0]) - 1 == len(
+            flatten_gen
+        ), f"len(non_mask_intervals[0]): {len(non_mask_intervals[0])}, len(flatten_gen): {len(flatten_gen)}"
+        res = []
+        for orig_interval, gen in zip(non_mask_intervals[0], flatten_gen):
+            res.append(y[0, :, orig_interval[0] : orig_interval[1]])
+            res.append(gen)
+        res.append(y[0, :, non_mask_intervals[0][-1][0] : non_mask_intervals[0][-1][1]])
+        res = torch.cat(res, dim=1).unsqueeze(0)  # [K,new_T] -> [1, K, new_T]
+
+        expected_y_len = (
+            y_len
+            - sum([item[1] - item[0] for item in mask_intervals[0]])
+            + sum([item - self.args.n_codebooks for item in num_gen])
+        )
+        assert res.shape == torch.Size(
+            (1, self.args.n_codebooks, expected_y_len)
+        ), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len - sum([item[1] - item[0] for item in mask_interval]) + sum([item - self.args.n_codebooks for item in num_gen]): {y_len}-{sum([item[1] - item[0] for item in mask_interval])} + {sum([item - self.args.n_codebooks for item in num_gen])}"
+
+        if self.args.special_first:
+            res = res - int(self.args.n_special)
+
+        return res
+
+    def inference_tts(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        top_k: int = -100,
+        top_p: float = 1.0,
+        temperature: float = 1.0,
+        stop_repetition: int = 3,
+        kvcache: int = 1,
+        silence_tokens: list[int] = [1388, 1898, 131],
+        *kargs,
+    ) -> torch.Tensor:
+        """
+        different from inference_tts, this implementation uses kvcache, which should have significant speed up
+        Args:
+          x:
+            A 2-D tensor of shape (1, L).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, K).
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          top_p: (`optional`) float
+            For Neucleus sampling
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+        """
+        eog_inference = self.args.eos if self.args.eos > 0 else self.args.eog
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        if self.args.special_first:
+            y = y + int(self.args.n_special)
+        y = y.transpose(2, 1)  # [1,T,K] -> [1,K,T]
+        assert (
+            y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks
+        ), y.shape  # there is no padding
+
+        # make x attention mask and x_input
+        x_attention_mask = (
+            torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1)
+            .bool()
+            .to(x.device)
+        )
+        # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+
+        y_len = y.shape[2]
+        y_lens = torch.LongTensor([y_len]).to(y.device)
+
+        # rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario
+        rearranged_y = [[y[0]]]
+        assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][
+            0
+        ].shape
+
+        # shift y to create the delayed pattern
+        shifted_y, patterns = self.shift(
+            rearranged_y
+        )  # each element [K S], patterns is not used, as we directly use the original input y
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
+        assert len(shifted_y[0]) == 1, len(shifted_y[0])
+
+        # below is different from forward or inference
+        # where we cut this shifted part
+        shifted_y[0][0] = shifted_y[0][0][:, : -(self.args.n_codebooks - 1)]
+        assert (
+            not (
+                shifted_y[0][0][self.args.n_codebooks :] == self.args.empty_token
+            ).any()
+            and not (shifted_y[0][0][self.args.n_codebooks :] == self.args.eog).any()
+        ), shifted_y[0][0]
+
+        # next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that
+        # next section is concate tensors of each sample to one tensor, which we also don't need
+        cated_y = shifted_y[0][0].unsqueeze(-1)  # [K,S]->[K,S,B]
+        new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device)
+        assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1))
+        assert not (cated_y == self.args.audio_pad_token).any(), cated_y
+
+        # replace tokens in y with the embeddings, add sum codebooks up
+        embedded_y = torch.stack(
+            [self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)],
+            dim=0,
+        )  # [K, S, B, D]
+        assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
+        assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
+        embedded_y = embedded_y.sum(dim=0)  # [K,S,B,D]->[S,B,D]
+        embedded_y = embedded_y.transpose(1, 0)  # [S,B,D]->[B,S,D]
+
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_attention_mask = (
+            torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1)
+            .bool()
+            .to(y.device)
+        )
+
+        x_padding_mask = torch.full((1, x_lens[0]), False).to(x.device)
+        y_padding_mask = torch.full((1, new_y_lens[0]), False).to(y.device)
+
+        # entering the generation stage
+        # starting from line 708
+        codebook_eog = [False] * self.args.n_codebooks
+        generated = []  # doesn't contain any empty token, contain eog
+        cur_generated = []
+        # say 0 is empty, 4 is eog
+        # tensor([[ 1,  2,  3,  4,  0,  0],
+        #         [ 0,  1,  2,  3,  4,  0],
+        #         [ 0,  0,  1,  2,  3,  4]])
+        num_gen = []
+        cur_num_gen = 0
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        logging.info(
+            f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default"
+        )
+        consec_silence_count = 0
+        prev_token = None
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+
+        # prepare the cache placeholder
+        # n_layers, 2, bsz, num_heads, src_len, head_dim
+        past = (
+            torch.ones(
+                [self.args.num_decoder_layers, 2, x.shape[0]],
+                device=x.device,
+                dtype=torch.float32,
+            )
+            if kvcache
+            else None
+        )
+
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        def sample_helper(
+            n_eog,
+            logits,
+            codebook_eog,
+            top_k,
+            top_p,
+            temperature,
+            prev_token,
+            consec_silence_count,
+            stop_repetition,
+            silence_tokens,
+            cur_num_gen,
+        ):
+            if n_eog == 0:
+                logits_adjust = logits
+                for jj in range(1, self.args.n_codebooks):
+                    logits_adjust[jj][eog_inference] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                if (
+                    cur_num_gen <= self.args.encodec_sr // 5
+                ):  # this shouldn't happen, but just in case the model stopped too early
+                    logits_adjust[0][eog_inference] = -10000
+                ##################### silence repetition handling #####################
+                if (
+                    stop_repetition > 0
+                    and prev_token in silence_tokens
+                    and consec_silence_count > stop_repetition
+                ):
+                    if logits_adjust[0, prev_token] < 0:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (
+                            consec_silence_count - (stop_repetition - 1)
+                        )
+                    else:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (
+                            consec_silence_count - (stop_repetition - 1)
+                        )
+                ##################### silence repetition handling #####################
+                samples = topk_sampling(
+                    logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                )  # [K, 1]
+                assert samples.shape == torch.Size(
+                    (self.args.n_codebooks, 1)
+                ), f"samples.shape: {samples.shape}"
+                if cur_num_gen < self.args.n_codebooks - 1:
+                    for jj in range(1, self.args.n_codebooks - cur_num_gen):
+                        samples[-jj, 0] = self.args.empty_token
+
+                if (
+                    samples[0, 0] == eog_inference
+                    or torch.argmax(logits[0], dim=-1) == eog_inference
+                    or y_input.shape[1] > x_lens[0] * (self.args.encodec_sr // 5)
+                ):  # last one means y is already too long, shouldn't happen, but put it here
+                    samples[0, 0] = eog_inference
+                    codebook_eog[0] = True
+                ##################### silence repetition handling #####################
+                if samples[0, 0] in silence_tokens and samples[0, 0] == prev_token:
+                    consec_silence_count += 1
+                else:
+                    consec_silence_count = 0
+                prev_token = samples[0, 0]
+                ##################### silence repetition handling #####################
+                return samples, codebook_eog, prev_token, consec_silence_count
+            else:
+                assert (
+                    sum(codebook_eog[i] for i in range(n_eog)) == n_eog
+                ), f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
+                logits_adjust = logits
+                for jj in range(n_eog + 1, self.args.n_codebooks):
+                    logits_adjust[jj][eog_inference] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                samples = topk_sampling(
+                    logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                )  # [K, 1]
+                for jj in range(n_eog):
+                    samples[jj, 0] = self.args.empty_token
+                samples[n_eog, 0] = eog_inference
+                codebook_eog[n_eog] = True
+                return samples, codebook_eog, prev_token, consec_silence_count
+
+        while True:
+            y_out, present = self.dec_forward(
+                x_input,
+                x_lens,
+                x_attention_mask,
+                x_padding_mask,
+                y_input,
+                new_y_lens,
+                y_attention_mask,
+                y_padding_mask,
+                past=past,
+            )
+            if past != None:
+                past = (
+                    torch.cat([past, present.to(past.dtype)], dim=-2)
+                    if past.ndim > 3
+                    else present.to(past.dtype)
+                )
+
+            y_out = y_out[:, -1:]  # only take the last token
+            logits = torch.stack(
+                [self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)],
+                dim=1,
+            )  # [B K S card], B==S==1, so [1 K 1 card]
+            logits = logits.squeeze(0).squeeze(1)  # [K card]
+            assert logits.shape == torch.Size(
+                (self.args.n_codebooks, self.n_audio_tokens[0])
+            ), f"{logits.shape}"
+
+            n_eog = sum(codebook_eog)
+            assert n_eog < self.args.n_codebooks
+            if (
+                self.args.eos > 0
+            ):  # if we are using end-of-sentence token (which is used by default), eog shouldn't be used here, as there is no masked spans
+                for jj in range(self.args.n_codebooks):
+                    logits[jj][self.args.eog] = -10000.0
+
+            samples, codebook_eog, prev_token, consec_silence_count = sample_helper(
+                n_eog,
+                logits,
+                codebook_eog,
+                top_k,
+                top_p,
+                temperature,
+                prev_token,
+                consec_silence_count,
+                stop_repetition,
+                silence_tokens,
+                cur_num_gen,
+            )
+
+            cur_num_gen += 1
+            cur_generated.append(samples.squeeze(-1))  # [K,1] -> [K]
+
+            # samples.shape is [K,1]
+            # ge samples_emb
+            samples_emb = torch.stack(
+                [
+                    self.audio_embedding[k](samples[k])
+                    for k in range(self.args.n_codebooks)
+                ],
+                dim=0,
+            )  # [K,1,D]
+            samples_emb = samples_emb.sum(dim=0, keepdim=True)  # [1,1,D]
+
+            if (
+                sum(codebook_eog) == self.args.n_codebooks
+            ):  # generation for the current span is done
+                codebook_eog = [False] * self.args.n_codebooks
+                num_gen.append(cur_num_gen)
+                cur_num_gen = 0
+                generated.append(cur_generated)
+                cur_generated = []
+                break
+            else:
+                assert samples_emb.shape == torch.Size(
+                    (1, 1, self.args.d_model)
+                ), f"samples_emb.shape: {samples_emb.shape}"
+
+            embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
+            y_input = self.audio_positional_embedding(embedded_y)  # [B T D]
+            # make attention mask and padding mask
+            y_attention_mask = (
+                torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1)
+                .bool()
+                .to(y.device)
+            )
+            new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
+            y_padding_mask = torch.full((1, new_y_lens[0]), False).to(y.device)
+
+        assert len(generated) == 1, f"len(generated): {len(generated)}"
+
+        # revert the pattern
+        flatten_gen = []
+        for l, orig_span in enumerate(generated):
+            span = torch.stack(orig_span, dim=0)  # [T, K]
+            span = span.transpose(1, 0)  # [K, T]
+            assert span.shape[0] == self.args.n_codebooks, span.shape
+            unshifted_span = []
+            for j, s in enumerate(span):
+                start_from = j
+                end_at = -(self.args.n_codebooks - start_from)
+                unshifted_span.append(s[start_from:end_at])
+            unshifted_span = torch.stack(unshifted_span, dim=0)
+
+            assert (
+                unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks
+            ), f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
+
+            flatten_gen.append(unshifted_span)
+        assert len(flatten_gen) == 1, len(flatten_gen)
+
+        # combine
+        res = [y[0], flatten_gen[0]]
+        res = torch.cat(res, dim=1).unsqueeze(0)  # [K, new_t] -> [1, K, new_T]
+
+        expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen])
+        assert res.shape == torch.Size(
+            (1, self.args.n_codebooks, expected_y_len)
+        ), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}"
+
+        if self.args.special_first:
+            res = res - int(self.args.n_special)
+            flatten_gen = flatten_gen - int(self.args.n_special)
+
+        return res, flatten_gen[0].unsqueeze(0)
+
+    def inference_tts_batch(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        top_k: int = -100,
+        top_p: float = 1.0,
+        temperature: float = 1.0,
+        stop_repetition: int = 3,
+        kvcache: int = 1,
+        batch_size: int = 5,
+        silence_tokens: list[int] = [1388, 1898, 131],
+        *kargs,
+    ) -> torch.Tensor:
+        """
+        have a batch size when forward passing, but they are equivalant to same example but different random seed, therefore as long as one example generated eog, we can drop all other samlpes
+        different from inference_tts, this implementation uses kvcache, which should have significant speed up
+        Args:
+          x:
+            A 2-D tensor of shape (1, L).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, K).
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          top_p: (`optional`) float
+            For Neucleus sampling
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+        """
+        eog_inference = self.args.eos if self.args.eos > 0 else self.args.eog
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        if self.args.special_first:
+            y = y + int(self.args.n_special)
+        y = y.transpose(2, 1)  # [1,T,K] -> [1,K,T]
+        assert (
+            y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks
+        ), y.shape  # there is no padding
+
+        # make x attention mask and x_input
+        x_attention_mask = (
+            torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1)
+            .bool()
+            .to(x.device)
+        )
+        # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+
+        y_len = y.shape[2]
+        y_lens = torch.LongTensor([y_len]).to(y.device)
+
+        # rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario
+        rearranged_y = [[y[0]]]
+        assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][
+            0
+        ].shape
+
+        # shift y to create the delayed pattern
+        shifted_y, patterns = self.shift(
+            rearranged_y
+        )  # each element [K S], patterns is not used, as we directly use the original input y
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
+        assert len(shifted_y[0]) == 1, len(shifted_y[0])
+
+        # below is different from forward or inference
+        # where we cut this shifted part
+        shifted_y[0][0] = shifted_y[0][0][:, : -(self.args.n_codebooks - 1)]
+        assert (
+            not (
+                shifted_y[0][0][self.args.n_codebooks :] == self.args.empty_token
+            ).any()
+            and not (shifted_y[0][0][self.args.n_codebooks :] == self.args.eog).any()
+        ), shifted_y[0][0]
+
+        # next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that
+        # next section is concate tensors of each sample to one tensor, which we also don't need
+        cated_y = shifted_y[0][0].unsqueeze(-1)  # [K,S]->[K,S,B]
+        new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device)
+        assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1))
+        assert not (cated_y == self.args.audio_pad_token).any(), cated_y
+
+        # replace tokens in y with the embeddings, add sum codebooks up
+        embedded_y = torch.stack(
+            [self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)],
+            dim=0,
+        )  # [K, S, B, D]
+        assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
+        assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
+        embedded_y = embedded_y.sum(dim=0)  # [K,S,B,D]->[S,B,D]
+        embedded_y = embedded_y.transpose(1, 0)  # [S,B,D]->[B,S,D]
+
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_attention_mask = (
+            torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1)
+            .bool()
+            .to(y.device)
+        )
+
+        x_padding_mask = torch.full((1, x_lens[0]), False).to(x.device)
+        y_padding_mask = torch.full((1, new_y_lens[0]), False).to(y.device)
+
+        # entering the generation stage
+        # starting from line 708
+        codebook_eog = [False] * self.args.n_codebooks
+        generated = []  # doesn't contain any empty token, contain eog
+        cur_generated = [[] for _ in range(batch_size)]
+        # say 0 is empty, 4 is eog
+        # tensor([[ 1,  2,  3,  4,  0,  0],
+        #         [ 0,  1,  2,  3,  4,  0],
+        #         [ 0,  0,  1,  2,  3,  4]])
+        num_gen = []
+        cur_num_gen = 0
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        logging.info(
+            f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default"
+        )
+        consec_silence_counts = [0 for _ in range(batch_size)]
+        prev_tokens = [None for _ in range(batch_size)]
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+
+        # prepare the cache placeholder
+        # n_layers, 2, bsz, num_heads, src_len, head_dim
+        past = (
+            torch.ones(
+                [self.args.num_decoder_layers, 2, x.shape[0]],
+                device=x.device,
+                dtype=torch.float32,
+            )
+            if kvcache
+            else None
+        )
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        keep = None  # NOTE: this very important, tells which sample to keep
+
+        def sample_helper(
+            n_eog,
+            logits,
+            codebook_eog,
+            top_k,
+            top_p,
+            temperature,
+            prev_tokens,
+            consec_silence_counts,
+            stop_repetition,
+            silence_tokens,
+            cur_num_gen,
+            keep,
+        ):
+            if n_eog == 0:
+                logits_adjust = logits
+                for jj in range(1, self.args.n_codebooks):
+                    logits_adjust[:, jj, eog_inference] = -10000
+                    logits_adjust[:, jj, self.args.empty_token] = -10000
+                if (
+                    cur_num_gen <= self.args.encodec_sr // 5
+                ):  # this shouldn't happen, but just in case the model stopped too early
+                    logits_adjust[:, :, eog_inference] = -10000
+                ##################### silence repetition handling #####################
+                for b in range(batch_size):
+                    prev_token = prev_tokens[b]
+                    consec_silence_count = consec_silence_counts[b]
+                    if (
+                        stop_repetition > 0
+                        and prev_token in silence_tokens
+                        and consec_silence_count > stop_repetition
+                    ):
+                        if logits_adjust[b, 0, prev_token] < 0:
+                            logits_adjust[b, 0, prev_token] = logits_adjust[
+                                b, 0, prev_token
+                            ] * (consec_silence_count - (stop_repetition - 1))
+                        else:
+                            logits_adjust[b, 0, prev_token] = logits_adjust[
+                                b, 0, prev_token
+                            ] / (consec_silence_count - (stop_repetition - 1))
+                ##################### silence repetition handling #####################
+                samples = topk_sampling(
+                    logits_adjust.reshape(
+                        batch_size * self.args.n_codebooks, logits_adjust.shape[-1]
+                    ),
+                    top_k=top_k,
+                    top_p=top_p,
+                    temperature=temperature,
+                )  # [B*K, 1]
+                samples = samples.reshape(batch_size, self.args.n_codebooks, 1)
+                assert samples.shape == torch.Size(
+                    (batch_size, self.args.n_codebooks, 1)
+                ), f"samples.shape: {samples.shape}"
+                for b in range(batch_size):
+                    if cur_num_gen < self.args.n_codebooks - 1:
+                        for jj in range(1, self.args.n_codebooks - cur_num_gen):
+                            samples[b, -jj, 0] = self.args.empty_token
+
+                    if (
+                        samples[b, 0, 0] == eog_inference
+                        or torch.argmax(logits[b, 0], dim=-1) == eog_inference
+                        or y_input.shape[1] > x_lens[b] * (self.args.encodec_sr // 5)
+                    ):  # last one means y is already too long, shouldn't happen, but put it here
+                        samples[b, 0, 0] = eog_inference
+                        codebook_eog[0] = True
+                        keep = b  # NOTE keep is a very important variable, we only return this one, note that if eog shows up in two samples, keep will be overwritten by the later one (or the last one)
+                    ##################### silence repetition handling #####################
+                    if (
+                        samples[b, 0, 0] in silence_tokens
+                        and samples[b, 0, 0] == prev_tokens[b]
+                    ):
+                        consec_silence_counts[b] += 1
+                    else:
+                        consec_silence_counts[b] = 0
+                    prev_tokens[b] = samples[b, 0, 0]
+                ##################### silence repetition handling #####################
+                return samples, codebook_eog, prev_tokens, consec_silence_counts, keep
+            else:
+                assert (
+                    sum(codebook_eog[i] for i in range(n_eog)) == n_eog
+                ), f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
+                logits_adjust = logits
+                for jj in range(n_eog + 1, self.args.n_codebooks):
+                    logits_adjust[:, jj, eog_inference] = -10000
+                    logits_adjust[:, jj, self.args.empty_token] = -10000
+                samples = topk_sampling(
+                    logits_adjust.reshape(
+                        batch_size * self.args.n_codebooks, logits_adjust.shape[-1]
+                    ),
+                    top_k=top_k,
+                    top_p=top_p,
+                    temperature=temperature,
+                )  # [B, K, 1]
+                samples = samples.reshape(batch_size, self.args.n_codebooks, 1)
+                for jj in range(n_eog):
+                    samples[keep, jj, 0] = self.args.empty_token
+                samples[keep, n_eog, 0] = eog_inference
+                codebook_eog[n_eog] = True
+                return samples, codebook_eog, prev_tokens, consec_silence_counts, keep
+
+        while True:
+            # if cur_num_gen > 0, should have everything in kvcache, so only pass in the last token
+            # in the first generation step, we repeat each tensor to make their first dimension of length the batch size
+            if cur_num_gen == 0:
+                assert x_input.ndim == 3 and x_input.shape[0] == 1, x_input.shape
+                assert (
+                    x_padding_mask.ndim == 2 and x_padding_mask.shape[0] == 1
+                ), x_padding_mask.shape
+                assert (
+                    y_input.ndim == 3
+                    and y_input.shape[0] == 1
+                    and y_input.shape[1] == new_y_lens[0]
+                ), y_input.shape
+                assert (
+                    embedded_y.ndim == 3
+                    and embedded_y.shape[0] == 1
+                    and embedded_y.shape[1] == new_y_lens[0]
+                ), embedded_y.shape
+                x_input = x_input.repeat(batch_size, 1, 1)
+                x_lens = x_lens.repeat(batch_size)
+                # x_attention_mask = x_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension
+                x_padding_mask = x_padding_mask.repeat(batch_size, 1)
+                y_input = y_input.repeat(batch_size, 1, 1)
+                new_y_lens = new_y_lens.repeat(batch_size)
+                # y_attention_mask = y_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension
+                y_padding_mask = y_padding_mask.repeat(batch_size, 1)
+                embedded_y = embedded_y.repeat(
+                    batch_size, 1, 1
+                )  # will be used to concat with newly generated token embedding
+                past = past.repeat(1, 1, batch_size) if past != None else None
+            else:
+                assert (
+                    x_input.shape[0] == batch_size
+                    and x_padding_mask.shape[0] == batch_size
+                    and y_input.shape[0] == batch_size
+                    and new_y_lens.shape[0] == batch_size
+                ), f"x_input.shape: {x_input.shape}, x_padding_mask.shape: {x_padding_mask.shape}, y_input.shape: {y_input.shape}, new_y_lens.shape: {new_y_lens.shape}"
+            y_out, present = self.dec_forward(
+                x_input,
+                x_lens,
+                x_attention_mask,
+                x_padding_mask,
+                y_input,
+                new_y_lens,
+                y_attention_mask,
+                y_padding_mask,
+                past=past,
+            )
+            if past != None:
+                past = (
+                    torch.cat([past, present.to(past.dtype)], dim=-2)
+                    if past.ndim > 3
+                    else present.to(past.dtype)
+                )
+
+            # if no eog emerges, y_out should have batch size of batch_size
+            if sum(codebook_eog) == 0:
+                assert y_out.shape[0] == batch_size and y_out.ndim == 3, y_out.shape
+            y_out = y_out[:, -1:]  # only take the last token
+            logits = torch.stack(
+                [self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)],
+                dim=1,
+            )  # [B K S card], S==1, so [B K 1 card]
+            logits = logits.squeeze(2)  # [B K card]
+            assert logits.shape == torch.Size(
+                (batch_size, self.args.n_codebooks, self.n_audio_tokens[0])
+            ), f"{logits.shape}"
+
+            n_eog = sum(codebook_eog)
+            if self.args.eos > 0:
+                for jj in range(self.args.n_codebooks):
+                    logits[:, jj, self.args.eog] = -10000.0
+            samples, codebook_eog, prev_tokens, consec_silence_counts, keep = (
+                sample_helper(
+                    n_eog,
+                    logits,
+                    codebook_eog,
+                    top_k,
+                    top_p,
+                    temperature,
+                    prev_tokens,
+                    consec_silence_counts,
+                    stop_repetition,
+                    silence_tokens,
+                    cur_num_gen,
+                    keep,
+                )
+            )
+
+            cur_num_gen += 1
+            if sum(codebook_eog) == 0:  # no eog yet, keep batch_size of samples
+                assert keep == None
+                for b in range(batch_size):
+                    cur_generated[b].append(samples[b].squeeze(-1))
+            elif sum(codebook_eog) == 1:  # the first eog just showed up in this step
+                assert keep != None
+                cur_generated = cur_generated[keep]
+                cur_generated.append(samples[keep].squeeze(-1))
+            else:  # we are generating the rest eogs for the 'keep' sample
+                cur_generated.append(samples[keep].squeeze(-1))
+
+            # samples.shape is [K,1]
+            # ge samples_emb
+            samples_emb = torch.stack(
+                [
+                    self.audio_embedding[k](samples[:, k])
+                    for k in range(self.args.n_codebooks)
+                ],
+                dim=1,
+            )  # [B, K,1,D]
+            assert samples_emb.shape == torch.Size(
+                [batch_size, self.args.n_codebooks, 1, self.args.d_model]
+            )
+            samples_emb = samples_emb.sum(dim=1, keepdim=False)  # [B,1,D]
+            if (
+                sum(codebook_eog) == self.args.n_codebooks
+            ):  # generation for the current span is done
+                codebook_eog = [False] * self.args.n_codebooks
+                num_gen.append(cur_num_gen)
+                cur_num_gen = 0
+                generated.append(cur_generated)
+                cur_generated = [[] for _ in range(batch_size)]
+                break
+            else:
+                assert samples_emb.shape == torch.Size(
+                    (batch_size, 1, self.args.d_model)
+                ), f"samples_emb.shape: {samples_emb.shape}"
+
+            embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
+            y_input = self.audio_positional_embedding(embedded_y)  # [B T D]
+            # make attention mask and padding mask
+            y_attention_mask = (
+                torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1)
+                .bool()
+                .to(y.device)
+            )
+            new_y_lens = (
+                torch.LongTensor([y_input.shape[1]]).to(y.device).repeat(batch_size)
+            )
+            y_padding_mask = torch.full((batch_size, new_y_lens[0]), False).to(y.device)
+
+        assert len(generated) == 1, f"len(generated): {len(generated)}"
+
+        # revert the pattern
+        flatten_gen = []
+        for l, orig_span in enumerate(generated):
+            span = torch.stack(orig_span, dim=0)  # [T, K]
+            span = span.transpose(1, 0)  # [K, T]
+            assert span.shape[0] == self.args.n_codebooks, span.shape
+            unshifted_span = []
+            for j, s in enumerate(span):
+                start_from = j
+                end_at = -(self.args.n_codebooks - start_from)
+                unshifted_span.append(s[start_from:end_at])
+            unshifted_span = torch.stack(unshifted_span, dim=0)
+
+            assert (
+                unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks
+            ), f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
+
+            flatten_gen.append(unshifted_span)
+        assert len(flatten_gen) == 1, len(flatten_gen)
+
+        # combine
+        res = [y[0], flatten_gen[0]]
+        res = torch.cat(res, dim=1).unsqueeze(0)  # [K, new_t] -> [1, K, new_T]
+
+        expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen])
+        assert res.shape == torch.Size(
+            (1, self.args.n_codebooks, expected_y_len)
+        ), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}"
+
+        if self.args.special_first:
+            res = res - int(self.args.n_special)
+            flatten_gen = flatten_gen - int(self.args.n_special)
+
+        return res, flatten_gen[0].unsqueeze(0)

src/model/modules/voicecraftconfig.py

@@ -0,0 +1,37 @@
+class VoiceCraftConfig:
+
+    def __init__(
+        self,
+        model_name="330M_TTSEnhanced.pth",  # "gigaHalfLibri330M_TTSEnhanced_max16s.pth",
+        encodec="encodec_4cb2048_giga.th",
+        top_k=0,
+        top_p=0.9,
+        temperature=1,
+        kvcache=1,
+        codec_sr=50,
+        codec_audio_sr=16000,
+        silence_tokens=[1388, 1898, 131],
+        stop_repetition=3,
+        sample_batch_size=2,
+        seed=1,
+        cut_off_sec=7.87,
+        voice_audio_path="84_121550_000074_000000.wav",
+        voice_audio_transcript="But when I had approached so near to them The common object, which the sense deceives, Lost not by distance any of its marks",
+        **kwargs,
+    ):
+        super().__init__()
+        self.model_name = model_name
+        self.encodec = encodec
+        self.top_k = top_k
+        self.top_p = top_p
+        self.temperature = temperature
+        self.kvcache = kvcache
+        self.codec_sr = codec_sr
+        self.codec_audio_sr = codec_audio_sr
+        self.silence_tokens = silence_tokens
+        self.stop_repetition = stop_repetition
+        self.sample_batch_size = sample_batch_size
+        self.seed = seed
+        self.cut_off_sec = cut_off_sec
+        self.voice_audio_path = voice_audio_path
+        self.voice_audio_transcript = voice_audio_transcript

src/utils/image_utils.py

@@ -0,0 +1,100 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import logging
+import os
+from io import BytesIO
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Tuple, Union
+
+import PIL
+import numpy as np
+import requests
+from packaging import version
+
+
+def _is_numpy(x):
+    return isinstance(x, np.ndarray)
+
+
+def is_numpy_array(x):
+    """
+    Tests if `x` is a numpy array or not.
+    """
+    return _is_numpy(x)
+
+
+def is_pil_image(img):
+    return isinstance(img, PIL.Image.Image)
+
+
+def is_valid_image(img):
+    return is_pil_image(img) or is_numpy_array(img)
+
+
+def valid_images(imgs):
+    # If we have an list of images, make sure every image is valid
+    if isinstance(imgs, (list, tuple)):
+        for img in imgs:
+            if not valid_images(img):
+                return False
+    # If not a list of tuple, we have been given a single image or batched tensor of images
+    elif not is_valid_image(imgs):
+        return False
+    return True
+
+
+def is_batched(img):
+    if isinstance(img, (list, tuple)):
+        return is_valid_image(img[0])
+    return False
+
+
+def is_scaled_image(image: np.ndarray) -> bool:
+    """
+    Checks to see whether the pixel values have already been rescaled to [0, 1].
+    """
+    if image.dtype == np.uint8:
+        return False
+
+    # It's possible the image has pixel values in [0, 255] but is of floating type
+    return np.min(image) >= 0 and np.max(image) <= 1
+
+
+def make_batched_images(images):
+    """
+    Accepts images in list or nested list format, and makes a list of images for preprocessing.
+
+    Args:
+        images (`Union[List[List[ImageInput]], List[ImageInput], ImageInput]`):
+            The input image.
+
+    Returns:
+        list: A list of images.
+    """
+    if (
+        isinstance(images, (list, tuple))
+        and isinstance(images[0], (list, tuple))
+        and is_valid_image(images[0][0])
+    ):
+        return [img for img_list in images for img in img_list]
+
+    elif isinstance(images, (list, tuple)) and is_valid_image(images[0]):
+        return images
+
+    elif is_valid_image(images):
+        return [images]
+
+    raise ValueError(f"Could not make batched video from {images}")

src/utils/model_utils.py

@@ -0,0 +1,73 @@
+import json
+import os
+from typing import Optional
+from PIL import Image
+
+
+from src.model.modules.imagecraftconfig import ImageCraftConfig
+from src.model.modules.imagecraftprocessor import (
+    ImageCraftProcessor,
+)
+
+
+def move_inputs_to_device(model_inputs: dict, device: str):
+    model_inputs = {k: v.to(device) for k, v in model_inputs.items()}
+    return model_inputs
+
+
+def get_model_inputs(
+    processor: ImageCraftProcessor,
+    prompt: str,
+    image: Image,
+    suffix: Optional[str] = None,
+    device: str = "cuda",
+):
+    images = [image]
+    prompts = [prompt]
+    if suffix is not None:
+        suffix = [suffix]
+    model_inputs = processor(text=prompts, images=images)
+    model_inputs = move_inputs_to_device(model_inputs, device)
+    return model_inputs
+
+
+def get_config(config_file="config.json"):
+    config = None
+    with open(config_file, "r") as f:
+        model_config_file = json.load(f)
+        config = ImageCraftConfig(**model_config_file)
+
+    return config
+
+
+# def load_hf_model(model_path: str, device: str) -> Tuple[ImageCraft, AutoTokenizer]:
+
+#     # Load the tokenizer
+#     tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="right")
+#     assert tokenizer.padding_side == "right"
+
+#     # Find all the *.safetensors files
+#     safetensors_files = glob.glob(os.path.join(model_path, "*.safetensors"))
+
+#     # ... and load them one by one in the tensors dictionary
+#     tensors = {}
+#     for safetensors_file in safetensors_files:
+#         with safe_open(safetensors_file, framework="pt", device="cpu") as f:
+#             for key in f.keys():
+#                 tensors[key] = f.get_tensor(key)
+
+#     # Load the model's config
+#     with open(os.path.join(model_path, "config.json"), "r") as f:
+#         model_config_file = json.load(f)
+#         config = ImageCraftConfig(**model_config_file)
+
+#     # Create the model using the configuration
+#     model = ImageCraft(config).to(device)
+
+#     # Load the state dict of the model
+#     model.load_state_dict(tensors, strict=False)
+
+#     # Tie weights
+#     model.tie_weights()
+
+#     return (model, tokenizer)

src/utils/tools.py

@@ -0,0 +1,22 @@
+import yaml
+import pickle
+
+
+def load_config():
+    # Read in the configuration file
+    with open("config.yaml") as p:
+        config = yaml.safe_load(p)
+    return config
+
+
+def pickle_dump(path, variable):
+    # Serialize data from memory to file
+    with open(path, "wb") as handle:
+        pickle.dump(variable, handle)
+
+
+def pickle_load(path):
+    # Read and load serialized data from file
+    with open(path, "rb") as handle:
+        loaded = pickle.load(handle)
+    return loaded

src/utils/util.py

@@ -0,0 +1,305 @@
+import io
+import os
+from pathlib import Path
+from tempfile import TemporaryDirectory
+import torch
+import torchaudio
+import random
+import numpy as np
+from PIL import Image
+from urllib.parse import urlparse
+from os.path import exists
+import re
+from num2words import num2words
+import uuid
+
+from typing import List, Optional, Dict, Union, Tuple, Iterable
+
+from src.utils.image_utils import is_valid_image
+
+
+IMAGENET_STANDARD_MEAN = [0.5, 0.5, 0.5]
+IMAGENET_STANDARD_STD = [0.5, 0.5, 0.5]
+
+
+def is_local(url):
+    url_parsed = urlparse(url)
+    if url_parsed.scheme in ("file", ""):
+        return exists(url_parsed.path)
+    return False
+
+
+def replace_numbers_with_words(sentence):
+    sentence = re.sub(r"(\d+)", r" \1 ", sentence)
+
+    def replace_with_words(match):
+        num = match.group(0)
+        try:
+            return num2words(num)
+        except:
+            return num
+
+    return re.sub(r"\b\d+\b", replace_with_words, sentence)
+
+
+def save_to_buffer(audio_tensors, codec_audio_sr):
+
+    result = torch.cat(audio_tensors, 1)
+    buffer = io.BytesIO()
+    torchaudio.save(buffer, result, int(codec_audio_sr), format="wav")
+    buffer.seek(0)
+    return buffer.read()
+
+
+def save_to_file(audio_tensors, codec_audio_sr):
+    generated_audio_dir = f"media/voicecraft/generated"
+    Path(generated_audio_dir).mkdir(parents=True, exist_ok=True)
+    filename = f"{generated_audio_dir}/{str(uuid.uuid4())}.wav"
+    tensors = torch.cat(audio_tensors, 1)
+    torchaudio.save(filename, tensors, int(codec_audio_sr), format="wav")
+    return filename
+
+
+def split_line_to_sentences(line):
+    line = line.strip().capitalize()
+    line = line + "." if line and line[-1] not in (".", "!", "?") else line
+    sentences = re.findall(r"\w+.*?[.?!]", line.replace("\n", " "), flags=re.S)
+    return sentences
+
+
+def seed_everything(seed=1):
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+def add_image_tokens_to_prompt(prefix_prompt, bos_token, image_seq_length, image_token):
+    return f"{image_token * image_seq_length}{bos_token}{prefix_prompt}\n"
+
+
+def rescale(
+    image: np.ndarray, scale: float, dtype: np.dtype = np.float32
+) -> np.ndarray:
+    rescaled_image = image * scale
+    rescaled_image = rescaled_image.astype(dtype)
+    return rescaled_image
+
+
+def resize(
+    image: Image,
+    size: Tuple[int, int],
+    resample: Image.Resampling = None,
+    reducing_gap: Optional[int] = None,
+) -> np.ndarray:
+    height, width = size
+    resized_image = image.resize(
+        (width, height), resample=resample, reducing_gap=reducing_gap
+    )
+    return resized_image
+
+
+def normalize(
+    image: np.ndarray,
+    mean: Union[float, Iterable[float]],
+    std: Union[float, Iterable[float]],
+) -> np.ndarray:
+    mean = np.array(mean, dtype=image.dtype)
+    std = np.array(std, dtype=image.dtype)
+    image = (image - mean) / std
+    return image
+
+
+def process_images(
+    images: List[Image.Image],
+    size: Dict[str, int] = None,
+    resample: Image.Resampling = None,
+    rescale_factor: float = None,
+    image_mean: Optional[Union[float, List[float]]] = None,
+    image_std: Optional[Union[float, List[float]]] = None,
+) -> List[np.ndarray]:
+    height, width = size[0], size[1]
+    images = [
+        resize(image=image, size=(height, width), resample=resample) for image in images
+    ]
+    # Convert each image to a numpy array
+    images = [np.array(image) for image in images]
+    # Rescale the pixel values to be in the range [0, 1]
+    images = [rescale(image, scale=rescale_factor) for image in images]
+    # Normalize the images to have mean 0 and standard deviation 1
+    images = [normalize(image, mean=image_mean, std=image_std) for image in images]
+    # Move the channel dimension to the first dimension. The model expects images in the format [Channel, Height, Width]
+    images = [image.transpose(2, 0, 1) for image in images]
+    return images
+
+
+def sample_top_p(probs: torch.Tensor, p: float):
+    # (B, vocab_size)
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    # (B, vocab_size)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    # (B, vocab_size)
+    # (Substracting "probs_sort" shifts the cumulative sum by 1 position to the right before masking)
+    mask = probs_sum - probs_sort > p
+    # Zero out all the probabilities of tokens that are not selected by the Top P
+    probs_sort[mask] = 0.0
+    # Redistribute the probabilities so that they sum up to 1.
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    # Sample a token (its index) from the top p distribution
+    next_token = torch.multinomial(probs_sort, num_samples=1)
+    # Get the token position in the vocabulary corresponding to the sampled index
+    next_token = torch.gather(probs_idx, -1, next_token)
+    return next_token
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """
+    Args:
+      lengths:
+        A 1-D tensor containing sentence lengths.
+      max_len:
+        The length of masks.
+    Returns:
+      Return a 2-D bool tensor, where masked positions
+      are filled with `True` and non-masked positions are
+      filled with `False`.
+    >>> lengths = torch.tensor([1, 3, 2, 5])
+    >>> make_pad_mask(lengths)
+    tensor([[False,  True,  True,  True,  True],
+            [False, False, False,  True,  True],
+            [False, False,  True,  True,  True],
+            [False, False, False, False, False]])
+    """
+    assert lengths.ndim == 1, lengths.ndim
+    max_len = max(max_len, lengths.max())
+    n = lengths.size(0)
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
+
+    return expaned_lengths >= lengths.unsqueeze(-1)
+
+
+def _prepare_4d_causal_attention_mask_with_cache_position(
+    attention_mask: torch.Tensor,
+    sequence_length: int,
+    target_length: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    min_dtype: float,
+    cache_position: torch.Tensor,
+    batch_size: int,
+    is_training: bool = False,
+    token_type_ids: torch.Tensor = None,
+):
+    """
+    Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+    `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+    Args:
+        attention_mask (`torch.Tensor`):
+            A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
+        sequence_length (`int`):
+            The sequence length being processed.
+        target_length (`int`):
+            The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
+        dtype (`torch.dtype`):
+            The dtype to use for the 4D attention mask.
+        device (`torch.device`):
+            The device to plcae the 4D attention mask on.
+        min_dtype (`float`):
+            The minimum value representable with the dtype `dtype`.
+        cache_position (`torch.Tensor`):
+            Indices depicting the position of the input sequence tokens in the sequence.
+        batch_size (`torch.Tensor`):
+            Batch size.
+        is_training (`bool`):
+            Whether the model is in training mode or in inference. The condition is checked by presence/absence of `token_type_ids/labels`
+    """
+    if attention_mask is not None and attention_mask.dim() == 4:
+        # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+        causal_mask = attention_mask
+    else:
+        causal_mask = torch.full(
+            (sequence_length, target_length),
+            fill_value=min_dtype,
+            dtype=dtype,
+            device=device,
+        )
+        # Causal diagonal mask only if training, otherwise attend to the whole prefix. Training-specific attn for prefix is handled below
+        if sequence_length != 1:
+            if is_training:
+                causal_mask = torch.triu(causal_mask, diagonal=1)
+            else:
+                causal_mask[:, :sequence_length] = 0.0
+
+        causal_mask *= torch.arange(
+            target_length, device=cache_position.device
+        ) > cache_position.reshape(-1, 1)
+        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+        if attention_mask is not None:
+            causal_mask = (
+                causal_mask.clone()
+            )  # copy to contiguous memory for in-place edit
+            mask_length = attention_mask.shape[-1]
+            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[
+                :, None, None, :
+            ].to(causal_mask.device)
+            padding_mask = padding_mask == 0
+            causal_mask[:, :, :, :mask_length] = causal_mask[
+                :, :, :, :mask_length
+            ].masked_fill(padding_mask, min_dtype)
+            # we are training thus we need to create a full mask on the image + prefix but causal on suffix
+            if is_training:
+                causal_mask[:, :, :, :mask_length] = causal_mask[
+                    :, :, :, :mask_length
+                ].masked_fill(
+                    token_type_ids[:, None, None, :].to(causal_mask.device) == 0, 0
+                )
+    return causal_mask
+
+
+# Copied from transformers.models.idefics2.processing_idefics2.is_url
+def is_url(val) -> bool:
+    return isinstance(val, str) and val.startswith("http")
+
+
+# Copied from transformers.models.idefics2.processing_idefics2.is_image_or_image_url
+def is_image_or_image_url(elem):
+    return is_url(elem) or is_valid_image(elem)
+
+
+def _is_str_or_image(elem):
+    return isinstance(elem, (str)) or is_image_or_image_url(elem)
+
+
+def generate_partial_autoregressive_mask(sz, start, end):
+    mask = torch.zeros(sz, sz).bool()
+    mask[start:end, start:end] = torch.triu(
+        torch.ones(end - start, end - start, dtype=torch.bool), diagonal=1
+    )
+    mask[:start, start:end] = True
+    mask[end:, start:end] = True
+    return mask
+
+
+def build_string_from_input(prompt, bos_token, image_seq_len, image_token, num_images):
+
+    return f"{image_token * image_seq_len * num_images}{bos_token}{prompt}\n"
+
+
+def is_torchdynamo_compiling():
+
+    try:
+        import torch
+
+        return torch.compiler.is_compiling()
+    except Exception:
+        try:
+            import torch._dynamo as dynamo  # noqa: F401
+
+            return dynamo.is_compiling()
+        except Exception:
+            return False