bark/mode_load.py

from bark.model_fine import FineGPT, FineGPTConfig
from bark.model import GPT, GPTConfig
from huggingface_hub import hf_hub_download
from typing import Type, TypeVar
from transformers import BertTokenizer
from scipy.special import softmax
from encodec import EncodecModel
import logging
import torch
import os
import torch.nn as nn
import torch.nn.functional as F
import re
import numpy as np
import contextlib
import funcy
import tqdm
import numpy as np
logger = logging.getLogger('bark_model_load')


def set_bark_model_load_logger(l):
    global logger
    logger = l



if (
        torch.cuda.is_available() and
        hasattr(torch.cuda, "amp") and
        hasattr(torch.cuda.amp, "autocast") and
        hasattr(torch.cuda, "is_bf16_supported") and
        torch.cuda.is_bf16_supported()
):
    autocast = funcy.partial(torch.cuda.amp.autocast, dtype=torch.bfloat16)
else:
    @contextlib.contextmanager
    def autocast():
        yield


class InferenceContext:
    def __init__(self, benchmark=False):
        # we can't expect inputs to be the same length, so disable benchmarking by default
        self._chosen_cudnn_benchmark = benchmark
        self._cudnn_benchmark = None

    def __enter__(self):
        self._cudnn_benchmark = torch.backends.cudnn.benchmark
        torch.backends.cudnn.benchmark = self._chosen_cudnn_benchmark

    def __exit__(self, exc_type, exc_value, exc_traceback):
        torch.backends.cudnn.benchmark = self._cudnn_benchmark


@contextlib.contextmanager
def _inference_mode():
    with InferenceContext(), torch.inference_mode(), torch.no_grad(), autocast():
        yield


# 定义一个泛型类型变量
G = TypeVar('G', bound='GPT')
C = TypeVar('C', bound='GPTConfig')


class ModelType:
    def __init__(self, model_type: str, model_class: Type[G], model_config: Type[C]):
        """
        模型类型适配
        :param model_type:
        :param model_class:
        :param model_config:
        """
        self.model_type = model_type
        self.model_class = model_class
        self.model_config = model_config


CONTEXT_WINDOW_SIZE = 1024

SEMANTIC_RATE_HZ = 49.9
SEMANTIC_VOCAB_SIZE = 10_000

CODEBOOK_SIZE = 1024
N_COARSE_CODEBOOKS = 2
N_FINE_CODEBOOKS = 8
COARSE_RATE_HZ = 75

SAMPLE_RATE = 24_000

REMOTE_MODEL_PATHS = {
    "text": {
        "repo_id": "suno/bark",
        "file_name": "text_2.pt",
    },
    "coarse": {
        "repo_id": "suno/bark",
        "file_name": "coarse_2.pt",
    },
    "fine": {
        "repo_id": "suno/bark",
        "file_name": "fine_2.pt",
    },
}
SUPPORTED_LANGS = [
    ("English", "en"),
    ("German", "de"),
    ("Spanish", "es"),
    ("French", "fr"),
    ("Hindi", "hi"),
    ("Italian", "it"),
    ("Japanese", "ja"),
    ("Korean", "ko"),
    ("Polish", "pl"),
    ("Portuguese", "pt"),
    ("Russian", "ru"),
    ("Turkish", "tr"),
    ("Chinese", "zh"),
]
# init prompt_speaker_np
ALLOWED_PROMPTS = {"announcer"}
for _, lang in SUPPORTED_LANGS:
    for prefix in ("", f"v2{os.path.sep}"):
        for n in range(10):
            ALLOWED_PROMPTS.add(f"{prefix}{lang}_speaker_{n}")

TEXT_ENCODING_OFFSET = 10_048
SEMANTIC_PAD_TOKEN = 10_000
TEXT_PAD_TOKEN = 129_595
SEMANTIC_INFER_TOKEN = 129_599

COARSE_SEMANTIC_PAD_TOKEN = 12_048
COARSE_INFER_TOKEN = 12_050


def _download(self, from_hf_path, file_name, local_dir):
    os.makedirs(local_dir, exist_ok=True)
    hf_hub_download(repo_id=from_hf_path, filename=file_name, local_dir=local_dir)


