upload code

app.py

@@ -0,0 +1,253 @@
+import spaces
+
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+import os
+import random
+import shutil
+import pickle
+import gradio as gr
+import soundfile as sf
+from pathlib import Path
+
+import torch
+import torchaudio
+
+from huggingface_hub import hf_hub_download
+
+from infer import load_model, eval_model
+from spkr import SpeakerEmbedding
+
+
+spkr_model = SpeakerEmbedding(device="cuda")
+model, tokenizer, tokenizer_voila, model_type = load_model("maitrix-org/Voila-chat", "maitrix-org/Voila-Tokenizer")
+default_ref_file = "examples/character_ref_emb_demo.pkl"
+default_ref_name = "Homer Simpson"
+million_voice_ref_file = hf_hub_download(repo_id="maitrix-org/Voila-million-voice", filename="character_ref_emb_chunk0.pkl", repo_type="dataset")
+
+instruction = "You are a smart AI agent created by Maitrix.org."
+save_path = os.environ.get("GRADIO_TEMP_DIR", tempfile.gettempdir())
+
+intro = """**Voila**
+
+For more demos, please goto [https://voila.maitrix.org](https://voila.maitrix.org)."""
+
+default_ref_emb_mask_list = pickle.load(open(default_ref_file, "rb"))
+million_voice_ref_emb_mask_list = pickle.load(open(million_voice_ref_file, "rb"))
+
+def get_ref_embs(ref_audio):
+    wav, sr = torchaudio.load(ref_audio)
+    ref_embs = spkr_model(wav, sr).cpu()
+    return ref_embs
+
+def delete_directory(request: gr.Request):
+    if not request.session_hash:
+        return
+    user_dir = Path(f"{save_path}/{str(request.session_hash)}")
+    if user_dir.exists():
+        shutil.rmtree(str(user_dir))
+
+def add_message(history, message):
+    history.append({"role": "user", "content": {"path": message}})
+    return history, gr.Audio(value=None), gr.Button(interactive=False)
+
+def call_bot(history, ref_embs, request: gr.Request):
+    formated_history = {
+        "instruction": instruction,
+        "conversations": [{'from': item["role"], 'audio': {"file": item["content"][0]}} for item in history],
+    }
+    formated_history["conversations"].append({"from": "assistant"})
+    print(formated_history)
+    ref_embs = torch.tensor(ref_embs, dtype=torch.float32, device="cuda")
+    ref_embs_mask = torch.tensor([1], device="cuda")
+    out = eval_model(model, tokenizer, tokenizer_voila, model_type, "chat_aiao", formated_history, ref_embs, ref_embs_mask, max_new_tokens=512)
+    if 'audio' in out:
+        wav, sr = out['audio']
+
+        user_dir = Path(f"{save_path}/{str(request.session_hash)}")
+        user_dir.mkdir(exist_ok=True)
+        save_name = f"{user_dir}/{len(history)}.wav"
+        sf.write(save_name, wav, sr)
+
+        history.append({"role": "assistant", "content": {"path": save_name}})
+    else:
+        history.append({"role": "assistant", "content": {"text": out['text']}})
+
+    return history
+
+def run_tts(text, ref_embs):
+    formated_history = {
+        "instruction": "",
+        "conversations": [{'from': "user", 'text': text}],
+    }
+    formated_history["conversations"].append({"from": "assistant"})
+    ref_embs = torch.tensor(ref_embs, dtype=torch.float32, device="cuda")
+    ref_embs_mask = torch.tensor([1], device="cuda")
+    out = eval_model(model, tokenizer, tokenizer_voila, model_type, "chat_tts", formated_history, ref_embs, ref_embs_mask, max_new_tokens=512)
+    if 'audio' in out:
+        wav, sr = out['audio']
+        return (sr, wav)
+    else:
+        raise Exception("No audio output")
+
+def run_asr(audio):
+    formated_history = {
+        "instruction": "",
+        "conversations": [{'from': "user", 'audio': {"file": audio}}],
+    }
+    formated_history["conversations"].append({"from": "assistant"})
+    out = eval_model(model, tokenizer, tokenizer_voila, model_type, "chat_asr", formated_history, None, None, max_new_tokens=512)
+    if 'text' in out:
+        return out['text']
+    else:
+        raise Exception("No text output")
+
+
+def markdown_ref_name(ref_name):
+    return f"### Current voice id: {ref_name}"
+
+def random_million_voice():
+    voice_id = random.choice(list(million_voice_ref_emb_mask_list.keys()))
+    return markdown_ref_name(voice_id), million_voice_ref_emb_mask_list[voice_id]
+
+def get_ref_modules(cur_ref_embs):
+    with gr.Row() as ref_row:
+        with gr.Row():
+            current_ref_name = gr.Markdown(markdown_ref_name(default_ref_name))
+        with gr.Row() as ref_name_row:
+            with gr.Column(scale=2, min_width=160):
+                ref_name_dropdown = gr.Dropdown(
+                    choices=list(default_ref_emb_mask_list.keys()),
+                    value=default_ref_name,
+                    label="Reference voice",
+                    min_width=160,
+                )
+            with gr.Column(scale=1, min_width=80):
+                random_ref_button = gr.Button(
+                    "Random from Million Voice", size="md",
+                )
+        with gr.Row(visible=False) as ref_audio_row:
+            with gr.Column(scale=2, min_width=80):
+                ref_audio = gr.Audio(
+                    sources=["microphone", "upload"],
+                    type="filepath",
+                    show_label=False,
+                    min_width=80,
+                )
+            with gr.Column(scale=1, min_width=80):
+                change_ref_button = gr.Button(
+                    "Change voice",
+                    interactive=False,
+                    min_width=80,
+                )
+    ref_name_dropdown.change(
+        lambda x: (markdown_ref_name(x), default_ref_emb_mask_list[x]),
+        ref_name_dropdown,
+        [current_ref_name, cur_ref_embs]
+    )
+    random_ref_button.click(
+        random_million_voice,
+        None,
+        [current_ref_name, cur_ref_embs],
+    )
+    ref_audio.input(lambda: gr.Button(interactive=True), None, change_ref_button)
+    # If custom ref voice checkbox is checked, show the Audio component to record or upload a reference voice
+    custom_ref_voice = gr.Checkbox(label="Use custom voice", value=False)
+    # Checked: enable audio and button
+    # Unchecked: disable audio and button
+    def custom_ref_voice_change(x, cur_ref_embs, cur_ref_embs_mask):
+        if not x:
+            cur_ref_embs = default_ref_emb_mask_list[default_ref_name]
+        return [gr.Row(visible=not x), gr.Audio(value=None), gr.Row(visible=x), markdown_ref_name("Custom voice"), cur_ref_embs]
+    custom_ref_voice.change(
+        custom_ref_voice_change,
+        [custom_ref_voice, cur_ref_embs],
+        [ref_name_row, ref_audio, ref_audio_row, current_ref_name, cur_ref_embs]
+    )
+    # When change ref button is clicked, get the reference voice and update the reference voice state
+    change_ref_button.click(
+        lambda: gr.Button(interactive=False), None, [change_ref_button]
+    ).then(
+        get_ref_embs, ref_audio, cur_ref_embs
+    )
+    return ref_row
+
+def get_chat_tab():
+    cur_ref_embs = gr.State(default_ref_emb_mask_list[default_ref_name])
+    with gr.Row() as chat_tab:
+        with gr.Column(scale=1):
+            ref_row = get_ref_modules(cur_ref_embs)
+            # Voice chat input
+            chat_input = gr.Audio(
+                sources=["microphone", "upload"],
+                type="filepath",
+                show_label=False,
+            )
+            submit = gr.Button("Submit", interactive=False)
+            gr.Markdown(intro)
+        with gr.Column(scale=9):
+            chatbot = gr.Chatbot(
+                elem_id="chatbot",
+                type="messages",
+                bubble_full_width=False,
+                scale=1,
+                show_copy_button=False,
+                avatar_images=(
+                    None,  # os.path.join("files", "avatar.png"),
+                    None, # os.path.join("files", "avatar.png"),
+                ),
+            )
+
+        chat_input.input(lambda: gr.Button(interactive=True), None, submit)
+        submit.click(
+            add_message, [chatbot, chat_input], [chatbot, chat_input, submit]
+        ).then(
+            call_bot, [chatbot, cur_ref_embs], chatbot, api_name="bot_response"
+        )
+    return chat_tab
+
+def get_tts_tab():
+    cur_ref_embs = gr.State(default_ref_emb_mask_list[default_ref_name])
+    with gr.Row() as tts_tab:
+        with gr.Column(scale=1):
+            ref_row = get_ref_modules(cur_ref_embs)
+            gr.Markdown(intro)
+        with gr.Column(scale=9):
+            tts_output = gr.Audio(label="TTS output", interactive=False)
+            with gr.Row():
+                text_input = gr.Textbox(label="Text", placeholder="Text to TTS")
+                submit = gr.Button("Submit")
+        submit.click(
+            run_tts, [text_input, cur_ref_embs], tts_output
+        )
+    return tts_tab
+
+def get_asr_tab():
+    with gr.Row() as asr_tab:
+        with gr.Column():
+            asr_input = gr.Audio(
+                label="ASR input",
+                sources=["microphone", "upload"],
+                type="filepath",
+            )
+            submit = gr.Button("Submit")
+            gr.Markdown(intro)
+        with gr.Column():
+            asr_output = gr.Textbox(label="ASR output", interactive=False)
+    submit.click(
+        run_asr, [asr_input], asr_output
+    )
+    return asr_tab
+
+with gr.Blocks(fill_height=True) as demo:
+    with gr.Tab("Chat"):
+        chat_tab = get_chat_tab()
+    with gr.Tab("TTS"):
+        tts_tab = get_tts_tab()
+    with gr.Tab("ASR"):
+        asr_tab = get_asr_tab()
+    demo.unload(delete_directory)
+
+if __name__ == "__main__":
+    demo.launch()

audio_transformer.py

