TTS/tts/layers/bark/inference_funcs.py

import logging
import os
import re
from glob import glob
from typing import Dict, List

import librosa
import numpy as np
import torch
import torchaudio
import tqdm
from encodec.utils import convert_audio
from scipy.special import softmax
from torch.nn import functional as F

from TTS.tts.layers.bark.hubert.hubert_manager import HubertManager
from TTS.tts.layers.bark.hubert.kmeans_hubert import CustomHubert
from TTS.tts.layers.bark.hubert.tokenizer import HubertTokenizer
from TTS.tts.layers.bark.load_model import clear_cuda_cache, inference_mode

logger = logging.getLogger(__name__)


def _tokenize(tokenizer, text):
    return tokenizer.encode(text, add_special_tokens=False)


def _detokenize(tokenizer, enc_text):
    return tokenizer.decode(enc_text)


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


def get_voices(extra_voice_dirs: List[str] = []):  # pylint: disable=dangerous-default-value
    dirs = extra_voice_dirs
    voices: Dict[str, List[str]] = {}
    for d in dirs:
        subs = os.listdir(d)
        for sub in subs:
            subj = os.path.join(d, sub)
            if os.path.isdir(subj):
                voices[sub] = list(glob(f"{subj}/*.npz"))
                # fetch audio files if no npz files are found
                if len(voices[sub]) == 0:
                    voices[sub] = list(glob(f"{subj}/*.wav")) + list(glob(f"{subj}/*.mp3"))
    return voices


def load_npz(npz_file):
    x_history = np.load(npz_file)
    semantic = x_history["semantic_prompt"]
    coarse = x_history["coarse_prompt"]
    fine = x_history["fine_prompt"]
    return semantic, coarse, fine


def load_voice(model, voice: str, extra_voice_dirs: List[str] = []):  # pylint: disable=dangerous-default-value
    if voice == "random":
        return None, None, None

    voices = get_voices(extra_voice_dirs)
    paths = voices[voice]

    # bark only uses a single sample for cloning
    if len(paths) > 1:
        raise ValueError(f"Voice {voice} has multiple paths: {paths}")

    try:
        path = voices[voice]
    except KeyError as e:
        raise KeyError(f"Voice {voice} not found in {extra_voice_dirs}") from e

    if len(paths) == 1 and paths[0].endswith(".npz"):
        return load_npz(path[0])

    audio_path = paths[0]
    # replace the file extension with .npz
    output_path = os.path.splitext(audio_path)[0] + ".npz"
    generate_voice(audio=audio_path, model=model, output_path=output_path)
    return load_voice(model, voice, extra_voice_dirs)


def zero_crossing_rate(audio, frame_length=1024, hop_length=512):
    zero_crossings = np.sum(np.abs(np.diff(np.sign(audio))) / 2)
    total_frames = 1 + int((len(audio) - frame_length) / hop_length)
    return zero_crossings / total_frames


def compute_spectral_contrast(audio_data, sample_rate, n_bands=6, fmin=200.0):
    spectral_contrast = librosa.feature.spectral_contrast(y=audio_data, sr=sample_rate, n_bands=n_bands, fmin=fmin)
    return np.mean(spectral_contrast)


def compute_average_bass_energy(audio_data, sample_rate, max_bass_freq=250):
    stft = librosa.stft(audio_data)
    power_spectrogram = np.abs(stft) ** 2
    frequencies = librosa.fft_frequencies(sr=sample_rate, n_fft=stft.shape[0])
    bass_mask = frequencies <= max_bass_freq
    bass_energy = power_spectrogram[np.ix_(bass_mask, np.arange(power_spectrogram.shape[1]))].mean()
    return bass_energy