def _load_codec_model(device):
    model = EncodecModel.encodec_model_24khz()
    model.set_target_bandwidth(6.0)
    model.eval()
    model.to(device)
    return model


def _normalize_whitespace(text):
    return re.sub(r"\s+", " ", text).strip()


def _load_history_prompt(history_prompt_dir: str, history_prompt_input: str):
    # make sure this works on non-ubuntu

    history_prompt_input = os.path.join(*history_prompt_input.split("/"))
    if history_prompt_input not in ALLOWED_PROMPTS:
        raise ValueError("history prompt not found")
    history_prompt = np.load(
        os.path.join(history_prompt_dir, "assets", "prompts", f"{history_prompt_input}.npz")
    )
    return history_prompt


def _flatten_codebooks(arr, offset_size=CODEBOOK_SIZE):
    assert len(arr.shape) == 2
    arr = arr.copy()
    if offset_size is not None:
        for n in range(1, arr.shape[0]):
            arr[n, :] += offset_size * n
    flat_arr = arr.ravel("F")
    return flat_arr


class ModelCheckPointInfo:
    _model_type: ModelType
    _model: nn.Module
    _model_path: str

    def __init__(self, model_type: ModelType):
        self._model_type = model_type

    @property
    def model_type(self) -> ModelType:
        return self._model_type

    @property
    def model(self) -> nn.Module:
        return self._model

    @property
    def model_path(self) -> str:
        return self._model_path

    @model.setter
    def model(self, value):
        self._model = value

    @model_path.setter
    def model_path(self, value):
        self._model_path = value