@@ -0,0 +1,354 @@
+import math
+from typing import Optional
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+
+from einops import rearrange
+
+
+@dataclass
+class LocalArgs:
+    codebook_size: int = 2048
+    num_codebooks: int = 4
+
+# Modified from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L105
+class KVCache(nn.Module):
+    def __init__(
+        self, n_layer, batch_size, max_seq_len, n_heads, head_dim, dtype, device
+    ):
+        super().__init__()
+        cache_shape = (n_layer, batch_size, n_heads, max_seq_len, head_dim)
+        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype, device=device))
+        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype, device=device))
+
+    def update(self, layer_idx, input_pos, k_val, v_val):
+        # k_val: [B, H, S, D]
+
+        k_out = self.k_cache
+        v_out = self.v_cache
+        k_out[layer_idx, :, :, input_pos:input_pos+1] = k_val
+        v_out[layer_idx, :, :, input_pos:input_pos+1] = v_val
+
+        return k_out[layer_idx], v_out[layer_idx]
+
+# Modified from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L756
+def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000) -> Tensor:
+    freqs = 1.0 / (
+        base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)
+    )
+    t = torch.arange(seq_len, device=freqs.device)
+    freqs = torch.outer(t, freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+    cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+    return cache
+
+# Copied from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L767
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+    xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+    freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+    x_out2 = torch.stack(
+        [
+            xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+            xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+        ],
+        -1,
+    )
+
+    x_out2 = x_out2.flatten(3)
+    return x_out2.type_as(x)
+
+# Copied from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L742
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x: Tensor) -> Tensor:
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+# Copied from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L731
+class FeedForward(nn.Module):
+    def __init__(self, config: LocalArgs) -> None:
+        super().__init__()
+        self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+# Modified from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L615
+class Attention(nn.Module):
+    def __init__(self, config: LocalArgs, layer_idx: int, use_sdpa: bool = True):
+        super().__init__()
+        assert config.dim % config.n_head == 0
+        self.layer_idx = layer_idx
+
+        total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+        # key, query, value projections for all heads, but in a batch
+        self.wqkv = nn.Linear(
+            config.dim, total_head_dim, bias=config.attention_qkv_bias
+        )
+        self.wo = nn.Linear(config.dim, config.dim, bias=False)
+
+        self.dropout = config.dropout
+        self.n_head = config.n_head
+        self.head_dim = config.head_dim
+        self.n_local_heads = config.n_local_heads
+        self.dim = config.dim
+        self.use_sdpa = use_sdpa
+        self._register_load_state_dict_pre_hook(self.load_hook)
+
+    def load_hook(self, state_dict, prefix, *args):
+        if prefix + "wq.weight" in state_dict:
+            wq = state_dict.pop(prefix + "wq.weight")
+            wk = state_dict.pop(prefix + "wk.weight")
+            wv = state_dict.pop(prefix + "wv.weight")
+            state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+    def forward(
+        self,
+        x: Tensor,
+        freqs_cis: Tensor,
+        mask: Tensor,
+        input_pos: Optional[int] = None,
+        kv_cache: Optional[KVCache] = None,
+    ) -> Tensor:
+        bsz, seqlen, _ = x.shape
+
+        kv_size = self.n_local_heads * self.head_dim
+        q, k, v = self.wqkv(x).split([self.dim, kv_size, kv_size], dim=-1)
+
+        q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+        k = k.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+        v = v.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+
+        q = apply_rotary_emb(q, freqs_cis)
+        k = apply_rotary_emb(k, freqs_cis)
+
+        q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+        if kv_cache is not None:
+            k, v = kv_cache.update(self.layer_idx, input_pos, k, v)
+
+        k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+        v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+
+        if self.use_sdpa:
+            if mask is None:
+                with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+                    y = F.scaled_dot_product_attention(
+                        q,
+                        k,
+                        v,
+                        dropout_p=self.dropout if self.training else 0.0,
+                        is_causal=True,
+                        # No third party attn_mask here to use flash_attention
+                    )
+            else:
+                y = F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=mask,
+                    dropout_p=self.dropout if self.training else 0.0,
+                )
+        else:
+            y = self.eq_scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=mask,
+                dropout_p=self.dropout if self.training else 0.0,
+            )
+
+        y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim)
+
+        return self.wo(y)
+
+    def eq_scaled_dot_product_attention(
+        self,
+        query,
+        key,
+        value,
+        attn_mask=None,
+        dropout_p=0.0,
+    ) -> torch.Tensor:
+        # This is a standard scaled dot product attention
+        # It's low efficient, but it doesn't raise cuda error
+
+        L, S = query.size(-2), key.size(-2)
+        scale_factor = 1 / math.sqrt(query.size(-1))
+        attn_bias = torch.zeros(1, 1, L, S, dtype=query.dtype, device=query.device)
+
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+            else:
+                attn_bias += attn_mask
+
+        attn_weight = query @ key.transpose(-2, -1) * scale_factor
+        attn_weight += attn_bias
+        attn_weight = torch.softmax(attn_weight, dim=-1)
+        attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+
+        return attn_weight @ value
+
+# Copied from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L599
+class TransformerBlock(nn.Module):
+    def __init__(self, config: LocalArgs, layer_idx: int, use_sdpa: bool = True) -> None:
+        super().__init__()
+        self.attention = Attention(config, layer_idx, use_sdpa=use_sdpa)
+        self.feed_forward = FeedForward(config)
+        self.ffn_norm = RMSNorm(config.dim, config.norm_eps)
+        self.attention_norm = RMSNorm(config.dim, config.norm_eps)
+
+    def forward(
+        self, x: Tensor, freqs_cis: Tensor, mask: Tensor, input_pos: int = None, kv_cache: KVCache = None
+    ) -> Tensor:
+        h = x + self.attention(self.attention_norm(x), freqs_cis, mask, input_pos, kv_cache)
+        out = h + self.feed_forward(self.ffn_norm(h))
+        return out
+
+# Modified from https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py#L470
+class AudioTransformer(nn.Module):
+    def __init__(self, config, use_sdpa: bool = False):
+        super().__init__()
+        self.config = LocalArgs()
+        self.config.codebook_size = config.codebook_size
+        self.config.num_codebooks = config.num_codebooks
+        if hasattr(config, "min_audio_token_id"):
+            self.config.min_audio_token_id = config.min_audio_token_id
+            self.config.max_audio_token_id = config.max_audio_token_id
+        self.config.n_layer = 4
+        self.config.dim = 1024
+        self.config.n_head = 32
+        self.config.n_local_heads = 32
+        self.config.intermediate_size = 2816
+        self.config.head_dim = self.config.dim // self.config.n_head
+        self.config.norm_eps = 1e-5
+        self.config.attention_qkv_bias = False
+        self.config.dropout = 0.0
+
+        self.embeddings = nn.Embedding(self.config.codebook_size, self.config.dim)
+        if self.config.dim != config.hidden_size:
+            self.input_proj = nn.Linear(config.hidden_size, self.config.dim, bias=False)
+        else:
+            self.input_proj = nn.Identity()
+        self.layers = nn.ModuleList(
+                TransformerBlock(self.config, layer_idx, use_sdpa=use_sdpa) for layer_idx in range(self.config.n_layer)
+        )
+        self.norm = RMSNorm(self.config.dim, eps=self.config.norm_eps)
+        self.token_head = nn.Linear(self.config.dim, self.config.codebook_size, bias=False)
+        self.gradient_checkpointing = False
+
+        self.register_buffer(
+            "freqs_cis",
+            precompute_freqs_cis(self.config.num_codebooks, self.config.dim // self.config.n_head, 10000),
+            persistent=False,
+        )
+        self.register_buffer(
+            "attention_mask",
+            torch.tril(torch.ones(self.config.num_codebooks, self.config.num_codebooks, dtype=torch.bool)),
+            persistent=False,
+        )
+
+    def run_model(self, hidden_states, freqs_cis, attention_mask, input_pos: int = None, kv_cache: KVCache = None):
+        for layer in self.layers:
+            # TODO: gradient_checkpointing is disabled because of bug
+            if False: # self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    freqs_cis,
+                    attention_mask,
+                    use_reentrant=True,
+                )
+            else:
+                hidden_states = layer(hidden_states, freqs_cis, attention_mask, input_pos, kv_cache)
+        hidden_states = self.norm(hidden_states)
+        logits = self.token_head(hidden_states)
+        return logits.float()
+
+    # inp: [bs, hidden_size]
+    # labels: [bs, num_codebooks]
+    # logits: [bs, num_codebooks, codebook_size]
+    def forward(self, inp, labels):
+        bs = inp.shape[0]
+
+        hidden_states = self.input_proj(inp)
+        if self.freqs_cis.dtype != hidden_states.dtype:
+            self.freqs_cis = self.freqs_cis.to(dtype=hidden_states.dtype)
+        if labels is not None:
+            # Training mode
+            # Get embedding
+            assert bs == labels.shape[0] and labels.shape[1] == self.config.num_codebooks, f"Labels shape error: {labels.shape}"
+            hidden_states = [hidden_states[:, None, :], self.embeddings(labels[..., :-1]).to(hidden_states.dtype)]
+            hidden_states = torch.cat(hidden_states, dim=1) # [bs, num_codebooks, hidden_size]
+            # Run attention layers
+            logits = self.run_model(hidden_states, self.freqs_cis, self.attention_mask)
+        else:
+            # Inference mode
+            raise RuntimeError(f"Please call function \"inference\" in inference mode")
+        return logits
+
+    # inp: [bs, seq_len, hidden_size]
+    # out_tokens: [bs, 1, num_codebooks]
+    @torch.inference_mode()
+    def inference(self, inp, temperature=0, top_k=0):
+        # Only use the last hidden states for token computation
+        inp = inp[:, -1:, :]
+
+        bs = inp.shape[0]
+        if self.freqs_cis.dtype != inp.dtype:
+            self.freqs_cis = self.freqs_cis.to(dtype=inp.dtype)
+
+        inp = self.input_proj(inp)
+
+        # Inference mode
+        kv_cache = KVCache(
+                self.config.n_layer,
+                bs,
+                self.config.num_codebooks,
+                self.config.n_head,
+                self.config.head_dim,
+                dtype=inp.dtype,
+                device=inp.device,
+        )
+        # Generate one token per step
+        out_tokens = []
+        for input_pos in range(self.config.num_codebooks):
+            inp = inp.reshape(bs, 1, self.config.dim)
+            local_freqs_cis = self.freqs_cis[input_pos]
+            local_mask = self.attention_mask[None, None, input_pos, :self.config.num_codebooks]
+
+            logits = self.run_model(inp, local_freqs_cis, local_mask, input_pos, kv_cache)
+            logits = logits.squeeze(dim=1)
+
+            # Apply temperature and top-k
+            if temperature > 0:
+                logits = logits / temperature
+            if top_k > 0:
+                top_k = min(top_k, 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 = logits.masked_fill(indices_to_remove, -float("Inf"))
+
+            # Do sample
+            probs = nn.functional.softmax(logits, dim=-1)
+            next_tokens = torch.multinomial(probs, num_samples=1)
+
+            next_tokens = next_tokens.reshape(bs, 1, 1)
+            inp = self.embeddings(next_tokens)
+            out_tokens.append(next_tokens)
+
+        return torch.cat(out_tokens, dim=-1)

examples/character_ref_emb_demo.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a869512a59e4447c19ecb283d6e0097bf71eaf57e8fa98712afd7c41acbbb554
+size 23264

infer.py

@@ -0,0 +1,198 @@
+import os
+import argparse
+import random
+import jsonlines
+import soundfile as sf
+import json
+import copy
+import torch
+from pathlib import Path
+from threading import Thread
+
+import torchaudio
+from transformers import AutoTokenizer
+
+from model import VoilaAudioAlphaModel, VoilaModel, VoilaAutonomousModel
+from spkr import SpeakerEmbedding
+from voila_tokenizer import VoilaTokenizer
+from tokenize_func import (
+    voila_input_format,
+    AUDIO_TOKEN_FORMAT,
+    DEFAULT_AUDIO_TOKEN,
+    DEFAULT_ASSISTANT_TOKEN,
+)
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+def load_model(model_name, audio_tokenizer_path):
+    disable_torch_init()
+
+    if "Voila-audio" in model_name:
+        model_type = "audio"
+        cls = VoilaAudioAlphaModel
+    elif "Voila-auto" in model_name:
+        model_type = "autonomous"
+        cls = VoilaAutonomousModel
+    else:
+        model_type = "token"
+        cls = VoilaModel
+
+    model = cls.from_pretrained(
+        model_name,
+        torch_dtype=torch.bfloat16,
+        use_flash_attention_2=True,
+        use_cache=True,
+    )
+    model = model.cuda()
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer_voila = VoilaTokenizer(model_path=audio_tokenizer_path, device="cuda")
+    return model, tokenizer, tokenizer_voila, model_type
+
+def is_audio_output_task(task_type):
+    return task_type.endswith("ao") or "aiao" in task_type or "tts" in task_type
+
+def eval_model(model, tokenizer, tokenizer_voila, model_type, task_type, history, ref_embs, ref_embs_mask, max_new_tokens=512):
+    # step1: initializing
+    num_codebooks = model.config.num_codebooks
+    codebook_size = model.config.codebook_size
+
+    AUDIO_MIN_TOKEN_ID = tokenizer.convert_tokens_to_ids(AUDIO_TOKEN_FORMAT.format(0))
+    assert isinstance(AUDIO_MIN_TOKEN_ID, int)
+    AUDIO_MAX_TOKEN_ID = tokenizer.convert_tokens_to_ids(AUDIO_TOKEN_FORMAT.format(codebook_size*num_codebooks-1))
+    assert isinstance(AUDIO_MAX_TOKEN_ID, int)
+    AUDIO_TOKEN_ID = tokenizer.convert_tokens_to_ids(DEFAULT_AUDIO_TOKEN)
+    assert isinstance(AUDIO_TOKEN_ID, int)
+    ASSISTANT_TOKEN_ID = tokenizer.convert_tokens_to_ids(DEFAULT_ASSISTANT_TOKEN)
+    assert isinstance(ASSISTANT_TOKEN_ID, int)
+
+    # step2: set infer config
+    data_cfg = {
+        "input_type": model_type,
+        "task_type": task_type,
+        "num_codebooks": num_codebooks,
+        "codebook_size": codebook_size,
+    }
+
+    # step3: infer
+    input_ids, audio_datas, audio_data_masks, streaming_user_input_audio_tokens = voila_input_format(history, tokenizer, tokenizer_voila, data_cfg)
+
+    # prepare user_streaming_generator to simulate streaming user input
+    def get_input_generator(all_tokens):
+        assert all_tokens is not None
+        for i in range(len(all_tokens[0])):
+            yield all_tokens[:,i]
+
+    if model_type == "autonomous":
+        input_generator = get_input_generator(torch.as_tensor(streaming_user_input_audio_tokens).cuda())
+        input_ids = [torch.as_tensor([input]).transpose(1,2).cuda() for input in input_ids]            # transpose to [bs, seq, num_codebooks]
+        input_ids = torch.cat(input_ids, dim=2)            # concat to [bs, seq, num_codebooks*2]
+    else:
+        input_ids = torch.as_tensor([input_ids]).transpose(1,2).cuda()      # transpose to [bs, seq, num_codebooks]
+    gen_params = {
+        "input_ids": input_ids,
+        "ref_embs": ref_embs,
+        "ref_embs_mask": ref_embs_mask,
+        "max_new_tokens": max_new_tokens,
+        "pad_token_id": tokenizer.pad_token_id,
+        "eos_token_id": tokenizer.eos_token_id,
+        "llm_audio_token_id": AUDIO_TOKEN_ID,
+        "min_audio_token_id": AUDIO_MIN_TOKEN_ID,
+        "temperature": 0.2,
+        "top_k": 50,
+        "audio_temperature": 0.8,
+        "audio_top_k": 50,
+    }
+    if model_type == "audio":
+        audio_datas = torch.tensor([audio_datas], dtype=torch.bfloat16).cuda()
+        audio_data_masks = torch.tensor([audio_data_masks]).cuda()
+        gen_params["audio_datas"] = audio_datas
+        gen_params["audio_data_masks"] = audio_data_masks
+    elif model_type == "autonomous":
+        gen_params["input_generator"] = input_generator
+        gen_params["llm_assistant_token_id"] = ASSISTANT_TOKEN_ID
+    print(f"Input str: {tokenizer.decode(input_ids[0, :, 0])}")
+    with torch.inference_mode():
+        outputs = model.run_generate(**gen_params)
+
+        if model_type == "autonomous":
+            outputs = outputs.chunk(2, dim=2)[1]
+        outputs = outputs[0].cpu().tolist()
+
+        predict_outputs = outputs[input_ids.shape[1]:]
+        text_outputs = []
+        audio_outputs = []
+        for _ in range(num_codebooks):
+            audio_outputs.append([])
+        for item in predict_outputs:
+            if item[0] >= AUDIO_MIN_TOKEN_ID and item[0] <= AUDIO_MAX_TOKEN_ID:
+                for n, at in enumerate(item):
+                    audio_outputs[n].append((at - AUDIO_MIN_TOKEN_ID)%codebook_size)
+            elif item[0] != tokenizer.eos_token_id:
+                text_outputs.append(item[0])
+
+        out ={
+            'text': tokenizer.decode(text_outputs),
+        }
+        if is_audio_output_task(task_type):
+            audio_values = tokenizer_voila.decode(torch.tensor(audio_outputs).cuda())
+            out['audio'] = (audio_values.detach().cpu().numpy(), 16000)
+        return out
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--instruction", type=str, default="")
+    parser.add_argument("--input-text", type=str, default=None)
+    parser.add_argument("--input-audio", type=str, default=None)
+    parser.add_argument("--result-path", type=str, default="output")
+    parser.add_argument("--ref-audio", type=str, default="examples/test1.mp3")
+    parser.add_argument("--model-name", type=str, default="maitrix-org/Voila-chat")
+    parser.add_argument("--audio-tokenizer-path", type=str, default="maitrix-org/Voila-Tokenizer")
+    parser.add_argument("--task-type", type=str, default="chat_aiao")
+    args = parser.parse_args()
+
+    assert args.model_name in [
+        "maitrix-org/Voila-audio-alpha",
+        "maitrix-org/Voila-base",
+        "maitrix-org/Voila-chat",
+        "maitrix-org/Voila-autonomous-preview",
+    ]
+
+    # step0: Model loading
+    model, tokenizer, tokenizer_voila, model_type = load_model(args.model_name, args.audio_tokenizer_path)
+
+    # step1: prepare inputs
+    Path(args.result_path).mkdir(exist_ok=True, parents=True)
+    history = {
+        "instruction": args.instruction,
+        "conversations": [],
+    }
+    if args.input_text is not None:
+        history["conversations"].append({"from": "user", "text": args.input_text})
+    elif args.input_audio is not None:
+        history["conversations"].append({"from": "user", "audio": {"file": args.input_audio}})
+    else:
+        raise Exception("Please provide atleast one of --input-text and --input-audio")
+    history["conversations"].append({"from": "assistant"})
+
+    # step2: encode ref
+    ref_embs, ref_embs_mask = None, None
+    if is_audio_output_task(args.task_type):
+        spkr_model = SpeakerEmbedding(device="cuda")
+        wav, sr = torchaudio.load(args.ref_audio)
+        ref_embs = spkr_model(wav, sr)
+        ref_embs_mask = torch.tensor([1]).cuda()
+
+    out = eval_model(model, tokenizer, tokenizer_voila, model_type, args.task_type, history, ref_embs, ref_embs_mask)
+    print(f"Output str: {out['text']}")
+    if 'audio' in out:
+        wav, sr = out['audio']
+        save_name = f"{args.result_path}/out.wav"
+        sf.write(save_name, wav, sr)

model.py

@@ -0,0 +1,1397 @@
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union, Dict, Any
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.utils import ModelOutput, logging
+from transformers.models.llama.modeling_llama import LlamaModel, LlamaPreTrainedModel
+
+from audio_transformer import AudioTransformer
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from https://github.com/pytorch/audio/blob/main/src/torchaudio/models/wav2vec2/components.py#L43
+class LayerNorm(torch.nn.LayerNorm):
+    """Layer norm with transpose"""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        x = input.transpose(-2, -1)
+        x = torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.transpose(-2, -1)
+        return x
+
+# Copied from https://github.com/pytorch/audio/blob/main/src/torchaudio/models/wav2vec2/components.py#L53
+class ConvLayerBlock(torch.nn.Module):
+    """Convolution unit of FeatureExtractor"""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+        bias: bool,
+        layer_norm: Optional[torch.nn.Module],
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.layer_norm = layer_norm
+        self.conv = torch.nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            bias=bias,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Args:
+            x (Tensor): Shape: ``[batch, in_channels, in_frame]``.
+        Returns:
+            Tensor: Shape ``[batch, out_channels, out_frames]``.
+            Optional[Tensor]: Shape ``[batch, ]``.
+        """
+        x = self.conv(x)
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+        x = torch.nn.functional.gelu(x)
+
+        return x
+
+# Copied from https://github.com/pytorch/audio/blob/main/src/torchaudio/models/wav2vec2/components.py#L146
+class FeatureProjection(torch.nn.Module):
+    """Layer that connects FeatureExtractor and Encoder
+
+    Projects features to encoder dimension.
+
+    Args:
+        in_features (int): Input feature dim.
+        out_features (int): Output feature dim.
+        dropout (float): Dropout probability.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        dropout=0.1,
+    ):
+        super().__init__()
+        self.layer_norm = torch.nn.LayerNorm(in_features)
+        self.projection = torch.nn.Linear(
+            in_features,
+            out_features,
+        )
+        self.dropout = torch.nn.Dropout(dropout)
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor):
+                Feature Tensor. shape: ``[batch, frame, in_feature]``
+        Returns:
+            Tensor: Projected features. ``[batch, frame, out_feature]``.
+        """
+        x = self.layer_norm(x)
+        x = self.projection(x)
+        x = self.dropout(x)
+        return x
+
+# Modified from https://github.com/pytorch/audio/blob/main/src/torchaudio/models/wav2vec2/components.py#L102
+class FeatureExtractor(torch.nn.Module):
+    """Extract features from audio
+
+    Args:
+        conv_layers (nn.ModuleList):
+            convolution layers
+    """
+
+    def __init__(
+        self,
+        shapes=[(512, 10, 5), (512, 3, 2), (512, 3, 2), (512, 3, 2), (512, 3, 2), (512, 2, 2), (512, 2, 2)],
+        bias=False,
+        norm_mode="group_norm",
+    ):
+        super().__init__()
+        if norm_mode not in ["group_norm", "layer_norm"]:
+            raise ValueError("Invalid norm mode")
+        blocks = []
+        in_channels = 1
+        for i, (out_channels, kernel_size, stride) in enumerate(shapes):
+            normalization = None
+            if norm_mode == "group_norm" and i == 0:
+                normalization = torch.nn.GroupNorm(
+                    num_groups=out_channels,
+                    num_channels=out_channels,
+                    affine=True,
+                )
+            elif norm_mode == "layer_norm":
+                normalization = LayerNorm(
+                    normalized_shape=out_channels,
+                    elementwise_affine=True,
+                )
+            blocks.append(
+                ConvLayerBlock(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    bias=bias,
+                    layer_norm=normalization,
+                )
+            )
+            in_channels = out_channels
+        self.conv_layers = torch.nn.ModuleList(blocks)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Args:
+            x (Tensor):
+                Input Tensor representing a batch of audio,
+                shape: ``[batch, time]``.
+
+        Returns:
+            Tensor:
+                The resulting feature, shape: ``[batch, frame, feature]``
+            Optional[Tensor]:
+                Valid length of each output sample. shape: ``[batch, ]``.
+        """
+        if x.ndim != 2:
+            raise ValueError(f"Expected the input Tensor to be 2D (batch, time). Found: {list(x.shape)}")
+
+        x = x.unsqueeze(1)  # (batch, channel==1, frame)
+        for layer in self.conv_layers:
+            x = layer(x)  # (batch, feature, frame)
+        x = x.transpose(1, 2)  # (batch, frame, feature)
+        return x
+
+# Modified from https://github.com/pytorch/audio/blob/main/src/torchaudio/models/wav2vec2/components.py#L102
+class FeatureExtractorAdapter(torch.nn.Module):
+    """Extract features from audio
+
+    Args:
+        conv_layers (nn.ModuleList):
+            convolution layers
+    """
+
+    def __init__(
+        self,
+        shapes=(512, 512, 2, 2),
+        hidden_size=2048,
+        bias=False,
+        norm_mode="group_norm",
+    ):
+        super().__init__()
+        if norm_mode not in ["group_norm", "layer_norm"]:
+            raise ValueError("Invalid norm mode")
+        in_channels, out_channels, kernel_size, stride = shapes
+        normalization = LayerNorm(
+            normalized_shape=out_channels,
+            elementwise_affine=True,
+        )
+        self.conv_layers = ConvLayerBlock(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    bias=False,
+                    layer_norm=normalization,
+                )
+        self.feat_proj = FeatureProjection(out_channels, hidden_size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Args:
+            x (Tensor):
+                Input Tensor representing a batch of audio,
+                shape: ``[batch, time]``.
+
+        Returns:
+            Tensor:
+                The resulting feature, shape: ``[batch, frame, feature]``
+            Optional[Tensor]:
+                Valid length of each output sample. shape: ``[batch, ]``.
+        """
+        x = x.transpose(1, 2)  # (batch, feature, frame)
+        x = self.conv_layers(x)  # (batch, feature, frame)
+        x = x.transpose(1, 2)  # (batch, frame, feature)
+        x = self.feat_proj(x)
+        return x
+
+@dataclass
+class VoilaOutput(ModelOutput):
+    """
+    Modified from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_outputs.py#L678
+
+    Base class for Voila outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            The hidden state of the last attention layer.
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    last_hidden_state: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    voila_pred: Optional[torch.FloatTensor] = None
+
+
+# Modified from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1103
+class VoilaModel(LlamaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.pad_vocab_size_multiple = 64
+
+        self.ref_emb_linear = nn.Linear(256, config.hidden_size, bias=True)
+        self.audio_transformer = AudioTransformer(config, use_sdpa=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        audio_labels: Optional[torch.LongTensor] = None,
+        ref_embs: Optional[List[torch.Tensor]] = None,
+        ref_embs_mask: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        num_logits_to_keep: int = 0,
+    ) -> Union[Tuple, VoilaOutput]:
+        r"""
+        Args:
+            input_ids: [bs, seq_len, num_codebooks]
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if inputs_embeds is None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+        assert len(inputs_embeds.shape) == 4
+        if len(inputs_embeds.shape) == 4:
+            inputs_embeds = inputs_embeds.mean(dim=2)
+
+        if self.training or \
+                (past_key_values is None and ref_embs is not None) or \
+                (past_key_values is not None and past_key_values.get_seq_length() < 4 and ref_embs is not None):
+            ref_embs = self.ref_emb_linear(ref_embs.to(self.ref_emb_linear.weight.dtype))
+            ref_embs = ref_embs * ref_embs_mask.unsqueeze(-1).unsqueeze(-1)
+            # (padding_left,padding_right,padding_top,padding_bottom,padding_front,padding_back)
+            padding = (0, 0, 4, inputs_embeds.shape[1] - 5, 0, 0)
+            ref_embs = torch.nn.functional.pad(ref_embs, padding, mode='constant', value=0.0)
+            inputs_embeds = inputs_embeds + ref_embs
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        if self.config.pretraining_tp > 1:
+            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
+            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
+            logits = torch.cat(logits, dim=-1)
+        else:
+            # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+            logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+
+        loss = None
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return VoilaOutput(
+            loss=loss,
+            logits=logits,
+            last_hidden_state=hidden_states,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _prepare_inputs_for_generation(
+        self, input_ids, ref_embs=None, ref_embs_mask=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None and past_key_values.get_seq_length() > 0:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_cache_shape()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is None and \
+                (past_key_values is None or past_key_values.get_seq_length() <= 0):
+            inputs_embeds = self.model.embed_tokens(input_ids)
+        if inputs_embeds is not None and \
+                (past_key_values is None or past_key_values.get_seq_length() <= 0):
+            model_inputs = {"inputs_embeds": inputs_embeds, "ref_embs": ref_embs, "ref_embs_mask": ref_embs_mask}
+        else:
+            model_inputs = {"input_ids": input_ids, "ref_embs": None}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs,
+        model_kwargs: Dict[str, Any],
+        num_new_token: int = 1,
+    ) -> Dict[str, Any]:
+        # update past_key_values
+        model_kwargs["past_key_values"] = outputs.past_key_values
+
+        # update attention mask
+        if "attention_mask" in model_kwargs:
+            attention_mask = model_kwargs["attention_mask"]
+            model_kwargs["attention_mask"] = torch.cat(
+                [attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_token))], dim=-1
+            )
+
+        return model_kwargs
+
+    def _prepare_attention_mask_for_generation(
+        self,
+        inputs: torch.Tensor,
+        pad_token_id: Optional[int],
+        eos_token_id: Optional[Union[int, List[int]]],
+    ) -> torch.LongTensor:
+        is_input_ids = len(inputs.shape) == 2 and inputs.dtype in [torch.int, torch.long]
+        is_pad_token_in_inputs = (pad_token_id is not None) and (pad_token_id in inputs)
+        if isinstance(eos_token_id, int):
+            eos_token_id = [eos_token_id]
+        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (pad_token_id not in eos_token_id)
+
+        # Check if input is input_ids and padded -> only then is attention_mask defined
+        if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id:
+            return inputs.ne(pad_token_id).long()
+        else:
+            return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device)
+
+    @torch.inference_mode()
+    def run_generate(
+        self,
+        input_ids: torch.LongTensor,
+        ref_embs: Optional[List[torch.Tensor]] = None,
+        ref_embs_mask: Optional[torch.LongTensor] = None,
+        max_new_tokens: Optional[int] = 128,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        streamer: Optional["BaseStreamer"] = None,
+        llm_audio_token_id: Optional[int] = None,
+        min_audio_token_id: Optional[int] = None,
+        temperature=0.2,
+        top_k=50,
+        audio_temperature=0.2,
+        audio_top_k=50,
+    ):
+        assert eos_token_id is not None and pad_token_id is not None, "eos_token_id and pad_token_id are required for inference"
+        assert llm_audio_token_id is not None and min_audio_token_id is not None, "llm_audio_token_id and min_audio_token_id are required for inference"
+        assert len(input_ids.shape) == 2 or len(input_ids.shape) == 3, f"input_ids is supposed to be [batch, seq_len] or [batch, seq_len, num_codebooks], and got {input_ids.shape}"
+
+        eos_token_id_tensor = torch.tensor([eos_token_id]).to(input_ids.device)
+
+        # keep track of which sequences are already finished
+        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
+
+        # Extend input_ids with additional num_codebooks dim
+        if len(input_ids.shape) == 2:
+            input_ids = input_ids[:, :, None].expand(1, 1, self.config.num_codebooks)
+
+        this_peer_finished = False  # used by synced_gpus only
+        max_length = input_ids.shape[1] + max_new_tokens
+
+        model_kwargs = {
+            "use_cache": True,
+            "past_key_values": DynamicCache(),
+            "attention_mask": self._prepare_attention_mask_for_generation(
+                input_ids, pad_token_id, eos_token_id
+            ),
+        }
+        # auto-regressive generation
+        while True:
+            # prepare model inputs
+            model_inputs = self._prepare_inputs_for_generation(
+                input_ids,
+                ref_embs=ref_embs,
+                ref_embs_mask=ref_embs_mask,
+                **model_kwargs
+            )
+
+            # forward pass to get next token
+            outputs = self(
+                **model_inputs,
+                return_dict=True,
+            )
+            audio_tokens = self.audio_transformer.inference(
+                outputs.last_hidden_state,
+                temperature=audio_temperature,
+                top_k=audio_top_k,
+            )
+            audio_tokens = torch.stack(
+                [
+                    audio_tokens[:, :, ci] + min_audio_token_id + ci*self.config.codebook_size
+                    for ci in range(self.config.num_codebooks)
+                ],
+                dim=2,
+            )
+
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # pre-process distribution
+            # Apply temperature and top-k
+            if temperature > 0:
+                next_token_logits = next_token_logits / temperature
+            if top_k > 0:
+                top_k = min(top_k, next_token_logits.size(-1))  # Safety check
+                # Remove all tokens with a probability less than the last token of the top-k
+                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                next_token_logits = next_token_logits.masked_fill(indices_to_remove, -float("Inf"))
+
+            # sample
+            probs = nn.functional.softmax(next_token_logits, dim=-1)
+            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+            # finished sentences should have their next token be a padding token
+            if eos_token_id is not None:
+                if pad_token_id is None:
+                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+
+            # Append NUM_CODEBOOK text tokens or audio_tokens
+            if len(next_tokens.shape) == 1:
+                next_tokens = next_tokens[:, None, None].expand(-1, 1, self.config.num_codebooks)
+            next_tokens = torch.where(next_tokens==llm_audio_token_id, audio_tokens, next_tokens)
+
+            input_ids = torch.cat([input_ids, next_tokens], dim=1)
+            if streamer is not None:
+                streamer.put(next_tokens.cpu())
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs, model_kwargs
+            )
+
+            # if eos_token was found in one sentence, set sentence to finished
+            if eos_token_id_tensor is not None:
+                unfinished_sequences = unfinished_sequences.mul(
+                    next_tokens[:, :, 0].ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=1)
+                )
+
+                # stop when each sentence is finished
+                if unfinished_sequences.max() == 0:
+                    this_peer_finished = True
+
+            # stop if we exceed the maximum length
+            if input_ids.shape[1] >= max_length:
+                this_peer_finished = True
+
+            if this_peer_finished:
+                break
+
+        if streamer is not None:
+            streamer.end()
+
+        return input_ids
+
+
+# Modified from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1103
+class VoilaAudioAlphaModel(LlamaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.pad_vocab_size_multiple = 64
+
+
+        self.ref_emb_linear = nn.Linear(256, config.hidden_size, bias=True)
+        self.audio_transformer = AudioTransformer(config, use_sdpa=False)
+
+        self.feature_extractor = FeatureExtractor()
+        self.audio_feature_extractor_adapter = FeatureExtractorAdapter(hidden_size=config.hidden_size)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        audio_labels: Optional[torch.LongTensor] = None,
+        ref_embs: Optional[List[torch.Tensor]] = None,
+        ref_embs_mask: Optional[torch.LongTensor] = None,
+        audio_datas: Optional[torch.FloatTensor] = None,
+        audio_data_masks: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        num_logits_to_keep: int = 0,
+    ) -> Union[Tuple, VoilaOutput]:
+        r"""
+        Args:
+            input_ids: [bs, seq_len, num_codebooks]
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if inputs_embeds is None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+        assert len(inputs_embeds.shape) == 4
+        if len(inputs_embeds.shape) == 4:
+            inputs_embeds = inputs_embeds.mean(dim=2)
+
+        if self.training or \
+                (past_key_values is None and ref_embs is not None) or \
+                (past_key_values is not None and past_key_values.get_seq_length() < 4 and ref_embs is not None):
+            ref_embs = self.ref_emb_linear(ref_embs.to(self.ref_emb_linear.weight.dtype))
+            ref_embs = ref_embs * ref_embs_mask.unsqueeze(-1).unsqueeze(-1)
+            # (padding_left,padding_right,padding_top,padding_bottom,padding_front,padding_back)
+            padding = (0, 0, 4, inputs_embeds.shape[1] - 5, 0, 0)
+            ref_embs = torch.nn.functional.pad(ref_embs, padding, mode='constant', value=0.0)
+            inputs_embeds = inputs_embeds + ref_embs
+
+        if self.training or audio_datas is not None:
+            audio_embeds = self.feature_extractor(audio_datas)
+            audio_embeds = self.audio_feature_extractor_adapter(audio_embeds)
+            audio_embeds = audio_embeds * audio_data_masks[..., None]
+            inputs_embeds = inputs_embeds + audio_embeds
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        if self.config.pretraining_tp > 1:
+            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
+            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
+            logits = torch.cat(logits, dim=-1)
+        else:
+            # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+            logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+
+        loss = None
+        if labels is not None:
+            # Upcast to float if we need to compute the loss to avoid potential precision issues
+            logits = logits.float()
+            # We shift tokens and labels in dataloader
+            shift_logits = logits.contiguous()
+            shift_labels = labels.contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if audio_labels is not None:
+            au_mask = (audio_labels >= 0).all(dim=-1)
+            au_hidden_states = hidden_states[au_mask]
+            au_audio_labels = audio_labels[au_mask]
+            if len(au_hidden_states) <= 0:
+                au_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+                au_audio_labels = torch.zeros_like(audio_labels).reshape(-1, self.config.num_codebooks)
+                loss_weight = 0.0
+            else:
+                loss_weight = 1.0
+            au_logits = self.audio_transformer(au_hidden_states, au_audio_labels)
+            # We shift tokens and labels in dataloader
+            shift_au_logits = au_logits.contiguous()
+            shift_audio_labels = au_audio_labels.contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_au_logits = shift_au_logits.view(-1, self.config.codebook_size)
+            shift_audio_labels = shift_audio_labels.view(-1)
+            # Enable model parallelism
+            shift_audio_labels = shift_audio_labels.to(shift_au_logits.device)
+            au_loss = loss_fct(shift_au_logits, shift_audio_labels)
+
+            loss += au_loss * loss_weight
+        else:
+            # au_tokens = self.audio_transformer.inference(hidden_states)
+            pass
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return VoilaOutput(
+            loss=loss,
+            logits=logits,
+            last_hidden_state=hidden_states,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _prepare_inputs_for_generation(
+        self, input_ids, ref_embs=None, ref_embs_mask=None, audio_datas=None, audio_data_masks=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None and past_key_values.get_seq_length() > 0:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_cache_shape()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is None and \
+                (past_key_values is None or past_key_values.get_seq_length() <= 0):
+            inputs_embeds = self.model.embed_tokens(input_ids)
+        if inputs_embeds is not None and \
+                (past_key_values is None or past_key_values.get_seq_length() <= 0):
+            model_inputs = {"inputs_embeds": inputs_embeds, "ref_embs": ref_embs, "ref_embs_mask": ref_embs_mask, "audio_datas": audio_datas, "audio_data_masks": audio_data_masks}
+        else:
+            model_inputs = {"input_ids": input_ids, "ref_embs": None, "audio_datas": None, "audio_data_masks": None}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs,
+        model_kwargs: Dict[str, Any],
+        num_new_token: int = 1,
+    ) -> Dict[str, Any]:
+        # update past_key_values
+        model_kwargs["past_key_values"] = outputs.past_key_values
+
+        # update attention mask
+        if "attention_mask" in model_kwargs:
+            attention_mask = model_kwargs["attention_mask"]
+            model_kwargs["attention_mask"] = torch.cat(
+                [attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_token))], dim=-1
+            )
+
+        return model_kwargs
+
+    def _prepare_attention_mask_for_generation(
+        self,
+        inputs: torch.Tensor,
+        pad_token_id: Optional[int],
+        eos_token_id: Optional[Union[int, List[int]]],
+    ) -> torch.LongTensor:
+        is_input_ids = len(inputs.shape) == 2 and inputs.dtype in [torch.int, torch.long]
+        is_pad_token_in_inputs = (pad_token_id is not None) and (pad_token_id in inputs)
+        if isinstance(eos_token_id, int):
+            eos_token_id = [eos_token_id]
+        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (pad_token_id not in eos_token_id)
+
+        # Check if input is input_ids and padded -> only then is attention_mask defined
+        if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id:
+            return inputs.ne(pad_token_id).long()
+        else:
+            return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device)
+
+    @torch.inference_mode()
+    def run_generate(
+        self,
+        input_ids: torch.LongTensor,
+        ref_embs: Optional[List[torch.Tensor]] = None,
+        ref_embs_mask: Optional[torch.LongTensor] = None,
+        audio_datas: Optional[torch.FloatTensor] = None,
+        audio_data_masks: Optional[torch.LongTensor] = None,
+        max_new_tokens: Optional[int] = 128,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        streamer: Optional["BaseStreamer"] = None,
+        llm_audio_token_id: Optional[int] = None,
+        min_audio_token_id: Optional[int] = None,
+        temperature=0.2,
+        top_k=50,
+        audio_temperature=0.2,
+        audio_top_k=50,
+    ):
+        assert eos_token_id is not None and pad_token_id is not None, "eos_token_id and pad_token_id are required for inference"
+        assert llm_audio_token_id is not None and min_audio_token_id is not None, "llm_audio_token_id and min_audio_token_id are required for inference"
+        assert len(input_ids.shape) == 2 or len(input_ids.shape) == 3, f"input_ids is supposed to be [batch, seq_len] or [batch, seq_len, num_codebooks], and got {input_ids.shape}"
+
+        eos_token_id_tensor = torch.tensor([eos_token_id]).to(input_ids.device)
+
+        # keep track of which sequences are already finished
+        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
+
+        # Extend input_ids with additional num_codebooks dim
+        if len(input_ids.shape) == 2:
+            input_ids = input_ids[:, :, None].expand(1, 1, self.config.num_codebooks)
+
+        this_peer_finished = False  # used by synced_gpus only
+        max_length = input_ids.shape[1] + max_new_tokens
+
+        model_kwargs = {
+            "use_cache": True,
+            "past_key_values": DynamicCache(),
+            "attention_mask": self._prepare_attention_mask_for_generation(
+                input_ids, pad_token_id, eos_token_id
+            ),
+        }
+        # auto-regressive generation
+        while True:
+            # prepare model inputs
+            model_inputs = self._prepare_inputs_for_generation(
+                input_ids,
+                ref_embs=ref_embs,
+                ref_embs_mask=ref_embs_mask,
+                audio_datas=audio_datas,
+                audio_data_masks=audio_data_masks,
+                **model_kwargs
+            )
+
+            # forward pass to get next token
+            outputs = self(
+                **model_inputs,
+                return_dict=True,
+            )
+            audio_tokens = self.audio_transformer.inference(
+                outputs.last_hidden_state,
+                temperature=audio_temperature,
+                top_k=audio_top_k,
+            )
+            audio_tokens = torch.stack(
+                [
+                    audio_tokens[:, :, ci] + min_audio_token_id + ci*self.config.codebook_size
+                    for ci in range(self.config.num_codebooks)
+                ],
+                dim=2,
+            )
+
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # pre-process distribution
+            # Apply temperature and top-k
+            if temperature > 0:
+                next_token_logits = next_token_logits / temperature
+            if top_k > 0:
+                top_k = min(top_k, next_token_logits.size(-1))  # Safety check
+                # Remove all tokens with a probability less than the last token of the top-k
+                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                next_token_logits = next_token_logits.masked_fill(indices_to_remove, -float("Inf"))
+
+            # sample
+            probs = nn.functional.softmax(next_token_logits, dim=-1)
+            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+            # finished sentences should have their next token be a padding token
+            if eos_token_id is not None:
+                if pad_token_id is None:
+                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+
+            # Append NUM_CODEBOOK text tokens or audio_tokens
+            if len(next_tokens.shape) == 1:
+                next_tokens = next_tokens[:, None, None].expand(-1, 1, self.config.num_codebooks)
+            next_tokens = torch.where(next_tokens==llm_audio_token_id, audio_tokens, next_tokens)
+
+            input_ids = torch.cat([input_ids, next_tokens], dim=1)
+            if streamer is not None:
+                streamer.put(next_tokens.cpu())
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs, model_kwargs
+            )
+
+            # if eos_token was found in one sentence, set sentence to finished
+            if eos_token_id_tensor is not None:
+                unfinished_sequences = unfinished_sequences.mul(
+                    next_tokens[:, :, 0].ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=1)
+                )
+
+                # stop when each sentence is finished
+                if unfinished_sequences.max() == 0:
+                    this_peer_finished = True
+
+            # stop if we exceed the maximum length
+            if input_ids.shape[1] >= max_length:
+                this_peer_finished = True
+
+            if this_peer_finished:
+                break
+
+        if streamer is not None:
+            streamer.end()
+
+        return input_ids
+
+
+# Modified from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1103
+class VoilaAutonomousModel(LlamaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.pad_vocab_size_multiple = 64
+
+        self.ref_emb_linear = nn.Linear(256, config.hidden_size, bias=True)
+        self.audio_transformer = AudioTransformer(config, use_sdpa=False)
+        self.voila_predictor = nn.Sequential(nn.Linear(config.hidden_size, 2, bias=True),)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        audio_labels: Optional[torch.LongTensor] = None,
+        voila_labels: Optional[torch.LongTensor] = None,
+        ref_embs: Optional[List[torch.Tensor]] = None,
+        ref_embs_mask: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        num_logits_to_keep: int = 0,
+    ) -> Union[Tuple, VoilaOutput]:
+        r"""
+        Args:
+            input_ids: [bs, seq_len, num_codebooks]
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if inputs_embeds is None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+        assert len(inputs_embeds.shape) == 4
+        if len(inputs_embeds.shape) == 4:
+            inputs_embeds = inputs_embeds.mean(dim=2)
+
+        if self.training or \
+                (past_key_values is None and ref_embs is not None) or \
+                (past_key_values is not None and past_key_values.get_seq_length() < 4 and ref_embs is not None):
+            ref_embs = self.ref_emb_linear(ref_embs.to(self.ref_emb_linear.weight.dtype))
+            ref_embs = ref_embs * ref_embs_mask.unsqueeze(-1).unsqueeze(-1)
+            # (padding_left,padding_right,padding_top,padding_bottom,padding_front,padding_back)
+            padding = (0, 0, 4, inputs_embeds.shape[1] - 5, 0, 0)
+            ref_embs = torch.nn.functional.pad(ref_embs, padding, mode='constant', value=0.0)
+            inputs_embeds = inputs_embeds + ref_embs
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        if self.config.pretraining_tp > 1:
+            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
+            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
+            logits = torch.cat(logits, dim=-1)
+        else:
+            # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+            logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+
+        # calc voila_predict_loss
+        voila_pred = self.voila_predictor(hidden_states)
+        voila_pred = voila_pred.float()
+
+        loss = None
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return VoilaOutput(
+            loss=loss,
+            logits=logits,
+            last_hidden_state=hidden_states,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            voila_pred=voila_pred,
+        )
+
+    def _prepare_inputs_for_generation(
+        self, input_ids, ref_embs=None, ref_embs_mask=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None and past_key_values.get_seq_length() > 0:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_cache_shape()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is None and \
+                (past_key_values is None or past_key_values.get_seq_length() <= 0):
+            inputs_embeds = self.model.embed_tokens(input_ids)
+        if inputs_embeds is not None and \
+                (past_key_values is None or past_key_values.get_seq_length() <= 0):
+            model_inputs = {"inputs_embeds": inputs_embeds, "ref_embs": ref_embs, "ref_embs_mask": ref_embs_mask}
+        else:
+            model_inputs = {"input_ids": input_ids, "ref_embs": None}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs,
+        model_kwargs: Dict[str, Any],
+        num_new_token: int = 1,
+    ) -> Dict[str, Any]:
+        # update past_key_values
+        model_kwargs["past_key_values"] = outputs.past_key_values
+
+        # update attention mask
+        if "attention_mask" in model_kwargs:
+            attention_mask = model_kwargs["attention_mask"]
+            model_kwargs["attention_mask"] = torch.cat(
+                [attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_token))], dim=-1
+            )
+
+        return model_kwargs
+
+    def _prepare_attention_mask_for_generation(
+        self,
+        inputs: torch.Tensor,
+        pad_token_id: Optional[int],
+        eos_token_id: Optional[Union[int, List[int]]],
+    ) -> torch.LongTensor:
+        is_input_ids = len(inputs.shape) == 2 and inputs.dtype in [torch.int, torch.long]
+        is_pad_token_in_inputs = (pad_token_id is not None) and (pad_token_id in inputs)
+        if isinstance(eos_token_id, int):
+            eos_token_id = [eos_token_id]
+        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (pad_token_id not in eos_token_id)
+
+        # Check if input is input_ids and padded -> only then is attention_mask defined
+        if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id:
+            return inputs.ne(pad_token_id).long()
+        else:
+            return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device)
+
+    @torch.inference_mode()
+    def run_generate(
+        self,
+        input_ids: torch.LongTensor,
+        input_generator,
+        ref_embs: Optional[List[torch.Tensor]] = None,
+        ref_embs_mask: Optional[torch.LongTensor] = None,
+        max_new_tokens: Optional[int] = 128,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        streamer: Optional["BaseStreamer"] = None,
+        llm_audio_token_id: Optional[int] = None,
+        min_audio_token_id: Optional[int] = None,
+        llm_assistant_token_id: Optional[int] = None,
+        temperature=0.2,
+        top_k=50,
+        audio_temperature=0.8,
+        audio_top_k=50,
+    ):
+        assert eos_token_id is not None and pad_token_id is not None, "eos_token_id and pad_token_id are required for inference"
+        assert llm_audio_token_id is not None and min_audio_token_id is not None, "llm_audio_token_id and min_audio_token_id are required for inference"
+        assert len(input_ids.shape) == 2 or len(input_ids.shape) == 3, f"input_ids is supposed to be [batch, seq_len] or [batch, seq_len, num_codebooks], and got {input_ids.shape}"
+
+        eos_token_id_tensor = torch.tensor([eos_token_id]).to(input_ids.device)
+
+        # keep track of which sequences are already finished
+        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
+
+        # Extend input_ids with additional num_codebooks dim
+        input_ids = input_ids.clone()
+        if len(input_ids.shape) == 2:
+            input_ids = input_ids[:, :, None].expand(1, 1, self.config.num_codebooks)
+
+        this_peer_finished = False  # used by synced_gpus only
+        max_length = input_ids.shape[1] + max_new_tokens
+
+        model_kwargs = {
+            "use_cache": True,
+            "past_key_values": DynamicCache(),
+            "attention_mask": self._prepare_attention_mask_for_generation(
+                input_ids, pad_token_id, eos_token_id
+            ),
+        }
+        speaking = False
+        # auto-regressive generation
+        while True:
+            # prepare model inputs
+            model_inputs = self._prepare_inputs_for_generation(
+                input_ids,
+                ref_embs=ref_embs,
+                ref_embs_mask=ref_embs_mask,
+                **model_kwargs
+            )
+
+            # forward pass to get next token
+            outputs = self(
+                **model_inputs,
+                return_dict=True,
+            )
+            audio_tokens = self.audio_transformer.inference(
+                outputs.last_hidden_state,
+                temperature=audio_temperature,
+                top_k=audio_top_k,
+            )
+            audio_tokens = torch.stack(
+                [
+                    audio_tokens[:, :, ci] + min_audio_token_id + ci*self.config.codebook_size
+                    for ci in range(self.config.num_codebooks)
+                ],
+                dim=2,
+            )
+
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # voila head output
+            voila_head_pred = outputs.voila_pred[:, -1, :]
+            voila_head_pred = torch.argmax(voila_head_pred, dim=-1)
+            voila_head_pred = voila_head_pred.cpu()[0].item()
+
+            # pre-process distribution
+            # Apply temperature and top-k
+            if temperature > 0:
+                next_token_logits = next_token_logits / temperature
+            if top_k > 0:
+                top_k = min(top_k, next_token_logits.size(-1))  # Safety check
+                # Remove all tokens with a probability less than the last token of the top-k
+                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                next_token_logits = next_token_logits.masked_fill(indices_to_remove, -float("Inf"))
+
+            # sample
+            probs = nn.functional.softmax(next_token_logits, dim=-1)
+            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+            # voila head pred == 1, use assistant token
+            if voila_head_pred == 1 and not speaking:
+                next_tokens[0] = llm_assistant_token_id
+                speaking = True
+            elif next_tokens[0] == eos_token_id:
+                speaking = False
+
+            # finished sentences should have their next token be a padding token
+            if eos_token_id is not None:
+                if pad_token_id is None:
+                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+
+            # Append NUM_CODEBOOK text tokens or audio_tokens
+            if len(next_tokens.shape) == 1:
+                next_tokens = next_tokens[:, None, None].expand(-1, 1, self.config.num_codebooks)
+            audio_token_mask = next_tokens == llm_audio_token_id
+            next_tokens = next_tokens * torch.logical_not(audio_token_mask) + audio_tokens * audio_token_mask
+
+            if audio_token_mask[0, 0, 0].item():
+                try:
+                    new_input_tokens = next(input_generator)
+                except:
+                    this_peer_finished = True
+                    break
+                new_input_tokens = new_input_tokens[None,None,:]
+            else:
+                new_input_tokens = next_tokens
+            new_input_tokens = torch.cat([new_input_tokens, next_tokens], dim=2)
+
+            input_ids = torch.cat([input_ids, new_input_tokens], dim=1)
+            if streamer is not None:
+                streamer.put(next_tokens.cpu())
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs, model_kwargs
+            )
+
+            # # if eos_token was found in one sentence, set sentence to finished
+            # if eos_token_id_tensor is not None:
+            #     unfinished_sequences = unfinished_sequences.mul(
+            #         next_tokens[:, :, 0].ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=1)
+            #     )
+
+            #     # stop when each sentence is finished
+            #     if unfinished_sequences.max() == 0:
+            #         this_peer_finished = True
+
+            # stop if we exceed the maximum length
+            if input_ids.shape[1] >= max_length:
+                this_peer_finished = True
+
+            if this_peer_finished:
+                break
+
+        if streamer is not None:
+            streamer.end()
+
+        return input_ids

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+torchaudio
+transformers
+soundfile
+librosa
+jsonlines
+gradio
+pyannote.audio

spkr.py

@@ -0,0 +1,50 @@
+import torch
+import torchaudio
+from torchaudio.functional import resample
+
+from pyannote.audio import Model
+from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
+
+
+class SpeakerEmbedding:
+    def __init__(self, model_path="pyannote/wespeaker-voxceleb-resnet34-LM", device="cuda"):
+        model = Model.from_pretrained(model_path).eval()
+
+        self.device = torch.device(device)
+        self.sample_rate = 16000
+        self.model = model.to(self.device)
+
+    @torch.no_grad()
+    def __call__(self, wav, sr):
+        wav = torch.tensor(wav, device=self.device)
+        if sr != self.sample_rate:
+            wav = resample(wav, sr, self.sample_rate)
+            sr = self.sample_rate
+
+        assert len(wav.shape) <= 3
+        is_batch = False
+        if len(wav.shape) == 3:
+            is_batch = True
+        elif len(wav.shape) == 2:
+            wav = wav[None, :, :]
+        else:
+            wav = wav[None, None, :]
+
+        with torch.inference_mode():
+            embeddings = self.model(wav)
+
+        if is_batch:
+            return embeddings
+        else:
+            return embeddings[0]
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--wav", type=str, required=True)
+    args = parser.parse_args()
+
+    model = SpeakerEmbedding(device="cuda")
+
+    wav, sr = torchaudio.load(args.wav)
+    print(model(wav, sr))

tokenize_func.py

@@ -0,0 +1,443 @@
+import io
+import copy
+import librosa
+import numpy as np
+
+
+AUDIO_TOKEN_FORMAT = "<|{}|>"
+
+DEFAULT_SYSTEM_START_TOKEN = "<SYSTEM>"
+DEFAULT_SYSTEM_END_TOKEN   = "</SYSTEM>"
+
+DEFAULT_TTS_REF_START_TOKEN = "<au_tts_ref_start>"
+DEFAULT_TTS_REF_END_TOKEN   = "<au_tts_ref_end>"
+DEFAULT_TTS_REF_TOKEN = "<au_tts_ref>"
+
+DEFAULT_CHAT_REF_START_TOKEN = "<au_chat_ref_start>"
+DEFAULT_CHAT_REF_END_TOKEN   = "<au_chat_ref_end>"
+DEFAULT_CHAT_REF_TOKEN = "<au_chat_ref>"
+
+DEFAULT_HUMAN_TOKEN = "<|HUMAN|>"
+DEFAULT_ASSISTANT_TOKEN = "<|VOILA|>"
+
+DEFAULT_AUDIO_TOKEN = "<au_token>"
+
+# ===================================
+# task special token
+# -----------------------------------
+TASK_ASR_TOKEN = "<asr>"
+TASK_TTS_TOKEN = "<tts>"
+TASK_CHAT_TOKEN = "<chat>"
+TASK_STREAM_CHAT_TOKEN = "<stream_chat>"
+
+TASK_ASR_TEXT_OUTPUT = "<asr_text_output>"
+TASK_TTS_AUDIO_OUTPUT = "<tts_audio_output>"
+TASK_CHAT_TEXT_OUTPUT = "<chat_text_output>"
+TASK_CHAT_AUDIO_OUTPUT = "<chat_audio_output>"
+
+CHAT_AUDIO_TEXT_SPLIT_TOKEN = "<chat_audio_text_split>"
+# ===================================
+
+PREPEND_LEN = 80
+SEG_LEN = 640
+AUDIO_SR = 16000
+
+TASK_TYPE_CONF = {
+    "chat_asr": TASK_ASR_TOKEN + TASK_ASR_TEXT_OUTPUT,
+    "chat_tts": TASK_TTS_TOKEN + TASK_TTS_AUDIO_OUTPUT,
+    "chat_tito": TASK_CHAT_TOKEN + TASK_CHAT_TEXT_OUTPUT,
+    "chat_tiao": TASK_CHAT_TOKEN + TASK_CHAT_AUDIO_OUTPUT,
+    "chat_aiao": TASK_CHAT_TOKEN + TASK_CHAT_AUDIO_OUTPUT,
+    "chat_atiao": TASK_CHAT_TOKEN + TASK_CHAT_AUDIO_OUTPUT,
+    "chat_aiao_auto": TASK_STREAM_CHAT_TOKEN + TASK_CHAT_AUDIO_OUTPUT,
+}
+
+
+def _get_zero_audio_pad(token_num):
+    return np.zeros(SEG_LEN*token_num)
+
+def _wrapper_audio_tokens(audio_tokens, num_codebooks, codebook_size):
+    ret_audio_tokens = []
+    for n in range(num_codebooks):
+        audio_token = audio_tokens[n]
+        ret_audio_tokens.append(''.join([AUDIO_TOKEN_FORMAT.format(au + n*codebook_size) if isinstance(au, int) else au for au in audio_token]))
+    return ret_audio_tokens
+
+def _wrapper_audio_tokens_autonomous(audio_tokens, num_codebooks, codebook_size, audio_token_min_id):
+    ret_audio_tokens = []
+    for n in range(num_codebooks):
+        audio_token = audio_tokens[n]
+        ret_audio_tokens.append([(au + n*codebook_size + audio_token_min_id) for au in audio_token])
+    return ret_audio_tokens
+
+# Item format
+# {
+#   "instruction": "",
+#   "conversations": [
+#     {
+#       "from": "user" or "assistant",
+#       "text": "",
+#       "audio": {
+#         "array": [],
+#         "sr": 16000,
+#         "bytes": "",
+#         "file": "",
+#       },
+#     }
+#   ],
+# }
+def _token_input_format(item, tokenizer, tokenizer_voila, dataset_cfg):
+    task_type = dataset_cfg["task_type"]
+    num_codebooks = dataset_cfg["num_codebooks"]
+    codebook_size = dataset_cfg["codebook_size"]
+
+    task_token = TASK_TYPE_CONF[task_type]
+
+    # Construct system message
+    system = item["instruction"]
+    if task_type in ["chat_aiao", "chat_atiao", "chat_tiao"]:
+        system = DEFAULT_CHAT_REF_START_TOKEN + DEFAULT_CHAT_REF_TOKEN + DEFAULT_CHAT_REF_END_TOKEN + system
+    elif task_type == "chat_tts":
+        system = DEFAULT_TTS_REF_START_TOKEN + DEFAULT_TTS_REF_TOKEN + DEFAULT_TTS_REF_END_TOKEN + system
+    else:
+        print (f"task type {task_type} do not use ref.")
+    system = task_token + system
+    system = DEFAULT_SYSTEM_START_TOKEN + system + DEFAULT_SYSTEM_END_TOKEN
+
+    # Get ids for system
+    system_ids = tokenizer.encode(system, add_special_tokens=False)
+
+    # Copy into num_codebooks input ids
+    input_ids_list = []
+    for _ in range(num_codebooks):
+        input_ids_list.append(copy.deepcopy(system_ids))
+
+    # Assemble conversations
+    for i, turn in enumerate(item["conversations"]):
+        if turn['from'] == 'assistant':
+            # task with audio token as input, prepare audio token
+            if task_type in ["chat_aiao", "chat_tts"]:
+                if "audio" not in turn:
+                    content = DEFAULT_ASSISTANT_TOKEN
+                    content_ids = tokenizer.encode(content, add_special_tokens=False)
+                    for n in range(num_codebooks):
+                        input_ids_list[n] += copy.deepcopy(content_ids)
+                else:
+                    # Load audio
+                    if 'array' in turn['audio']:
+                        assert "sr" in turn["audio"]
+                        if len(turn["audio"]['array'].shape) > 1:
+                            assert turn["audio"]['array'].shape[0] <= 2
+                            turn["audio"]['array'] = librosa.to_mono(turn["audio"]['array'])
+                        audio = librosa.resample(turn["audio"]['array'], orig_sr=turn["audio"]["sr"], target_sr=AUDIO_SR)
+                    elif "bytes" in turn['audio']:
+                        audio, _ = librosa.load(io.BytesIO(turn["audio"]['bytes']), sr=AUDIO_SR)
+                    elif "file" in turn['audio']:
+                        audio, _ = librosa.load(turn["audio"]['file'], sr=AUDIO_SR)
+                    else:
+                        raise Exception(f"No audio input for task {task_type}")
+
+                    # get audio token
+                    audio_tokens = tokenizer_voila.encode(audio, sr=AUDIO_SR)
+                    audio_tokens = audio_tokens.cpu().numpy().tolist()
+                    audio_tokens = _wrapper_audio_tokens(audio_tokens, num_codebooks, codebook_size)
+
+                    for n in range(num_codebooks):
+                        content = DEFAULT_ASSISTANT_TOKEN + audio_tokens[n] + tokenizer.eos_token
+                        content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                                max_length=tokenizer.model_max_length)
+                        input_ids_list[n] += content_ids
+
+            elif task_type in ["chat_tito", "chat_asr"]:
+                if "text" not in turn:
+                    content = DEFAULT_ASSISTANT_TOKEN
+                    content_ids = tokenizer.encode(content, add_special_tokens=False)
+                    for n in range(num_codebooks):
+                        input_ids_list[n] += copy.deepcopy(content_ids)
+                else:
+                    text = turn['text'].strip()
+                    content = DEFAULT_ASSISTANT_TOKEN + text + tokenizer.eos_token
+                    content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                                    max_length=tokenizer.model_max_length)
+                    for n in range(num_codebooks):
+                        input_ids_list[n] += copy.deepcopy(content_ids)
+            else:
+                raise ValueError (f"[Error] Invalid data type of {task_type}.")
+        else:
+            # task with audio token as input, prepare audio token
+            if task_type in ["chat_aiao", "chat_asr"]:
+                # Load audio
+                assert "audio" in turn
+                if 'array' in turn['audio']:
+                    assert "sr" in turn["audio"]
+                    if len(turn["audio"]['array'].shape) > 1:
+                        assert turn["audio"]['array'].shape[0] <= 2
+                        turn["audio"]['array'] = librosa.to_mono(turn["audio"]['array'])
+                    audio = librosa.resample(turn["audio"]['array'], orig_sr=turn["audio"]["sr"], target_sr=AUDIO_SR)
+                elif "bytes" in turn['audio']:
+                    audio, _ = librosa.load(io.BytesIO(turn["audio"]['bytes']), sr=AUDIO_SR)
+                elif "file" in turn['audio']:
+                    audio, _ = librosa.load(turn["audio"]['file'], sr=AUDIO_SR)
+                else:
+                    raise Exception(f"No audio input for task {task_type}")
+
+                # get audio token
+                audio_tokens = tokenizer_voila.encode(audio, sr=AUDIO_SR)
+                audio_tokens = audio_tokens.cpu().numpy().tolist()
+                audio_tokens = _wrapper_audio_tokens(audio_tokens, num_codebooks, codebook_size)
+
+                for n in range(num_codebooks):
+                    content = DEFAULT_HUMAN_TOKEN + audio_tokens[n]
+                    content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                            max_length=tokenizer.model_max_length)
+                    input_ids_list[n] += copy.deepcopy(content_ids)
+            elif task_type in ["chat_tito", "chat_tts"]:
+                text = turn['text'].strip()
+                content = DEFAULT_HUMAN_TOKEN + text
+                content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                                max_length=tokenizer.model_max_length)
+                for n in range(num_codebooks):
+                    input_ids_list[n] += copy.deepcopy(content_ids)
+            else:
+                raise ValueError (f"[Error] Invalid data type of {task_type}.")
+
+    for n in range(num_codebooks):
+        input_ids_list[n] = input_ids_list[n][:tokenizer.model_max_length]
+
+    return input_ids_list, None, None, None
+
+def _token_input_format_autonomous(item, tokenizer, tokenizer_voila, dataset_cfg):
+    task_type = dataset_cfg["task_type"]
+    num_codebooks = dataset_cfg["num_codebooks"]
+    codebook_size = dataset_cfg["codebook_size"]
+    assert task_type == "chat_aiao_auto", f"only support chat_aiao_auto, {task_type} is invalid"
+
+    DEFAULT_HUMAN_TOKEN_ID = tokenizer.convert_tokens_to_ids(DEFAULT_HUMAN_TOKEN)
+    assert isinstance(DEFAULT_HUMAN_TOKEN_ID, int), "DEFAULT_HUMAN_TOKEN_ID should be an integer"
+    AUDIO_MIN_TOKEN_ID = tokenizer.convert_tokens_to_ids(AUDIO_TOKEN_FORMAT.format(0))
+    assert isinstance(AUDIO_MIN_TOKEN_ID, int), "AUDIO_MIN_TOKEN_ID should be an integer"
+
+    task_token = TASK_TYPE_CONF[task_type]
+
+    # Construct system message
+    system = DEFAULT_CHAT_REF_START_TOKEN + DEFAULT_CHAT_REF_TOKEN + DEFAULT_CHAT_REF_END_TOKEN
+    system = task_token + system
+    system = DEFAULT_SYSTEM_START_TOKEN + system + DEFAULT_SYSTEM_END_TOKEN
+
+    # Get ids for system
+    system_ids_list = [[], []]
+    system_ids = tokenizer.encode(system, add_special_tokens=False)
+
+    # Insert instruction tokens into system prompt tokens
+    instruction = item["instruction"]
+    if instruction != "":
+        instruction_ids = tokenizer.encode(instruction, add_special_tokens=False)
+    else:
+        instruction_ids = []
+
+    system_ids_list[0] = system_ids[:-1] + instruction_ids + system_ids[-1:]
+    system_ids_list[1] = system_ids[:-1] + instruction_ids + system_ids[-1:]
+
+    # Copy into num_codebooks input ids
+    channel1_input_ids_list = [[] for _ in range(num_codebooks)]
+    channel2_input_ids_list = [[] for _ in range(num_codebooks)]
+    for n in range(num_codebooks):
+        channel1_input_ids_list[n] += copy.deepcopy(system_ids_list[0]) + [DEFAULT_HUMAN_TOKEN_ID]
+        channel2_input_ids_list[n] += copy.deepcopy(system_ids_list[1]) + [DEFAULT_HUMAN_TOKEN_ID]
+
+    # prepare audio token to simulate streaming input
+    audio_meta = item['conversations'][0]['audio']
+    if 'array' in audio_meta:
+        assert "sr" in audio_meta
+        if len(audio_meta['array'].shape) > 1:
+            assert audio_meta['array'].shape[0] <= 2
+            audio_meta['array'] = librosa.to_mono(audio_meta['array'])
+        audio = librosa.resample(audio_meta['array'], orig_sr=audio_meta["sr"], target_sr=AUDIO_SR)
+    elif "bytes" in audio_meta:
+        audio, _ = librosa.load(io.BytesIO(audio_meta['bytes']), sr=AUDIO_SR)
+    elif "file" in audio_meta:
+        audio, _ = librosa.load(audio_meta['file'], sr=AUDIO_SR)
+    else:
+        raise Exception(f"No audio input for task {task_type}")
+
+    # get audio token
+    streaming_user_input_audio_tokens = tokenizer_voila.encode(audio, sr=AUDIO_SR)
+    streaming_user_input_audio_tokens = streaming_user_input_audio_tokens.cpu().numpy().tolist()
+    streaming_user_input_audio_tokens = _wrapper_audio_tokens_autonomous(streaming_user_input_audio_tokens, num_codebooks, codebook_size, AUDIO_MIN_TOKEN_ID)
+
+    return [channel1_input_ids_list, channel2_input_ids_list], None, None, streaming_user_input_audio_tokens
+
+def _alpha_audio_input_format(item, tokenizer, dataset_cfg):
+    task_type = dataset_cfg["task_type"]
+    num_codebooks = dataset_cfg["num_codebooks"]
+    codebook_size = dataset_cfg["codebook_size"]
+
+    task_token = TASK_TYPE_CONF[task_type]
+
+    # Construct system message
+    system = item["instruction"]
+    if task_type in ["chat_aiao", "chat_atiao", "chat_tiao"]:
+        system = DEFAULT_CHAT_REF_START_TOKEN + DEFAULT_CHAT_REF_TOKEN + DEFAULT_CHAT_REF_END_TOKEN + system
+    elif task_type == "chat_tts":
+        system = DEFAULT_TTS_REF_START_TOKEN + DEFAULT_TTS_REF_TOKEN + DEFAULT_TTS_REF_END_TOKEN + system
+    else:
+        print (f"task type {task_type} do not use ref.")
+    system = task_token + system
+    system = DEFAULT_SYSTEM_START_TOKEN + system + DEFAULT_SYSTEM_END_TOKEN
+
+    # Get ids for system
+    system_ids = tokenizer.encode(system, add_special_tokens=False)
+
+    # Copy into num_codebooks input ids
+    input_ids_list = []
+    for _ in range(num_codebooks):
+        input_ids_list.append(copy.deepcopy(system_ids))
+
+    # Construct audio data and mask
+    audio_data = [np.array([0]*PREPEND_LEN)]
+    audio_data.append(_get_zero_audio_pad(len(system_ids)))
+    audio_data_mask = [0] * len(system_ids)
+
+    # Assemble conversations
+    for i, turn in enumerate(item["conversations"]):
+        if turn['from'] == 'assistant':
+            # task with audio token as input, prepare audio token
+            if task_type in ["chat_aiao"]:
+                if "audio" not in turn:
+                    content = DEFAULT_ASSISTANT_TOKEN
+                    content_ids = tokenizer.encode(content, add_special_tokens=False)
+                    for n in range(num_codebooks):
+                        input_ids_list[n] += copy.deepcopy(content_ids)
+                    # preprocess audio_data & audio_data_mask
+                    audio_data.append(_get_zero_audio_pad(len(content_ids)))
+                    audio_data_mask += [0] * len(content_ids)
+                else:
+                    # Load audio
+                    if 'array' in turn['audio']:
+                        assert "sr" in turn["audio"]
+                        if len(turn["audio"]['array'].shape) > 1:
+                            assert turn["audio"]['array'].shape[0] <= 2
+                            turn["audio"]['array'] = librosa.to_mono(turn["audio"]['array'])
+                        audio = librosa.resample(turn["audio"]['array'], orig_sr=turn["audio"]["sr"], target_sr=AUDIO_SR)
+                    elif "bytes" in turn['audio']:
+                        audio, _ = librosa.load(io.BytesIO(turn["audio"]['bytes']), sr=AUDIO_SR)
+                    elif "file" in turn['audio']:
+                        audio, _ = librosa.load(turn["audio"]['file'], sr=AUDIO_SR)
+                    else:
+                        raise Exception(f"No audio input for task {task_type}")
+
+                    # get audio token
+                    audio_token_num = int(len(audio) / SEG_LEN)
+                    audio_token = [DEFAULT_AUDIO_TOKEN] * audio_token_num
+                    audio_token = ''.join(audio_token)
+                    audio = audio[:SEG_LEN*audio_token_num]             # trim audio
+
+                    content = DEFAULT_ASSISTANT_TOKEN + audio_token + tokenizer.eos_token
+                    content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                                max_length=tokenizer.model_max_length)
+                    for n in range(num_codebooks):
+                        input_ids_list[n] += copy.deepcopy(content_ids)
+
+                    audio_data.append(_get_zero_audio_pad(1))
+                    audio_data_mask += [0]
+                    audio_data.append(audio)
+                    audio_data_mask += [1] * audio_token_num
+                    audio_data.append(_get_zero_audio_pad(1))
+                    audio_data_mask += [0]
+            elif task_type in ["chat_tito"]:
+                if "text" not in turn:
+                    content = DEFAULT_ASSISTANT_TOKEN
+                    content_ids = tokenizer.encode(content, add_special_tokens=False)
+                    for n in range(num_codebooks):
+                        input_ids_list[n] += copy.deepcopy(content_ids)
+                    # preprocess audio_data & audio_data_mask
+                    audio_data.append(_get_zero_audio_pad(len(content_ids)))
+                    audio_data_mask += [0] * len(content_ids)
+                else:
+                    text = turn['text'].strip()
+                    content = DEFAULT_ASSISTANT_TOKEN + text + tokenizer.eos_token
+                    content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                                    max_length=tokenizer.model_max_length)
+                    for n in range(num_codebooks):
+                        input_ids_list[n] += copy.deepcopy(content_ids)
+                    audio_data.append(_get_zero_audio_pad(len(content_ids)))
+                    audio_data_mask += [0] * len(content_ids)
+            else:
+                raise ValueError (f"[Error] Invalid data type of {task_type}.")
+        else:
+            # task with audio token as input, prepare audio token
+            if task_type in ["chat_aiao"]:
+                # Load audio
+                assert "audio" in turn
+                if 'array' in turn['audio']:
+                    assert "sr" in turn["audio"]
+                    if len(turn["audio"]['array'].shape) > 1:
+                        assert turn["audio"]['array'].shape[0] <= 2
+                        turn["audio"]['array'] = librosa.to_mono(turn["audio"]['array'])
+                    audio = librosa.resample(turn["audio"]['array'], orig_sr=turn["audio"]["sr"], target_sr=AUDIO_SR)
+                elif "bytes" in turn['audio']:
+                    audio, _ = librosa.load(io.BytesIO(turn["audio"]['bytes']), sr=AUDIO_SR)
+                elif "file" in turn['audio']:
+                    audio, _ = librosa.load(turn["audio"]['file'], sr=AUDIO_SR)
+                else:
+                    raise Exception(f"No audio input for task {task_type}")
+
+                # get audio token
+                audio_token_num = int(len(audio) / SEG_LEN)
+                audio_token = [DEFAULT_AUDIO_TOKEN] * audio_token_num
+                audio_token = ''.join(audio_token)
+                audio = audio[:SEG_LEN*audio_token_num]             # trim audio
+
+                content = DEFAULT_HUMAN_TOKEN + audio_token
+                content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                            max_length=tokenizer.model_max_length)
+                for n in range(num_codebooks):
+                    input_ids_list[n] += copy.deepcopy(content_ids)
+
+                audio_data.append(_get_zero_audio_pad(1))
+                audio_data_mask += [0]
+                audio_data.append(audio)
+                audio_data_mask += [1] * audio_token_num
+            elif task_type in ["chat_tito"]:
+                text = turn['text'].strip()
+                content = DEFAULT_HUMAN_TOKEN + text
+                content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
+                                                max_length=tokenizer.model_max_length)
+                for n in range(num_codebooks):
+                    input_ids_list[n] += copy.deepcopy(content_ids)
+                audio_data.append(_get_zero_audio_pad(len(content_ids)))
+                audio_data_mask += [0] * len(content_ids)
+            else:
+                raise ValueError (f"[Error] Invalid data type of {task_type}.")
+
+    for n in range(num_codebooks):
+        input_ids_list[n] = input_ids_list[n][:tokenizer.model_max_length]
+    audio_data_mask = audio_data_mask[:tokenizer.model_max_length]
+    audio_data = np.concatenate(audio_data)
+    audio_data = audio_data[:PREPEND_LEN + tokenizer.model_max_length*SEG_LEN]
+
+    return input_ids_list, audio_data, audio_data_mask, None
+
+# Item format
+# {
+#   "instruction": "",
+#   "conversations": [
+#     {
+#       "from": "user" or "assistant",
+#       "text": "",
+#       "audio": {
+#         "array": [],
+#         "sr": 16000,
+#         "bytes": "",
+#         "file": "",
+#       },
+#     }
+#   ],
+# }
+def voila_input_format(item, tokenizer, tokenizer_voila, dataset_cfg):
+    if dataset_cfg["input_type"] == "audio":
+        return _alpha_audio_input_format(item, tokenizer, dataset_cfg)
+    elif dataset_cfg["input_type"] == "autonomous":
+        return _token_input_format_autonomous(item, tokenizer, tokenizer_voila, dataset_cfg)
+    else:
+        return _token_input_format(item, tokenizer, tokenizer_voila, dataset_cfg)

voila_tokenizer.py

@@ -0,0 +1,71 @@
+import torch
+import torchaudio
+from torchaudio.functional import resample
+
+from transformers import AutoProcessor, EncodecModel
+
+
+ALL_BANDWIDTHS = [1.1]
+
+class VoilaTokenizer:
+    def __init__(
+        self,
+        model_path="maitrix-org/Voila-Tokenizer",
+        bandwidth_id=0,
+        device="cpu",
+    ):
+        self.device = torch.device(device)
+        self.bandwidth = ALL_BANDWIDTHS[bandwidth_id]
+        self.bandwidth_id = torch.tensor([bandwidth_id], device=device)
+
+        self.processor = AutoProcessor.from_pretrained(model_path)
+        self.model = EncodecModel.from_pretrained(model_path).to(device)
+
+        self.sampling_rate = self.processor.sampling_rate
+        self.model_version = self.model.config.model_version
+
+
+    @torch.no_grad()
+    def encode(self, wav, sr):
+        wav = torch.tensor(wav, dtype=torch.float32, device=self.device)
+        if sr != self.processor.sampling_rate:
+            wav = resample(wav, sr, self.processor.sampling_rate)
+            sr = self.processor.sampling_rate
+        if len(wav.shape) == 1:
+            wav = wav[None, None, :]
+        elif len(wav.shape) == 2:
+            assert wav.shape[0] == 1
+            wav = wav[None, :]
+        elif len(wav.shape) == 3:
+            assert wav.shape[0] == 1 and wav.shape[1] == 1
+
+        # inputs = self.processor(raw_audio=wav, sampling_rate=sr, return_tensors="pt")
+        encoder_outputs = self.model.encode(wav, bandwidth=self.bandwidth)
+        return encoder_outputs.audio_codes[0, 0]
+
+    @torch.no_grad()
+    def decode(self, audio_codes):
+        assert len(audio_codes.shape) == 2
+        audio_values = self.model.decode(audio_codes[None, None, :, :], [None])[0]
+        return audio_values[0, 0]
+
+if __name__ == '__main__':
+    import argparse
+    import soundfile as sf
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--wav", type=str)
+    args = parser.parse_args()
+
+    wav, sr = torchaudio.load(args.wav)
+    if len(wav.shape) > 1:
+        wav = wav[0]
+
+    model = VoilaTokenizer(device="cuda")
+
+    audio_codes = model.encode(wav, sr)
+    audio_values = model.decode(audio_codes).cpu().numpy()
+
+    tps = audio_codes.shape[-1] / (audio_values.shape[-1] / model.processor.sampling_rate)
+    print(audio_codes.shape, audio_values.shape, tps)
+    sf.write("audio_mt.wav", audio_values, model.processor.sampling_rate)