def generate_voice(
    audio,
    model,
    output_path,
):
    """Generate a new voice from a given audio and text prompt.

    Args:
        audio (np.ndarray): The audio to use as a base for the new voice.
        text (str): Transcription of the audio you are clonning.
        model (BarkModel): The BarkModel to use for generating the new voice.
        output_path (str): The path to save the generated voice to.
    """
    if isinstance(audio, str):
        audio, sr = torchaudio.load(audio)
        audio = convert_audio(audio, sr, model.config.sample_rate, model.encodec.channels)
        audio = audio.unsqueeze(0).to(model.device)

    with torch.no_grad():
        encoded_frames = model.encodec.encode(audio)
    codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1).squeeze()  # [n_q, T]

    # move codes to cpu
    codes = codes.cpu().numpy()

    # generate semantic tokens
    # Load the HuBERT model
    hubert_manager = HubertManager()
    # hubert_manager.make_sure_hubert_installed(model_path=model.config.LOCAL_MODEL_PATHS["hubert"])
    hubert_manager.make_sure_tokenizer_installed(model_path=model.config.LOCAL_MODEL_PATHS["hubert_tokenizer"])

    hubert_model = CustomHubert(checkpoint_path=model.config.LOCAL_MODEL_PATHS["hubert"]).to(model.device)

    # Load the CustomTokenizer model
    tokenizer = HubertTokenizer.load_from_checkpoint(
        model.config.LOCAL_MODEL_PATHS["hubert_tokenizer"], map_location=model.device
    )
    # semantic_tokens = model.text_to_semantic(
    #     text, max_gen_duration_s=seconds, top_k=50, top_p=0.95, temp=0.7
    # )  # not 100%
    semantic_vectors = hubert_model.forward(audio[0], input_sample_hz=model.config.sample_rate)
    semantic_tokens = tokenizer.get_token(semantic_vectors)
    semantic_tokens = semantic_tokens.cpu().numpy()

    np.savez(output_path, fine_prompt=codes, coarse_prompt=codes[:2, :], semantic_prompt=semantic_tokens)


def generate_text_semantic(
    text,
    model,
    history_prompt=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,
    base=None,
    use_kv_caching=True,
    **kwargs,  # pylint: disable=unused-argument
):
    """Generate semantic tokens from text.

    Args:
        text (str): The text to generate semantic tokens from.
        model (BarkModel): The BarkModel to use for generating the semantic tokens.
        history_prompt (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a prompt for the generation.
        temp (float): The temperature to use for the generation.
        top_k (int): The number of top tokens to consider for the generation.
        top_p (float): The cumulative probability to consider for the generation.
        silent (bool): Whether to silence the tqdm progress bar.
        min_eos_p (float): The minimum probability to consider for the end of sentence token.
        max_gen_duration_s (float): The maximum duration in seconds to generate for.
        allow_early_stop (bool): Whether to allow the generation to stop early.
        base (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a base for the generation.
        use_kv_caching (bool): Whether to use key-value caching for the generation.
        **kwargs: Additional keyword arguments. They are ignored.

    Returns:
        np.ndarray: The generated semantic tokens.
    """
    assert isinstance(text, str)
    text = _normalize_whitespace(text)
    assert len(text.strip()) > 0
    if all(v is not None for v in history_prompt) or base is not None:
        if history_prompt is not None:
            semantic_history = history_prompt[0]
        if base is not None:
            semantic_history = base[0]
        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() <= model.config.SEMANTIC_VOCAB_SIZE - 1
        )
    else:
        semantic_history = None
    encoded_text = np.array(_tokenize(model.tokenizer, text)) + model.config.TEXT_ENCODING_OFFSET
    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=model.config.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=model.config.SEMANTIC_PAD_TOKEN,
            mode="constant",
        )
    else:
        semantic_history = np.array([model.config.SEMANTIC_PAD_TOKEN] * 256)
    x = torch.from_numpy(
        np.hstack([encoded_text, semantic_history, np.array([model.config.SEMANTIC_INFER_TOKEN])]).astype(np.int64)
    )[None]
    assert x.shape[1] == 256 + 256 + 1
    with inference_mode():
        x = x.to(model.device)
        n_tot_steps = 768
        # custom tqdm updates since we don't know when eos will occur
        pbar = tqdm.tqdm(disable=silent, total=100)
        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.semantic_model(
                x_input, merge_context=True, use_cache=use_kv_caching, past_kv=kv_cache
            )
            relevant_logits = logits[0, 0, : model.config.SEMANTIC_VOCAB_SIZE]
            if allow_early_stop:
                relevant_logits = torch.hstack(
                    (relevant_logits, logits[0, 0, [model.config.SEMANTIC_PAD_TOKEN]])
                )  # eos
            if top_p is not None:
                # faster to convert to numpy
                logits_device = relevant_logits.device
                logits_dtype = relevant_logits.type()
                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(logits_device).type(logits_dtype)
            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 = torch.softmax(relevant_logits / temp, dim=-1)
            item_next = torch.multinomial(probs, num_samples=1)
            if allow_early_stop and (
                item_next == model.config.SEMANTIC_VOCAB_SIZE or (min_eos_p is not None and probs[-1] >= min_eos_p)
            ):
                # eos found, so break
                pbar.update(100 - pbar_state)
                break
            x = torch.cat((x, item_next[None]), dim=1)
            tot_generated_duration_s += 1 / model.config.SEMANTIC_RATE_HZ
            if max_gen_duration_s is not None and tot_generated_duration_s > max_gen_duration_s:
                pbar.update(100 - pbar_state)
                break
            if n == n_tot_steps - 1:
                pbar.update(100 - pbar_state)
                break
            del logits, relevant_logits, probs, item_next
            req_pbar_state = np.min([100, int(round(100 * n / n_tot_steps))])
            if req_pbar_state > pbar_state:
                pbar.update(req_pbar_state - pbar_state)
            pbar_state = req_pbar_state
        pbar.close()
        out = x.detach().cpu().numpy().squeeze()[256 + 256 + 1 :]
    assert all(out >= 0) and all(out < model.config.SEMANTIC_VOCAB_SIZE)
    clear_cuda_cache()
    return out