class BarkModelLoader:
    _text_model: ModelCheckPointInfo = ModelCheckPointInfo(
        model_type=ModelType(model_type="text_model", model_class=GPT, model_config=GPTConfig))
    _coarse_model: ModelCheckPointInfo = ModelCheckPointInfo(
        model_type=ModelType(model_type="coarse_model", model_class=GPT, model_config=GPTConfig))
    _fine_model: ModelCheckPointInfo = ModelCheckPointInfo(
        model_type=ModelType(model_type="fine_model", model_class=FineGPT, model_config=FineGPTConfig))

    _tokenizer: BertTokenizer
    _tokenizer_path: str = "bert-base-multilingual-cased"
    _encodec: EncodecModel

    def __init__(self, tokenizer_path: str, text_path: str, coarse_path: str, fine_path: str, device: str):

        if tokenizer_path:
            self._tokenizer_path = tokenizer_path
        logger.info(f"BertTokenizer load.")
        self._tokenizer = BertTokenizer.from_pretrained(self._tokenizer_path)
        logger.info(f"BertTokenizer loaded")

        self._text_model.model_path = text_path
        self._coarse_model.model_path = coarse_path
        self._fine_model.model_path = fine_path

        # if not os.path.exists(self._text_model.model_path):
        #     model_info = REMOTE_MODEL_PATHS['text']
        #     logger.info(f"text model not found, downloading into `{self._text_model.model_path}`.")
        #
        #     _download(model_info["repo_id"], model_info["file_name"], self._text_model.model_path)
        # if not os.path.exists(self._coarse_model.model_path):
        #     model_info = REMOTE_MODEL_PATHS['coarse']
        #     logger.info(f"coarse model not found, downloading into `{self._coarse_model.model_path}`.")
        #     _download(model_info["repo_id"], model_info["file_name"], self._coarse_model.model_path)
        # if not os.path.exists(self._fine_model.model_path):
        #     model_info = REMOTE_MODEL_PATHS['fine']
        #     logger.info(f"fine model not found, downloading into `{self._fine_model.model_path}`.")
        #     _download(model_info["repo_id"], model_info["file_name"], self._fine_model.model_path)

        self._load_moad(model_type=self._text_model.model_type, ckpt_path=self._text_model.model_path, device=device)
        self._load_moad(model_type=self._coarse_model.model_type, ckpt_path=self._coarse_model.model_path, device=device)
        self._load_moad(model_type=self._fine_model.model_type, ckpt_path=self._fine_model.model_path, device=device)

    def _load_moad(self, model_type: ModelType, ckpt_path: str, device: str):
        if not os.path.exists(self._fine_model.model_path):
            raise RuntimeError("loader model path is not exists")

        device = torch.device(device)
        checkpoint = torch.load(ckpt_path, map_location=device)
        # this is a hack
        model_args = checkpoint["model_args"]
        if "input_vocab_size" not in model_args:
            model_args["input_vocab_size"] = model_args["vocab_size"]
            model_args["output_vocab_size"] = model_args["vocab_size"]
            del model_args["vocab_size"]
        gptconf = model_type.model_config(**checkpoint["model_args"])
        model = model_type.model_class(gptconf)
        state_dict = checkpoint["model"]
        # fixup checkpoint
        unwanted_prefix = "_orig_mod."
        for k, v in list(state_dict.items()):
            if k.startswith(unwanted_prefix):
                state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
        extra_keys = set(state_dict.keys()) - set(model.state_dict().keys())
        extra_keys = set([k for k in extra_keys if not k.endswith(".attn.bias")])
        missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
        missing_keys = set([k for k in missing_keys if not k.endswith(".attn.bias")])
        if len(extra_keys) != 0:
            raise ValueError(f"extra keys found: {extra_keys}")
        if len(missing_keys) != 0:
            raise ValueError(f"missing keys: {missing_keys}")
        model.load_state_dict(state_dict, strict=False)
        n_params = model.get_num_params()
        val_loss = checkpoint["best_val_loss"].item()
        logger.info(
            f"model {model_type.model_type} loaded: {round(n_params / 1e6, 1)}M params, {round(val_loss, 3)} loss")
        model.eval()
        model.to(device)
        del checkpoint, state_dict
        if model_type.model_type == "text_model":
            self._text_model.model = model
        elif model_type.model_type == "coarse_model":
            self._coarse_model.model = model
        elif model_type.model_type == "fine_model":
            self._fine_model.model = model
        self._encodec = _load_codec_model(device)

    def generate_text_semantic(
            self,
            text,
            history_prompt=None,
            history_prompt_dir=None,
            temp=0.7,
            top_k=None,
            top_p=None,
            silent=False,
            min_eos_p=0.2,
            max_gen_duration_s=None,
            allow_early_stop=True,
            use_kv_caching=False,
    ):
        """Generate semantic tokens from text."""
        assert isinstance(text, str)
        text = _normalize_whitespace(text)
        assert len(text.strip()) > 0
        if history_prompt is not None:
            history_prompt = _load_history_prompt(history_prompt_dir=history_prompt_dir, history_prompt_input=history_prompt)
            semantic_history = history_prompt["semantic_prompt"]
            assert (
                    isinstance(semantic_history, np.ndarray)
                    and len(semantic_history.shape) == 1
                    and len(semantic_history) > 0
                    and semantic_history.min() >= 0
                    and semantic_history.max() <= SEMANTIC_VOCAB_SIZE - 1
            )
        else:
            semantic_history = None
        model = self._text_model.model
        tokenizer = self._tokenizer
        encoded_text = np.array(tokenizer.encode(text, add_special_tokens=False)) + TEXT_ENCODING_OFFSET
        # if OFFLOAD_CPU:
        #     model.to(models_devices["text"])
        device = next(model.parameters()).device
        if len(encoded_text) > 256:
            p = round((len(encoded_text) - 256) / len(encoded_text) * 100, 1)
            logger.warning(f"warning, text too long, lopping of last {p}%")
            encoded_text = encoded_text[:256]
        encoded_text = np.pad(
            encoded_text,
            (0, 256 - len(encoded_text)),
            constant_values=TEXT_PAD_TOKEN,
            mode="constant",
        )
        if semantic_history is not None:
            semantic_history = semantic_history.astype(np.int64)
            # lop off if history is too long, pad if needed
            semantic_history = semantic_history[-256:]
            semantic_history = np.pad(
                semantic_history,
                (0, 256 - len(semantic_history)),
                constant_values=SEMANTIC_PAD_TOKEN,
                mode="constant",
            )
        else:
            semantic_history = np.array([SEMANTIC_PAD_TOKEN] * 256)
        x = torch.from_numpy(
            np.hstack([
                encoded_text, semantic_history, np.array([SEMANTIC_INFER_TOKEN])
            ]).astype(np.int64)
        )[None]
        assert x.shape[1] == 256 + 256 + 1
        with _inference_mode():
            x = x.to(device)
            n_tot_steps = 768
            # custom tqdm updates since we don't know when eos will occur
            pbar = tqdm.tqdm(disable=silent, total=n_tot_steps)
            pbar_state = 0
            tot_generated_duration_s = 0
            kv_cache = None
            for n in range(n_tot_steps):
                if use_kv_caching and kv_cache is not None:
                    x_input = x[:, [-1]]
                else:
                    x_input = x
                logits, kv_cache = model(
                    x_input, merge_context=True, use_cache=use_kv_caching, past_kv=kv_cache
                )
                relevant_logits = logits[0, 0, :SEMANTIC_VOCAB_SIZE]
                if allow_early_stop:
                    relevant_logits = torch.hstack(
                        (relevant_logits, logits[0, 0, [SEMANTIC_PAD_TOKEN]])  # eos
                    )
                if top_p is not None:
                    # faster to convert to numpy
                    original_device = relevant_logits.device
                    relevant_logits = relevant_logits.detach().cpu().type(torch.float32).numpy()
                    sorted_indices = np.argsort(relevant_logits)[::-1]
                    sorted_logits = relevant_logits[sorted_indices]
                    cumulative_probs = np.cumsum(softmax(sorted_logits))
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].copy()
                    sorted_indices_to_remove[0] = False
                    relevant_logits[sorted_indices[sorted_indices_to_remove]] = -np.inf
                    relevant_logits = torch.from_numpy(relevant_logits)
                    relevant_logits = relevant_logits.to(original_device)
                if top_k is not None:
                    v, _ = torch.topk(relevant_logits, min(top_k, relevant_logits.size(-1)))
                    relevant_logits[relevant_logits < v[-1]] = -float("Inf")
                probs = F.softmax(relevant_logits / temp, dim=-1)
                item_next = torch.multinomial(probs, num_samples=1).to(torch.int32)
                if allow_early_stop and (
                        item_next == SEMANTIC_VOCAB_SIZE
                        or (min_eos_p is not None and probs[-1] >= min_eos_p)
                ):
                    # eos found, so break
                    pbar.update(n - pbar_state)
                    break
                x = torch.cat((x, item_next[None]), dim=1)
                tot_generated_duration_s += 1 / SEMANTIC_RATE_HZ
                if max_gen_duration_s is not None and tot_generated_duration_s > max_gen_duration_s:
                    pbar.update(n - pbar_state)
                    break
                if n == n_tot_steps - 1:
                    pbar.update(n - pbar_state)
                    break
                del logits, relevant_logits, probs, item_next

                if n > pbar_state:
                    if n > pbar.total:
                        pbar.total = n
                    pbar.update(n - pbar_state)
                pbar_state = n
            pbar.total = n
            pbar.refresh()
            pbar.close()
            out = x.detach().cpu().numpy().squeeze()[256 + 256 + 1:]

        assert all(0 <= out) and all(out < SEMANTIC_VOCAB_SIZE)

        return out

    def generate_coarse(
            self,
            x_semantic,
            history_prompt=None,
            history_prompt_dir=None,
            temp=0.7,
            top_k=None,
            top_p=None,
            silent=False,
            max_coarse_history=630,  # min 60 (faster), max 630 (more context)
            sliding_window_len=60,
            use_kv_caching=False,
    ):
        """Generate coarse audio codes from semantic tokens."""
        assert (
                isinstance(x_semantic, np.ndarray)
                and len(x_semantic.shape) == 1
                and len(x_semantic) > 0
                and x_semantic.min() >= 0
                and x_semantic.max() <= SEMANTIC_VOCAB_SIZE - 1
        )
        assert 60 <= max_coarse_history <= 630
        assert max_coarse_history + sliding_window_len <= 1024 - 256
        semantic_to_coarse_ratio = COARSE_RATE_HZ / SEMANTIC_RATE_HZ * N_COARSE_CODEBOOKS
        max_semantic_history = int(np.floor(max_coarse_history / semantic_to_coarse_ratio))
        if history_prompt is not None:
            history_prompt = _load_history_prompt(history_prompt_dir=history_prompt_dir,history_prompt_input=history_prompt)
            x_semantic_history = history_prompt["semantic_prompt"]
            x_coarse_history = history_prompt["coarse_prompt"]
            assert (
                    isinstance(x_semantic_history, np.ndarray)
                    and len(x_semantic_history.shape) == 1
                    and len(x_semantic_history) > 0
                    and x_semantic_history.min() >= 0
                    and x_semantic_history.max() <= SEMANTIC_VOCAB_SIZE - 1
                    and isinstance(x_coarse_history, np.ndarray)
                    and len(x_coarse_history.shape) == 2
                    and x_coarse_history.shape[0] == N_COARSE_CODEBOOKS
                    and x_coarse_history.shape[-1] >= 0
                    and x_coarse_history.min() >= 0
                    and x_coarse_history.max() <= CODEBOOK_SIZE - 1
                    and (
                            round(x_coarse_history.shape[-1] / len(x_semantic_history), 1)
                            == round(semantic_to_coarse_ratio / N_COARSE_CODEBOOKS, 1)
                    )
            )
            x_coarse_history = _flatten_codebooks(x_coarse_history) + SEMANTIC_VOCAB_SIZE
            # trim histories correctly
            n_semantic_hist_provided = np.min(
                [
                    max_semantic_history,
                    len(x_semantic_history) - len(x_semantic_history) % 2,
                    int(np.floor(len(x_coarse_history) / semantic_to_coarse_ratio)),
                ]
            )
            n_coarse_hist_provided = int(round(n_semantic_hist_provided * semantic_to_coarse_ratio))
            x_semantic_history = x_semantic_history[-n_semantic_hist_provided:].astype(np.int32)
            x_coarse_history = x_coarse_history[-n_coarse_hist_provided:].astype(np.int32)
            # TODO: bit of a hack for time alignment (sounds better)
            x_coarse_history = x_coarse_history[:-2]
        else:
            x_semantic_history = np.array([], dtype=np.int32)
            x_coarse_history = np.array([], dtype=np.int32)

        model = self._coarse_model.model
        # if OFFLOAD_CPU:
        #     model.to(models_devices["coarse"])
        device = next(model.parameters()).device
        # start loop
        n_steps = int(
            round(
                np.floor(len(x_semantic) * semantic_to_coarse_ratio / N_COARSE_CODEBOOKS)
                * N_COARSE_CODEBOOKS
            )
        )
        assert n_steps > 0 and n_steps % N_COARSE_CODEBOOKS == 0
        x_semantic = np.hstack([x_semantic_history, x_semantic]).astype(np.int32)
        x_coarse = x_coarse_history.astype(np.int32)
        base_semantic_idx = len(x_semantic_history)
        with _inference_mode():
            x_semantic_in = torch.from_numpy(x_semantic)[None].to(device)
            x_coarse_in = torch.from_numpy(x_coarse)[None].to(device)
            n_window_steps = int(np.ceil(n_steps / sliding_window_len))
            n_step = 0
            for _ in tqdm.tqdm(range(n_window_steps), total=n_window_steps, disable=silent):
                semantic_idx = base_semantic_idx + int(round(n_step / semantic_to_coarse_ratio))
                # pad from right side
                x_in = x_semantic_in[:, np.max([0, semantic_idx - max_semantic_history]):]
                x_in = x_in[:, :256]
                x_in = F.pad(
                    x_in,
                    (0, 256 - x_in.shape[-1]),
                    "constant",
                    COARSE_SEMANTIC_PAD_TOKEN,
                )
                x_in = torch.hstack(
                    [
                        x_in,
                        torch.tensor([COARSE_INFER_TOKEN])[None].to(device),
                        x_coarse_in[:, -max_coarse_history:],
                    ]
                )
                kv_cache = None
                for _ in range(sliding_window_len):
                    if n_step >= n_steps:
                        continue
                    is_major_step = n_step % N_COARSE_CODEBOOKS == 0

                    if use_kv_caching and kv_cache is not None:
                        x_input = x_in[:, [-1]]
                    else:
                        x_input = x_in

                    logits, kv_cache = model(x_input, use_cache=use_kv_caching, past_kv=kv_cache)
                    logit_start_idx = (
                            SEMANTIC_VOCAB_SIZE + (1 - int(is_major_step)) * CODEBOOK_SIZE
                    )
                    logit_end_idx = (
                            SEMANTIC_VOCAB_SIZE + (2 - int(is_major_step)) * CODEBOOK_SIZE
                    )
                    relevant_logits = logits[0, 0, logit_start_idx:logit_end_idx]
                    if top_p is not None:
                        # faster to convert to numpy
                        original_device = relevant_logits.device
                        relevant_logits = relevant_logits.detach().cpu().type(torch.float32).numpy()
                        sorted_indices = np.argsort(relevant_logits)[::-1]
                        sorted_logits = relevant_logits[sorted_indices]
                        cumulative_probs = np.cumsum(softmax(sorted_logits))
                        sorted_indices_to_remove = cumulative_probs > top_p
                        sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].copy()
                        sorted_indices_to_remove[0] = False
                        relevant_logits[sorted_indices[sorted_indices_to_remove]] = -np.inf
                        relevant_logits = torch.from_numpy(relevant_logits)
                        relevant_logits = relevant_logits.to(original_device)
                    if top_k is not None:
                        v, _ = torch.topk(relevant_logits, min(top_k, relevant_logits.size(-1)))
                        relevant_logits[relevant_logits < v[-1]] = -float("Inf")
                    probs = F.softmax(relevant_logits / temp, dim=-1)
                    item_next = torch.multinomial(probs, num_samples=1).to(torch.int32)
                    item_next += logit_start_idx
                    x_coarse_in = torch.cat((x_coarse_in, item_next[None]), dim=1)
                    x_in = torch.cat((x_in, item_next[None]), dim=1)
                    del logits, relevant_logits, probs, item_next
                    n_step += 1
                del x_in
            del x_semantic_in
        # if OFFLOAD_CPU:
        #     model.to("cpu")
        gen_coarse_arr = x_coarse_in.detach().cpu().numpy().squeeze()[len(x_coarse_history):]
        del x_coarse_in
        assert len(gen_coarse_arr) == n_steps
        gen_coarse_audio_arr = gen_coarse_arr.reshape(-1, N_COARSE_CODEBOOKS).T - SEMANTIC_VOCAB_SIZE
        for n in range(1, N_COARSE_CODEBOOKS):
            gen_coarse_audio_arr[n, :] -= n * CODEBOOK_SIZE
        # _clear_cuda_cache()
        return gen_coarse_audio_arr

    def generate_fine(
            self,
            x_coarse_gen,
            history_prompt=None,
            history_prompt_dir=None,
            temp=0.5,
            silent=True,
    ):
        """Generate full audio codes from coarse audio codes."""
        assert (
                isinstance(x_coarse_gen, np.ndarray)
                and len(x_coarse_gen.shape) == 2
                and 1 <= x_coarse_gen.shape[0] <= N_FINE_CODEBOOKS - 1
                and x_coarse_gen.shape[1] > 0
                and x_coarse_gen.min() >= 0
                and x_coarse_gen.max() <= CODEBOOK_SIZE - 1
        )
        if history_prompt is not None:
            history_prompt = _load_history_prompt(history_prompt_dir=history_prompt_dir,history_prompt_input=history_prompt)
            x_fine_history = history_prompt["fine_prompt"]
            assert (
                    isinstance(x_fine_history, np.ndarray)
                    and len(x_fine_history.shape) == 2
                    and x_fine_history.shape[0] == N_FINE_CODEBOOKS
                    and x_fine_history.shape[1] >= 0
                    and x_fine_history.min() >= 0
                    and x_fine_history.max() <= CODEBOOK_SIZE - 1
            )
        else:
            x_fine_history = None
        n_coarse = x_coarse_gen.shape[0]

        model = self._fine_model.model
        # if OFFLOAD_CPU:
        #     model.to(models_devices["fine"])
        device = next(model.parameters()).device
        # make input arr
        in_arr = np.vstack(
            [
                x_coarse_gen,
                np.zeros((N_FINE_CODEBOOKS - n_coarse, x_coarse_gen.shape[1]))
                + CODEBOOK_SIZE,  # padding
            ]
        ).astype(np.int32)
        # prepend history if available (max 512)
        if x_fine_history is not None:
            x_fine_history = x_fine_history.astype(np.int32)
            in_arr = np.hstack(
                [
                    x_fine_history[:, -512:].astype(np.int32),
                    in_arr,
                ]
            )
            n_history = x_fine_history[:, -512:].shape[1]
        else:
            n_history = 0
        n_remove_from_end = 0
        # need to pad if too short (since non-causal model)
        if in_arr.shape[1] < 1024:
            n_remove_from_end = 1024 - in_arr.shape[1]
            in_arr = np.hstack(
                [
                    in_arr,
                    np.zeros((N_FINE_CODEBOOKS, n_remove_from_end), dtype=np.int32) + CODEBOOK_SIZE,
                ]
            )
        # we can be lazy about fractional loop and just keep overwriting codebooks
        n_loops = np.max([0, int(np.ceil((x_coarse_gen.shape[1] - (1024 - n_history)) / 512))]) + 1
        with _inference_mode():
            in_arr = torch.tensor(in_arr.T).to(device)
            for n in tqdm.tqdm(range(n_loops), disable=silent):
                start_idx = np.min([n * 512, in_arr.shape[0] - 1024])
                start_fill_idx = np.min([n_history + n * 512, in_arr.shape[0] - 512])
                rel_start_fill_idx = start_fill_idx - start_idx
                in_buffer = in_arr[start_idx: start_idx + 1024, :][None]
                for nn in range(n_coarse, N_FINE_CODEBOOKS):
                    logits = model(nn, in_buffer)
                    if temp is None:
                        relevant_logits = logits[0, rel_start_fill_idx:, :CODEBOOK_SIZE]
                        codebook_preds = torch.argmax(relevant_logits, -1)
                    else:
                        relevant_logits = logits[0, :, :CODEBOOK_SIZE] / temp
                        probs = F.softmax(relevant_logits, dim=-1)
                        codebook_preds = torch.multinomial(
                            probs[rel_start_fill_idx:1024], num_samples=1
                        ).reshape(-1)
                    codebook_preds = codebook_preds.to(torch.int32)
                    in_buffer[0, rel_start_fill_idx:, nn] = codebook_preds
                    del logits, codebook_preds
                # transfer over info into model_in and convert to numpy
                for nn in range(n_coarse, N_FINE_CODEBOOKS):
                    in_arr[
                    start_fill_idx: start_fill_idx + (1024 - rel_start_fill_idx), nn
                    ] = in_buffer[0, rel_start_fill_idx:, nn]
                del in_buffer
            gen_fine_arr = in_arr.detach().cpu().numpy().squeeze().T
            del in_arr
        # if OFFLOAD_CPU:
        #     model.to("cpu")
        gen_fine_arr = gen_fine_arr[:, n_history:]
        if n_remove_from_end > 0:
            gen_fine_arr = gen_fine_arr[:, :-n_remove_from_end]
        assert gen_fine_arr.shape[-1] == x_coarse_gen.shape[-1]
        # _clear_cuda_cache()
        return gen_fine_arr

    def codec_decode(self, fine_tokens):
        """Turn quantized audio codes into audio array using encodec."""

        model = self._encodec
        device = next(model.parameters()).device
        arr = torch.from_numpy(fine_tokens)[None]
        arr = arr.to(device)
        arr = arr.transpose(0, 1)
        emb = model.quantizer.decode(arr)
        out = model.decoder(emb)
        audio_arr = out.detach().cpu().numpy().squeeze()
        del arr, emb, out

        return audio_arr


if __name__ == '__main__':
    bark_load = BarkModelLoader(tokenizer_path='/media/checkpoint/bark/bert-base-multilingual-cased/',
                                text_path='/media/checkpoint/bark/suno/bark_v0/text_2.pt',
                                coarse_path='/media/checkpoint/bark/suno/bark_v0/coarse_2.pt',
                                fine_path='/media/checkpoint/bark/suno/bark_v0/fine_2.pt')

    print(bark_load)