def _flatten_codebooks(arr, offset_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


def generate_coarse(
    x_semantic,
    model,
    history_prompt=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,
    base=None,
    use_kv_caching=True,
):
    """Generate coarse audio codes from semantic tokens.

    Args:
        x_semantic (np.ndarray): The semantic tokens to generate coarse audio codes from.
        model (BarkModel): The BarkModel to use for generating the coarse audio codes.
        history_prompt (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a prompt for the generation.
        temp (float): The temperature to use for the generation.
        top_k (int): The number of top tokens to consider for the generation.
        top_p (float): The cumulative probability to consider for the generation.
        silent (bool): Whether to silence the tqdm progress bar.
        max_coarse_history (int): The maximum number of coarse audio codes to use as history.
        sliding_window_len (int): The length of the sliding window to use for the generation.
        base (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a base for the generation.
        use_kv_caching (bool): Whether to use key-value caching for the generation.

    Returns:
        np.ndarray: The generated coarse audio codes.
    """
    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() <= model.config.SEMANTIC_VOCAB_SIZE - 1
    )
    assert 60 <= max_coarse_history <= 630
    assert max_coarse_history + sliding_window_len <= 1024 - 256
    semantic_to_coarse_ratio = (
        model.config.COARSE_RATE_HZ / model.config.SEMANTIC_RATE_HZ * model.config.N_COARSE_CODEBOOKS
    )
    max_semantic_history = int(np.floor(max_coarse_history / semantic_to_coarse_ratio))
    if all(v is not None for v in history_prompt) or base is not None:
        if history_prompt is not None:
            x_history = history_prompt
            x_semantic_history = x_history[0]
            x_coarse_history = x_history[1]
        if base is not None:
            x_semantic_history = base[0]
            x_coarse_history = base[1]
        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() <= model.config.SEMANTIC_VOCAB_SIZE - 1
            and isinstance(x_coarse_history, np.ndarray)
            and len(x_coarse_history.shape) == 2
            and x_coarse_history.shape[0] == model.config.N_COARSE_CODEBOOKS
            and x_coarse_history.shape[-1] >= 0
            and x_coarse_history.min() >= 0
            and x_coarse_history.max() <= model.config.CODEBOOK_SIZE - 1
            and (
                round(x_coarse_history.shape[-1] / len(x_semantic_history), 1)
                == round(semantic_to_coarse_ratio / model.config.N_COARSE_CODEBOOKS, 1)
            )
        )
        x_coarse_history = (
            _flatten_codebooks(x_coarse_history, model.config.CODEBOOK_SIZE) + model.config.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)
    # start loop
    n_steps = int(
        round(
            np.floor(len(x_semantic) * semantic_to_coarse_ratio / model.config.N_COARSE_CODEBOOKS)
            * model.config.N_COARSE_CODEBOOKS
        )
    )
    assert n_steps > 0 and n_steps % model.config.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(model.device)
        x_coarse_in = torch.from_numpy(x_coarse)[None].to(model.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",
                model.config.COARSE_SEMANTIC_PAD_TOKEN,
            )
            x_in = torch.hstack(
                [
                    x_in,
                    torch.tensor([model.config.COARSE_INFER_TOKEN])[None].to(model.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 % model.config.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.coarse_model(x_input, use_cache=use_kv_caching, past_kv=kv_cache)
                logit_start_idx = (
                    model.config.SEMANTIC_VOCAB_SIZE + (1 - int(is_major_step)) * model.config.CODEBOOK_SIZE
                )
                logit_end_idx = model.config.SEMANTIC_VOCAB_SIZE + (2 - int(is_major_step)) * model.config.CODEBOOK_SIZE
                relevant_logits = logits[0, 0, logit_start_idx:logit_end_idx]
                if top_p is not None:
                    # faster to convert to numpy
                    logits_device = relevant_logits.device
                    logits_dtype = relevant_logits.type()
                    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(torch.nn.functional.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(logits_device).type(logits_dtype)
                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 = torch.nn.functional.softmax(relevant_logits / temp, dim=-1)
                item_next = torch.multinomial(probs, num_samples=1)
                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
    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, model.config.N_COARSE_CODEBOOKS).T - model.config.SEMANTIC_VOCAB_SIZE
    )
    for n in range(1, model.config.N_COARSE_CODEBOOKS):
        gen_coarse_audio_arr[n, :] -= n * model.config.CODEBOOK_SIZE
    clear_cuda_cache()
    return gen_coarse_audio_arr


def generate_fine(
    x_coarse_gen,
    model,
    history_prompt=None,
    temp=0.5,
    silent=True,
    base=None,
):
    """Generate full audio codes from coarse audio codes.

    Args:
        x_coarse_gen (np.ndarray): The coarse audio codes to generate full audio codes from.
        model (BarkModel): The BarkModel to use for generating the full audio codes.
        history_prompt (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a prompt for the generation.
        temp (float): The temperature to use for the generation.
        silent (bool): Whether to silence the tqdm progress bar.
        base (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a base for the generation.

    Returns:
        np.ndarray: The generated full audio codes.
    """
    assert (
        isinstance(x_coarse_gen, np.ndarray)
        and len(x_coarse_gen.shape) == 2
        and 1 <= x_coarse_gen.shape[0] <= model.config.N_FINE_CODEBOOKS - 1
        and x_coarse_gen.shape[1] > 0
        and x_coarse_gen.min() >= 0
        and x_coarse_gen.max() <= model.config.CODEBOOK_SIZE - 1
    )
    if all(v is not None for v in history_prompt) or base is not None:
        if history_prompt is not None:
            x_fine_history = history_prompt[2]
        if base is not None:
            x_fine_history = base[2]
        assert (
            isinstance(x_fine_history, np.ndarray)
            and len(x_fine_history.shape) == 2
            and x_fine_history.shape[0] == model.config.N_FINE_CODEBOOKS
            and x_fine_history.shape[1] >= 0
            and x_fine_history.min() >= 0
            and x_fine_history.max() <= model.config.CODEBOOK_SIZE - 1
        )
    else:
        x_fine_history = None
    n_coarse = x_coarse_gen.shape[0]
    # make input arr
    in_arr = np.vstack(
        [
            x_coarse_gen,
            np.zeros((model.config.N_FINE_CODEBOOKS - n_coarse, x_coarse_gen.shape[1]))
            + model.config.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((model.config.N_FINE_CODEBOOKS, n_remove_from_end), dtype=np.int32)
                + model.config.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(model.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, model.config.N_FINE_CODEBOOKS):
                logits = model.fine_model(nn, in_buffer)
                if temp is None:
                    relevant_logits = logits[0, rel_start_fill_idx:, : model.config.CODEBOOK_SIZE]
                    codebook_preds = torch.argmax(relevant_logits, -1)
                else:
                    relevant_logits = logits[0, :, : model.config.CODEBOOK_SIZE] / temp
                    probs = F.softmax(relevant_logits, dim=-1)
                    codebook_preds = torch.hstack(
                        [torch.multinomial(probs[n], num_samples=1) for n in range(rel_start_fill_idx, 1024)]
                    )
                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, model.config.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
    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(fine_tokens, model):
    """Turn quantized audio codes into audio array using encodec."""
    arr = torch.from_numpy(fine_tokens)[None]
    arr = arr.to(model.device)
    arr = arr.transpose(0, 1)
    emb = model.encodec.quantizer.decode(arr)
    out = model.encodec.decoder(emb)
    audio_arr = out.detach().cpu().numpy().squeeze()
    return audio_arr