Duplicate from kevinwang676/Bark-with-Voice-Cloning

Co-authored-by: Kevin Wang <[email protected]>

.flake8

@@ -0,0 +1,17 @@
+[flake8]
+ignore = E203, E266, E501, W503
+max-line-length = 100
+import-order-style = google
+application-import-names = flake8
+select = B,C,E,F,W,T4,B9
+exclude =
+  .tox,
+  .git,
+  __pycache__,
+  docs,
+  sbatch,
+  .venv,
+  *.pyc,
+  *.egg-info,
+  .cache,
+  .eggs

.gitattributes

@@ -0,0 +1,44 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+SE_checkpoint.pth.tar filter=lfs diff=lfs merge=lfs -text
+best_model.pth.tar filter=lfs diff=lfs merge=lfs -text
+nana_longest_vocal.wav filter=lfs diff=lfs merge=lfs -text
+test.wav filter=lfs diff=lfs merge=lfs -text
+reference.wav filter=lfs diff=lfs merge=lfs -text
+ref.wav filter=lfs diff=lfs merge=lfs -text
+DeepFilterNet2/checkpoints/model_96.ckpt.best filter=lfs diff=lfs merge=lfs -text
+samples/dkitchen.wav filter=lfs diff=lfs merge=lfs -text
+samples/nriver.wav filter=lfs diff=lfs merge=lfs -text
+samples/scafe.wav filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,160 @@
+__pycache__/
+/outputs
+/speakers
+.vs
+*.npz
+*.wav
+*.npy
+.vs/
+/models
+/bark_ui_enhanced.egg-info
+/build/lib/bark
+*.pth
+*.pt
+*.zip
+
+# Own stuff
+*.wav
+*.png
+*.pdf
+out/
+export/
+DeepFilterNet/poetry.lock
+gradio_cached_examples/
+
+### Rust gitignore ###
+
+# Generated by Cargo
+# will have compiled files and executables
+debug/
+target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+### Python gitignore ###
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+typings
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# IDE
+.idea

DeepFilterNet2/checkpoints/model_96.ckpt.best

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

DeepFilterNet2/config.ini

@@ -0,0 +1,108 @@
+[train]
+seed = 43
+device = 
+model = deepfilternet2
+jit = false
+mask_only = false
+df_only = false
+batch_size = 96
+batch_size_eval = 128
+num_workers = 16
+max_sample_len_s = 3.0
+p_atten_lim = 0.0
+p_reverb = 0.1
+overfit = false
+max_epochs = 100
+log_freq = 100
+log_timings = False
+validation_criteria = loss
+validation_criteria_rule = min
+early_stopping_patience = 15
+global_ds_sampling_f = 1
+num_prefetch_batches = 8
+dataloader_snrs = -5,0,5,10,20,40
+detect_anomaly = false
+batch_size_scheduling = 0/8,1/16,2/24,5/32,10/64,20/128,40/9999
+start_eval = true
+validation_set_caching = false
+
+[df]
+sr = 48000
+fft_size = 960
+hop_size = 480
+nb_erb = 32
+nb_df = 96
+norm_tau = 1
+lsnr_max = 35
+lsnr_min = -15
+min_nb_erb_freqs = 2
+pad_mode = input_specf
+
+[deepfilternet]
+conv_lookahead = 2
+conv_ch = 64
+conv_depthwise = True
+emb_hidden_dim = 256
+emb_num_layers = 3
+gru_groups = 8
+linear_groups = 8
+conv_dec_mode = transposed
+convt_depthwise = True
+mask_pf = False
+df_order = 5
+df_lookahead = 2
+df_hidden_dim = 256
+df_num_layers = 2
+dfop_method = df
+group_shuffle = False
+conv_kernel = 1,3
+df_gru_skip = none
+df_output_layer = groupedlinear
+gru_type = squeeze
+df_pathway_kernel_size_t = 5
+df_n_iter = 1
+enc_concat = True
+conv_kernel_inp = 3,3
+
+[localsnrloss]
+factor = 1e-3
+
+[maskloss]
+factor = 0
+mask = iam
+gamma = 0.6
+gamma_pred = 0.6
+f_under = 1
+
+[spectralloss]
+factor_magnitude = 1000
+factor_complex = 1000
+gamma = 0.3
+
+[dfalphaloss]
+factor = 0.0
+
+[multiresspecloss]
+factor = 500
+factor_complex = 500
+gamma = 0.3
+fft_sizes = 256,512,1024
+
+[optim]
+lr = 0.001
+momentum = 0
+weight_decay = 1e-12
+weight_decay_end = 0.05
+optimizer = adamw
+lr_min = 1e-06
+lr_warmup = 0.0001
+warmup_epochs = 3
+lr_cycle_mul = 1.0
+lr_cycle_decay = 0.5
+lr_cycle_limit = 1
+lr_update_per_epoch = False
+lr_cycle_epochs = -1
+
+[sdrloss]
+factor = 0.0
+segmental_ws = 0

README.md

@@ -0,0 +1,14 @@
+---
+title: Bark with Voice Cloning
+emoji: 📊
+colorFrom: purple
+colorTo: purple
+sdk: gradio
+sdk_version: 3.34.0
+app_file: app.py
+pinned: true
+license: mit
+duplicated_from: kevinwang676/Bark-with-Voice-Cloning
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,674 @@
+from cProfile import label
+import dataclasses
+from distutils.command.check import check
+from doctest import Example
+import gradio as gr
+import os
+import sys
+import numpy as np
+import logging
+import torch
+import pytorch_seed
+import time
+
+
+import math
+import tempfile
+from typing import Optional, Tuple, Union
+
+import matplotlib.pyplot as plt
+from loguru import logger
+from PIL import Image
+from torch import Tensor
+from torchaudio.backend.common import AudioMetaData
+
+from df import config
+from df.enhance import enhance, init_df, load_audio, save_audio
+from df.io import resample
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model, df, _ = init_df("./DeepFilterNet2", config_allow_defaults=True)
+model = model.to(device=device).eval()
+
+fig_noisy: plt.Figure
+fig_enh: plt.Figure
+ax_noisy: plt.Axes
+ax_enh: plt.Axes
+fig_noisy, ax_noisy = plt.subplots(figsize=(15.2, 4))
+fig_noisy.set_tight_layout(True)
+fig_enh, ax_enh = plt.subplots(figsize=(15.2, 4))
+fig_enh.set_tight_layout(True)
+
+NOISES = {
+    "None": None,
+    "Kitchen": "samples/dkitchen.wav",
+    "Living Room": "samples/dliving.wav",
+    "River": "samples/nriver.wav",
+    "Cafe": "samples/scafe.wav",
+}
+
+
+from xml.sax import saxutils
+from bark.api import generate_with_settings
+from bark.api import save_as_prompt
+from util.settings import Settings
+#import nltk
+
+from bark import SAMPLE_RATE
+from cloning.clonevoice import clone_voice
+from bark.generation import SAMPLE_RATE, preload_models, _load_history_prompt, codec_decode
+from scipy.io.wavfile import write as write_wav
+from util.parseinput import split_and_recombine_text, build_ssml, is_ssml, create_clips_from_ssml
+from datetime import datetime
+from tqdm.auto import tqdm
+from util.helper import create_filename, add_id3_tag
+from swap_voice import swap_voice_from_audio
+from training.training_prepare import prepare_semantics_from_text, prepare_wavs_from_semantics
+from training.train import training_prepare_files, train
+
+
+# Denoise
+
+def mix_at_snr(clean, noise, snr, eps=1e-10):
+    """Mix clean and noise signal at a given SNR.
+    Args:
+        clean: 1D Tensor with the clean signal to mix.
+        noise: 1D Tensor of shape.
+        snr: Signal to noise ratio.
+    Returns:
+        clean: 1D Tensor with gain changed according to the snr.
+        noise: 1D Tensor with the combined noise channels.
+        mix: 1D Tensor with added clean and noise signals.
+    """
+    clean = torch.as_tensor(clean).mean(0, keepdim=True)
+    noise = torch.as_tensor(noise).mean(0, keepdim=True)
+    if noise.shape[1] < clean.shape[1]:
+        noise = noise.repeat((1, int(math.ceil(clean.shape[1] / noise.shape[1]))))
+    max_start = int(noise.shape[1] - clean.shape[1])
+    start = torch.randint(0, max_start, ()).item() if max_start > 0 else 0
+    logger.debug(f"start: {start}, {clean.shape}")
+    noise = noise[:, start : start + clean.shape[1]]
+    E_speech = torch.mean(clean.pow(2)) + eps
+    E_noise = torch.mean(noise.pow(2))
+    K = torch.sqrt((E_noise / E_speech) * 10 ** (snr / 10) + eps)
+    noise = noise / K
+    mixture = clean + noise
+    logger.debug("mixture: {mixture.shape}")
+    assert torch.isfinite(mixture).all()
+    max_m = mixture.abs().max()
+    if max_m > 1:
+        logger.warning(f"Clipping detected during mixing. Reducing gain by {1/max_m}")
+        clean, noise, mixture = clean / max_m, noise / max_m, mixture / max_m
+    return clean, noise, mixture
+
+
+def load_audio_gradio(
+    audio_or_file: Union[None, str, Tuple[int, np.ndarray]], sr: int
+) -> Optional[Tuple[Tensor, AudioMetaData]]:
+    if audio_or_file is None:
+        return None
+    if isinstance(audio_or_file, str):
+        if audio_or_file.lower() == "none":
+            return None
+        # First try default format
+        audio, meta = load_audio(audio_or_file, sr)
+    else:
+        meta = AudioMetaData(-1, -1, -1, -1, "")
+        assert isinstance(audio_or_file, (tuple, list))
+        meta.sample_rate, audio_np = audio_or_file
+        # Gradio documentation says, the shape is [samples, 2], but apparently sometimes its not.
+        audio_np = audio_np.reshape(audio_np.shape[0], -1).T
+        if audio_np.dtype == np.int16:
+            audio_np = (audio_np / (1 << 15)).astype(np.float32)
+        elif audio_np.dtype == np.int32:
+            audio_np = (audio_np / (1 << 31)).astype(np.float32)
+        audio = resample(torch.from_numpy(audio_np), meta.sample_rate, sr)
+    return audio, meta
+
+
+def demo_fn(speech_upl: str, noise_type: str, snr: int, mic_input: str):
+    if mic_input:
+        speech_upl = mic_input
+    sr = config("sr", 48000, int, section="df")
+    logger.info(f"Got parameters speech_upl: {speech_upl}, noise: {noise_type}, snr: {snr}")
+    snr = int(snr)
+    noise_fn = NOISES[noise_type]
+    meta = AudioMetaData(-1, -1, -1, -1, "")
+    max_s = 1000  # limit to 10 seconds
+    if speech_upl is not None:
+        sample, meta = load_audio(speech_upl, sr)
+        max_len = max_s * sr
+        if sample.shape[-1] > max_len:
+            start = torch.randint(0, sample.shape[-1] - max_len, ()).item()
+            sample = sample[..., start : start + max_len]
+    else:
+        sample, meta = load_audio("samples/p232_013_clean.wav", sr)
+        sample = sample[..., : max_s * sr]
+    if sample.dim() > 1 and sample.shape[0] > 1:
+        assert (
+            sample.shape[1] > sample.shape[0]
+        ), f"Expecting channels first, but got {sample.shape}"
+        sample = sample.mean(dim=0, keepdim=True)
+    logger.info(f"Loaded sample with shape {sample.shape}")
+    if noise_fn is not None:
+        noise, _ = load_audio(noise_fn, sr)  # type: ignore
+        logger.info(f"Loaded noise with shape {noise.shape}")
+        _, _, sample = mix_at_snr(sample, noise, snr)
+    logger.info("Start denoising audio")
+    enhanced = enhance(model, df, sample)
+    logger.info("Denoising finished")
+    lim = torch.linspace(0.0, 1.0, int(sr * 0.15)).unsqueeze(0)
+    lim = torch.cat((lim, torch.ones(1, enhanced.shape[1] - lim.shape[1])), dim=1)
+    enhanced = enhanced * lim
+    if meta.sample_rate != sr:
+        enhanced = resample(enhanced, sr, meta.sample_rate)
+        sample = resample(sample, sr, meta.sample_rate)
+        sr = meta.sample_rate
+    enhanced_wav = tempfile.NamedTemporaryFile(suffix="enhanced.wav", delete=False).name
+    save_audio(enhanced_wav, enhanced, sr)
+    logger.info(f"saved audios: {enhanced_wav}")
+    ax_noisy.clear()
+    ax_enh.clear()
+    # noisy_wav = gr.make_waveform(noisy_fn, bar_count=200)
+    # enh_wav = gr.make_waveform(enhanced_fn, bar_count=200)
+    return enhanced_wav
+
+
+def specshow(
+    spec,
+    ax=None,
+    title=None,
+    xlabel=None,
+    ylabel=None,
+    sr=48000,
+    n_fft=None,
+    hop=None,
+    t=None,
+    f=None,
+    vmin=-100,
+    vmax=0,
+    xlim=None,
+    ylim=None,
+    cmap="inferno",
+):
+    """Plots a spectrogram of shape [F, T]"""
+    spec_np = spec.cpu().numpy() if isinstance(spec, torch.Tensor) else spec
+    if ax is not None:
+        set_title = ax.set_title
+        set_xlabel = ax.set_xlabel
+        set_ylabel = ax.set_ylabel
+        set_xlim = ax.set_xlim
+        set_ylim = ax.set_ylim
+    else:
+        ax = plt
+        set_title = plt.title
+        set_xlabel = plt.xlabel
+        set_ylabel = plt.ylabel
+        set_xlim = plt.xlim
+        set_ylim = plt.ylim
+    if n_fft is None:
+        if spec.shape[0] % 2 == 0:
+            n_fft = spec.shape[0] * 2
+        else:
+            n_fft = (spec.shape[0] - 1) * 2
+    hop = hop or n_fft // 4
+    if t is None:
+        t = np.arange(0, spec_np.shape[-1]) * hop / sr
+    if f is None:
+        f = np.arange(0, spec_np.shape[0]) * sr // 2 / (n_fft // 2) / 1000
+    im = ax.pcolormesh(
+        t, f, spec_np, rasterized=True, shading="auto", vmin=vmin, vmax=vmax, cmap=cmap
+    )
+    if title is not None:
+        set_title(title)
+    if xlabel is not None:
+        set_xlabel(xlabel)
+    if ylabel is not None:
+        set_ylabel(ylabel)
+    if xlim is not None:
+        set_xlim(xlim)
+    if ylim is not None:
+        set_ylim(ylim)
+    return im
+
+
+def spec_im(
+    audio: torch.Tensor,
+    figsize=(15, 5),
+    colorbar=False,
+    colorbar_format=None,
+    figure=None,
+    labels=True,
+    **kwargs,
+) -> Image:
+    audio = torch.as_tensor(audio)
+    if labels:
+        kwargs.setdefault("xlabel", "Time [s]")
+        kwargs.setdefault("ylabel", "Frequency [Hz]")
+    n_fft = kwargs.setdefault("n_fft", 1024)
+    hop = kwargs.setdefault("hop", 512)
+    w = torch.hann_window(n_fft, device=audio.device)
+    spec = torch.stft(audio, n_fft, hop, window=w, return_complex=False)
+    spec = spec.div_(w.pow(2).sum())
+    spec = torch.view_as_complex(spec).abs().clamp_min(1e-12).log10().mul(10)
+    kwargs.setdefault("vmax", max(0.0, spec.max().item()))
+
+    if figure is None:
+        figure = plt.figure(figsize=figsize)
+        figure.set_tight_layout(True)
+    if spec.dim() > 2:
+        spec = spec.squeeze(0)
+    im = specshow(spec, **kwargs)
+    if colorbar:
+        ckwargs = {}
+        if "ax" in kwargs:
+            if colorbar_format is None:
+                if kwargs.get("vmin", None) is not None or kwargs.get("vmax", None) is not None:
+                    colorbar_format = "%+2.0f dB"
+            ckwargs = {"ax": kwargs["ax"]}
+        plt.colorbar(im, format=colorbar_format, **ckwargs)
+    figure.canvas.draw()
+    return Image.frombytes("RGB", figure.canvas.get_width_height(), figure.canvas.tostring_rgb())
+
+
+def toggle(choice):
+    if choice == "mic":
+        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
+    else:
+        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
+
+# Bark
+
+settings = Settings('config.yaml')
+
+def generate_text_to_speech(text, selected_speaker, text_temp, waveform_temp, eos_prob, quick_generation, complete_settings, seed, batchcount, progress=gr.Progress(track_tqdm=True)):
+    # Chunk the text into smaller pieces then combine the generated audio
+
+    # generation settings
+    if selected_speaker == 'None':
+        selected_speaker = None
+
+    voice_name = selected_speaker
+
+    if text == None or len(text) < 1:
+       if selected_speaker == None:
+            raise gr.Error('No text entered!')
+
+       # Extract audio data from speaker if no text and speaker selected
+       voicedata = _load_history_prompt(voice_name)
+       audio_arr = codec_decode(voicedata["fine_prompt"])
+       result = create_filename(settings.output_folder_path, "None", "extract",".wav")
+       save_wav(audio_arr, result)
+       return result
+
+    if batchcount < 1:
+        batchcount = 1
+
+
+    silenceshort = np.zeros(int((float(settings.silence_sentence) / 1000.0) * SAMPLE_RATE), dtype=np.int16)  # quarter second of silence
+    silencelong = np.zeros(int((float(settings.silence_speakers) / 1000.0) * SAMPLE_RATE), dtype=np.float32)  # half a second of silence
+    use_last_generation_as_history = "Use last generation as history" in complete_settings
+    save_last_generation = "Save generation as Voice" in complete_settings
+    for l in range(batchcount):
+        currentseed = seed
+        if seed != None and seed > 2**32 - 1:
+            logger.warning(f"Seed {seed} > 2**32 - 1 (max), setting to random")
+            currentseed = None
+        if currentseed == None or currentseed <= 0:
+            currentseed = np.random.default_rng().integers(1, 2**32 - 1)
+        assert(0 < currentseed and currentseed < 2**32)
+
+        progress(0, desc="Generating")
+
+        full_generation = None
+
+        all_parts = []
+        complete_text = ""
+        text = text.lstrip()
+        if is_ssml(text):
+            list_speak = create_clips_from_ssml(text)
+            prev_speaker = None
+            for i, clip in tqdm(enumerate(list_speak), total=len(list_speak)):
+                selected_speaker = clip[0]
+                # Add pause break between speakers
+                if i > 0 and selected_speaker != prev_speaker:
+                    all_parts += [silencelong.copy()]
+                prev_speaker = selected_speaker
+                text = clip[1]
+                text = saxutils.unescape(text)
+                if selected_speaker == "None":
+                    selected_speaker = None
+
+                print(f"\nGenerating Text ({i+1}/{len(list_speak)}) -> {selected_speaker} (Seed {currentseed}):`{text}`")
+                complete_text += text
+                with pytorch_seed.SavedRNG(currentseed):
+                    audio_array = generate_with_settings(text_prompt=text, voice_name=selected_speaker, semantic_temp=text_temp, coarse_temp=waveform_temp, eos_p=eos_prob)
+                    currentseed = torch.random.initial_seed()
+                if len(list_speak) > 1:
+                    filename = create_filename(settings.output_folder_path, currentseed, "audioclip",".wav")
+                    save_wav(audio_array, filename)
+                    add_id3_tag(filename, text, selected_speaker, currentseed)
+
+                all_parts += [audio_array]
+        else:
+            texts = split_and_recombine_text(text, settings.input_text_desired_length, settings.input_text_max_length)
+            for i, text in tqdm(enumerate(texts), total=len(texts)):
+                print(f"\nGenerating Text ({i+1}/{len(texts)}) -> {selected_speaker} (Seed {currentseed}):`{text}`")
+                complete_text += text
+                if quick_generation == True:
+                    with pytorch_seed.SavedRNG(currentseed):
+                        audio_array = generate_with_settings(text_prompt=text, voice_name=selected_speaker, semantic_temp=text_temp, coarse_temp=waveform_temp, eos_p=eos_prob)
+                        currentseed = torch.random.initial_seed()
+                else:
+                    full_output = use_last_generation_as_history or save_last_generation
+                    if full_output:
+                        full_generation, audio_array = generate_with_settings(text_prompt=text, voice_name=voice_name, semantic_temp=text_temp, coarse_temp=waveform_temp, eos_p=eos_prob, output_full=True)
+                    else:
+                        audio_array = generate_with_settings(text_prompt=text, voice_name=voice_name, semantic_temp=text_temp, coarse_temp=waveform_temp, eos_p=eos_prob)
+
+                # Noticed this in the HF Demo - convert to 16bit int -32767/32767 - most used audio format  
+                # audio_array = (audio_array * 32767).astype(np.int16)
+
+                if len(texts) > 1:
+                    filename = create_filename(settings.output_folder_path, currentseed, "audioclip",".wav")
+                    save_wav(audio_array, filename)
+                    add_id3_tag(filename, text, selected_speaker, currentseed)
+
+                if quick_generation == False and (save_last_generation == True or use_last_generation_as_history == True):
+                    # save to npz
+                    voice_name = create_filename(settings.output_folder_path, seed, "audioclip", ".npz")
+                    save_as_prompt(voice_name, full_generation)
+                    if use_last_generation_as_history:
+                        selected_speaker = voice_name
+
+                all_parts += [audio_array]
+                # Add short pause between sentences
+                if text[-1] in "!?.\n" and i > 1:
+                    all_parts += [silenceshort.copy()]
+
+        # save & play audio
+        result = create_filename(settings.output_folder_path, currentseed, "final",".wav")
+        save_wav(np.concatenate(all_parts), result)
+        # write id3 tag with text truncated to 60 chars, as a precaution...
+        add_id3_tag(result, complete_text, selected_speaker, currentseed)
+
+    return result
+
+
+
+def save_wav(audio_array, filename):
+    write_wav(filename, SAMPLE_RATE, audio_array)
+
+def save_voice(filename, semantic_prompt, coarse_prompt, fine_prompt):
+    np.savez_compressed(
+        filename,
+        semantic_prompt=semantic_prompt,
+        coarse_prompt=coarse_prompt,
+        fine_prompt=fine_prompt
+    )
+    
+
+def on_quick_gen_changed(checkbox):
+    if checkbox == False:
+        return gr.CheckboxGroup.update(visible=True)
+    return gr.CheckboxGroup.update(visible=False)
+
+def delete_output_files(checkbox_state):
+    if checkbox_state:
+        outputs_folder = os.path.join(os.getcwd(), settings.output_folder_path)
+        if os.path.exists(outputs_folder):
+            purgedir(outputs_folder)
+    return False
+
+
+# https://stackoverflow.com/a/54494779
+def purgedir(parent):
+    for root, dirs, files in os.walk(parent):                                      
+        for item in files:
+            # Delete subordinate files                                                 
+            filespec = os.path.join(root, item)
+            os.unlink(filespec)
+        for item in dirs:
+            # Recursively perform this operation for subordinate directories   
+            purgedir(os.path.join(root, item))
+
+def convert_text_to_ssml(text, selected_speaker):
+    return build_ssml(text, selected_speaker)
+
+
+def training_prepare(selected_step, num_text_generations, progress=gr.Progress(track_tqdm=True)):
+    if selected_step == prepare_training_list[0]:
+        prepare_semantics_from_text()
+    else:
+        prepare_wavs_from_semantics()
+    return None
+
+
+def start_training(save_model_epoch, max_epochs, progress=gr.Progress(track_tqdm=True)):
+    training_prepare_files("./training/data/", "./training/data/checkpoint/hubert_base_ls960.pt")
+    train("./training/data/", save_model_epoch, max_epochs)
+    return None
+
+
+
+def apply_settings(themes, input_server_name, input_server_port, input_server_public, input_desired_len, input_max_len, input_silence_break, input_silence_speaker):
+    settings.selected_theme = themes
+    settings.server_name = input_server_name
+    settings.server_port = input_server_port
+    settings.server_share = input_server_public
+    settings.input_text_desired_length = input_desired_len
+    settings.input_text_max_length = input_max_len
+    settings.silence_sentence = input_silence_break
+    settings.silence_speaker = input_silence_speaker
+    settings.save()
+
+def restart():
+    global restart_server
+    restart_server = True
+
+
+def create_version_html():
+    python_version = ".".join([str(x) for x in sys.version_info[0:3]])
+    versions_html = f"""
+python: <span title="{sys.version}">{python_version}</span>
+ • 
+torch: {getattr(torch, '__long_version__',torch.__version__)}
+ • 
+gradio: {gr.__version__}
+"""
+    return versions_html
+
+    
+
+logger = logging.getLogger(__name__)
+APPTITLE = "Bark Voice Cloning UI"
+
+
+autolaunch = False
+
+if len(sys.argv) > 1:
+    autolaunch = "-autolaunch" in sys.argv
+
+if torch.cuda.is_available() == False:
+    os.environ['BARK_FORCE_CPU'] = 'True'
+    logger.warning("No CUDA detected, fallback to CPU!")
+
+print(f'smallmodels={os.environ.get("SUNO_USE_SMALL_MODELS", False)}')
+print(f'enablemps={os.environ.get("SUNO_ENABLE_MPS", False)}')
+print(f'offloadcpu={os.environ.get("SUNO_OFFLOAD_CPU", False)}')
+print(f'forcecpu={os.environ.get("BARK_FORCE_CPU", False)}')
+print(f'autolaunch={autolaunch}\n\n')
+
+#print("Updating nltk\n")
+#nltk.download('punkt')
+
+print("Preloading Models\n")
+preload_models()
+
+available_themes = ["Default", "gradio/glass", "gradio/monochrome", "gradio/seafoam", "gradio/soft", "gstaff/xkcd", "freddyaboulton/dracula_revamped", "ysharma/steampunk"]
+tokenizer_language_list = ["de","en", "pl"]
+prepare_training_list = ["Step 1: Semantics from Text","Step 2: WAV from Semantics"]
+
+seed = -1
+server_name = settings.server_name
+if len(server_name) < 1:
+    server_name = None
+server_port = settings.server_port
+if server_port <= 0:
+    server_port = None
+global run_server
+global restart_server
+
+run_server = True
+
+while run_server:
+    # Collect all existing speakers/voices in dir
+    speakers_list = []
+
+    for root, dirs, files in os.walk("./bark/assets/prompts"):
+        for file in files:
+            if file.endswith(".npz"):
+                pathpart = root.replace("./bark/assets/prompts", "")
+                name = os.path.join(pathpart, file[:-4])
+                if name.startswith("/") or name.startswith("\\"):
+                     name = name[1:]
+                speakers_list.append(name)
+
+    speakers_list = sorted(speakers_list, key=lambda x: x.lower())
+    speakers_list.insert(0, 'None')
+
+    print(f'Launching {APPTITLE} Server')
+
+    # Create Gradio Blocks
+
+    with gr.Blocks(title=f"{APPTITLE}", mode=f"{APPTITLE}", theme=settings.selected_theme) as barkgui:
+        gr.Markdown("# <center>🐶🎶⭐ - Bark Voice Cloning</center>")
+        gr.Markdown("## <center>🤗 - If you like this space, please star my [github repo](https://github.com/KevinWang676/Bark-Voice-Cloning)</center>")
+        gr.Markdown("### <center>🎡 - Based on [bark-gui](https://github.com/C0untFloyd/bark-gui)</center>")
+        gr.Markdown(f""" You can duplicate and use it with a GPU: <a href="https://huggingface.co/spaces/{os.getenv('SPACE_ID')}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a>
+                         or open in [Colab](https://colab.research.google.com/github/KevinWang676/Bark-Voice-Cloning/blob/main/Bark_Voice_Cloning.ipynb) for quick start 🌟 P.S. Voice cloning needs a GPU, but TTS doesn't 😄
+                    """)
+
+        with gr.Tab("🎙️ - Clone Voice"):
+            with gr.Row():
+                input_audio_filename = gr.Audio(label="Input audio.wav", source="upload", type="filepath")
+            #transcription_text = gr.Textbox(label="Transcription Text", lines=1, placeholder="Enter Text of your Audio Sample here...")
+            with gr.Row():
+                with gr.Column():
+                    initialname = "/home/user/app/bark/assets/prompts/file"
+                    output_voice = gr.Textbox(label="Filename of trained Voice (do not change the initial name)", lines=1, placeholder=initialname, value=initialname, visible=False)
+                with gr.Column():
+                    tokenizerlang = gr.Dropdown(tokenizer_language_list, label="Base Language Tokenizer", value=tokenizer_language_list[1], visible=False)
+            with gr.Row():
+                clone_voice_button = gr.Button("Create Voice", variant="primary")
+            with gr.Row():
+                dummy = gr.Text(label="Progress")
+                npz_file = gr.File(label=".npz file")
+            speakers_list.insert(0, npz_file) # add prompt
+
+        with gr.Tab("🎵 - TTS"):
+            with gr.Row():
+                with gr.Column():
+                    placeholder = "Enter text here."
+                    input_text = gr.Textbox(label="Input Text", lines=4, placeholder=placeholder)
+                    convert_to_ssml_button = gr.Button("Convert Input Text to SSML")
+                with gr.Column():
+                        seedcomponent = gr.Number(label="Seed (default -1 = Random)", precision=0, value=-1)
+                        batchcount = gr.Number(label="Batch count", precision=0, value=1)
+           
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("[Voice Prompt Library](https://suno-ai.notion.site/8b8e8749ed514b0cbf3f699013548683?v=bc67cff786b04b50b3ceb756fd05f68c)")
+                    speaker = gr.Dropdown(speakers_list, value=speakers_list[0], label="Voice (Choose “file” if you wanna use the custom voice)")
+                    
+                with gr.Column():
+                    text_temp = gr.Slider(0.1, 1.0, value=0.6, label="Generation Temperature", info="1.0 more diverse, 0.1 more conservative")
+                    waveform_temp = gr.Slider(0.1, 1.0, value=0.7, label="Waveform temperature", info="1.0 more diverse, 0.1 more conservative")
+
+            with gr.Row():
+                with gr.Column():
+                    quick_gen_checkbox = gr.Checkbox(label="Quick Generation", value=True)
+                    settings_checkboxes = ["Use last generation as history", "Save generation as Voice"]
+                    complete_settings = gr.CheckboxGroup(choices=settings_checkboxes, value=settings_checkboxes, label="Detailed Generation Settings", type="value", interactive=True, visible=False)
+                with gr.Column():
+                    eos_prob = gr.Slider(0.0, 0.5, value=0.05, label="End of sentence probability")
+
+            with gr.Row():
+                with gr.Column():
+                    tts_create_button = gr.Button("Generate", variant="primary")
+                with gr.Column():
+                    hidden_checkbox = gr.Checkbox(visible=False)
+                    button_stop_generation = gr.Button("Stop generation")
+            with gr.Row():
+                output_audio = gr.Audio(label="Generated Audio", type="filepath")
+
+            with gr.Row():
+                with gr.Column():
+                    radio = gr.Radio(
+                        ["mic", "file"], value="file", label="How would you like to upload your audio?", visible=False
+                    )
+                    mic_input = gr.Mic(label="Input", type="filepath", visible=False)
+                    audio_file = output_audio
+                    inputs = [
+                        audio_file,
+                        gr.Dropdown(
+                            label="Add background noise",
+                            choices=list(NOISES.keys()),
+                            value="None", visible =False,
+                        ),
+                        gr.Dropdown(
+                            label="Noise Level (SNR)",
+                            choices=["-5", "0", "10", "20"],
+                            value="0", visible =False,
+                        ),
+                        mic_input,
+                    ]
+                    btn_denoise = gr.Button("Denoise", variant="primary")
+                with gr.Column():
+                    outputs = [
+                        gr.Audio(type="filepath", label="Enhanced audio"),
+                    ]
+            btn_denoise.click(fn=demo_fn, inputs=inputs, outputs=outputs)
+            radio.change(toggle, radio, [mic_input, audio_file])
+
+        with gr.Tab("🔮 - Voice Conversion"):
+            with gr.Row():
+                 swap_audio_filename = gr.Audio(label="Input audio.wav to swap voice", source="upload", type="filepath")
+            with gr.Row():
+                 with gr.Column():
+                     swap_tokenizer_lang = gr.Dropdown(tokenizer_language_list, label="Base Language Tokenizer", value=tokenizer_language_list[1])
+                     swap_seed = gr.Number(label="Seed (default -1 = Random)", precision=0, value=-1)
+                 with gr.Column():
+                     speaker_swap = gr.Dropdown(speakers_list, value=speakers_list[0], label="Voice (Choose “file” if you wanna use the custom voice)")
+                     swap_batchcount = gr.Number(label="Batch count", precision=0, value=1)
+            with gr.Row():
+                swap_voice_button = gr.Button("Generate", variant="primary")
+            with gr.Row():
+                output_swap = gr.Audio(label="Generated Audio", type="filepath")
+
+   
+        quick_gen_checkbox.change(fn=on_quick_gen_changed, inputs=quick_gen_checkbox, outputs=complete_settings)
+        convert_to_ssml_button.click(convert_text_to_ssml, inputs=[input_text, speaker],outputs=input_text)
+        gen_click = tts_create_button.click(generate_text_to_speech, inputs=[input_text, speaker, text_temp, waveform_temp, eos_prob, quick_gen_checkbox, complete_settings, seedcomponent, batchcount],outputs=output_audio)
+        button_stop_generation.click(fn=None, inputs=None, outputs=None, cancels=[gen_click])
+        
+
+
+        swap_voice_button.click(swap_voice_from_audio, inputs=[swap_audio_filename, speaker_swap, swap_tokenizer_lang, swap_seed, swap_batchcount], outputs=output_swap)
+        clone_voice_button.click(clone_voice, inputs=[input_audio_filename, output_voice], outputs=[dummy, npz_file])
+
+
+        restart_server = False
+        try:
+            barkgui.queue().launch(show_error=True)
+        except:
+            restart_server = True
+            run_server = False
+        try:
+            while restart_server == False:
+                time.sleep(1.0)
+        except (KeyboardInterrupt, OSError):
+            print("Keyboard interruption in main thread... closing server.")
+            run_server = False
+        barkgui.close()
+

bark/__init__.py

@@ -0,0 +1,2 @@
+from .api import generate_audio, text_to_semantic, semantic_to_waveform, save_as_prompt
+from .generation import SAMPLE_RATE, preload_models

bark/api.py

@@ -0,0 +1,158 @@
+from typing import Dict, Optional, Union
+
+import numpy as np
+
+from .generation import codec_decode, generate_coarse, generate_fine, generate_text_semantic
+
+
+def generate_with_settings(text_prompt, semantic_temp=0.6, eos_p=0.2, coarse_temp=0.7, fine_temp=0.5, voice_name=None, output_full=False):
+
+    # generation with more control
+    x_semantic = generate_text_semantic(
+        text_prompt,
+        history_prompt=voice_name,
+        temp=semantic_temp,
+        min_eos_p = eos_p,
+        use_kv_caching=True
+    )
+
+    x_coarse_gen = generate_coarse(
+        x_semantic,
+        history_prompt=voice_name,
+        temp=coarse_temp,
+        use_kv_caching=True
+    )
+    x_fine_gen = generate_fine(
+        x_coarse_gen,
+        history_prompt=voice_name,
+        temp=fine_temp,
+    )
+
+    if output_full:
+        full_generation = {
+            'semantic_prompt': x_semantic,
+            'coarse_prompt': x_coarse_gen,
+            'fine_prompt': x_fine_gen
+        }
+        return full_generation, codec_decode(x_fine_gen)
+    return codec_decode(x_fine_gen)
+
+
+def text_to_semantic(
+    text: str,
+    history_prompt: Optional[Union[Dict, str]] = None,
+    temp: float = 0.7,
+    silent: bool = False,
+):
+    """Generate semantic array from text.
+
+    Args:
+        text: text to be turned into audio
+        history_prompt: history choice for audio cloning
+        temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+        silent: disable progress bar
+
+    Returns:
+        numpy semantic array to be fed into `semantic_to_waveform`
+    """
+    x_semantic = generate_text_semantic(
+        text,
+        history_prompt=history_prompt,
+        temp=temp,
+        silent=silent,
+        use_kv_caching=True
+    )
+    return x_semantic
+
+
+def semantic_to_waveform(
+    semantic_tokens: np.ndarray,
+    history_prompt: Optional[Union[Dict, str]] = None,
+    temp: float = 0.7,
+    silent: bool = False,
+    output_full: bool = False,
+):
+    """Generate audio array from semantic input.
+
+    Args:
+        semantic_tokens: semantic token output from `text_to_semantic`
+        history_prompt: history choice for audio cloning
+        temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+        silent: disable progress bar
+        output_full: return full generation to be used as a history prompt
+
+    Returns:
+        numpy audio array at sample frequency 24khz
+    """
+    coarse_tokens = generate_coarse(
+        semantic_tokens,
+        history_prompt=history_prompt,
+        temp=temp,
+        silent=silent,
+        use_kv_caching=True
+    )
+    fine_tokens = generate_fine(
+        coarse_tokens,
+        history_prompt=history_prompt,
+        temp=0.5,
+    )
+    audio_arr = codec_decode(fine_tokens)
+    if output_full:
+        full_generation = {
+            "semantic_prompt": semantic_tokens,
+            "coarse_prompt": coarse_tokens,
+            "fine_prompt": fine_tokens,
+        }
+        return full_generation, audio_arr
+    return audio_arr
+
+
+def save_as_prompt(filepath, full_generation):
+    assert(filepath.endswith(".npz"))
+    assert(isinstance(full_generation, dict))
+    assert("semantic_prompt" in full_generation)
+    assert("coarse_prompt" in full_generation)
+    assert("fine_prompt" in full_generation)
+    np.savez(filepath, **full_generation)
+
+
+def generate_audio(
+    text: str,
+    history_prompt: Optional[Union[Dict, str]] = None,
+    text_temp: float = 0.7,
+    waveform_temp: float = 0.7,
+    silent: bool = False,
+    output_full: bool = False,
+):
+    """Generate audio array from input text.
+
+    Args:
+        text: text to be turned into audio
+        history_prompt: history choice for audio cloning
+        text_temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+        waveform_temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+        silent: disable progress bar
+        output_full: return full generation to be used as a history prompt
+
+    Returns:
+        numpy audio array at sample frequency 24khz
+    """
+    semantic_tokens = text_to_semantic(
+        text,
+        history_prompt=history_prompt,
+        temp=text_temp,
+        silent=silent,
+    )
+    out = semantic_to_waveform(
+        semantic_tokens,
+        history_prompt=history_prompt,
+        temp=waveform_temp,
+        silent=silent,
+        output_full=output_full,
+    )
+    if output_full:
+        full_generation, audio_arr = out
+        return full_generation, audio_arr
+    else:
+        audio_arr = out
+    return audio_arr

bark/assets/prompts/announcer.npz

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

bark/assets/prompts/v2/en_speaker_0.npz

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

bark/assets/prompts/v2/en_speaker_1.npz

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

bark/assets/prompts/v2/en_speaker_2.npz

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

bark/assets/prompts/v2/en_speaker_3.npz

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

bark/assets/prompts/v2/en_speaker_4.npz

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

bark/assets/prompts/v2/en_speaker_5.npz

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

bark/assets/prompts/v2/en_speaker_6.npz

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

bark/assets/prompts/v2/en_speaker_7.npz

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

bark/assets/prompts/v2/en_speaker_8.npz

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

bark/assets/prompts/v2/en_speaker_9.npz

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

bark/assets/prompts/v2/zh_speaker_0.npz

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

bark/assets/prompts/v2/zh_speaker_1.npz

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

bark/assets/prompts/v2/zh_speaker_2.npz

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

bark/assets/prompts/v2/zh_speaker_3.npz

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

bark/assets/prompts/v2/zh_speaker_4.npz

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

bark/assets/prompts/v2/zh_speaker_5.npz

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

bark/assets/prompts/v2/zh_speaker_6.npz

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

bark/assets/prompts/v2/zh_speaker_7.npz

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

bark/assets/prompts/v2/zh_speaker_8.npz

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

bark/assets/prompts/v2/zh_speaker_9.npz

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

bark/generation.py

@@ -0,0 +1,864 @@
+import contextlib
+import gc
+import os
+import re
+import requests
+import gc
+import sys
+
+from encodec import EncodecModel
+import funcy
+import logging
+import numpy as np
+from scipy.special import softmax
+import torch
+import torch.nn.functional as F
+import tqdm
+from transformers import BertTokenizer
+from huggingface_hub import hf_hub_download, hf_hub_url
+
+from .model import GPTConfig, GPT
+from .model_fine import FineGPT, FineGPTConfig
+from .settings import initenv
+
+initenv(sys.argv)
+global_force_cpu = os.environ.get("BARK_FORCE_CPU", False)
+if (
+    global_force_cpu != True and
+    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
+
+
+# hold models in global scope to lazy load
+global models
+models = {}
+
+global models_devices
+models_devices = {}
+
+
+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
+
+
+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"),
+]
+
+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}")
+
+
+logger = logging.getLogger(__name__)
+
+
+CUR_PATH = os.path.dirname(os.path.abspath(__file__))
+
+
+#default_cache_dir = os.path.join(os.path.expanduser("~"), ".cache")
+#CACHE_DIR = os.path.join(os.getenv("XDG_CACHE_HOME", default_cache_dir), "suno", "bark_v0")
+#CACHE_DIR = os.path.join(os.getcwd(), "models"
+CACHE_DIR = "./models"
+
+
+def _cast_bool_env_var(s):
+    return s.lower() in ('true', '1', 't')
+
+USE_SMALL_MODELS = _cast_bool_env_var(os.environ.get("SUNO_USE_SMALL_MODELS", "False"))
+GLOBAL_ENABLE_MPS = _cast_bool_env_var(os.environ.get("SUNO_ENABLE_MPS", "False"))
+OFFLOAD_CPU = _cast_bool_env_var(os.environ.get("SUNO_OFFLOAD_CPU", "False"))
+
+REMOTE_MODEL_PATHS = {
+    "text_small": {
+        "repo_id": "suno/bark",
+        "file_name": "text.pt",
+    },
+    "coarse_small": {
+        "repo_id": "suno/bark",
+        "file_name": "coarse.pt",
+    },
+    "fine_small": {
+        "repo_id": "suno/bark",
+        "file_name": "fine.pt",
+    },
+    "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",
+    },
+}
+
+
+if not hasattr(torch.nn.functional, 'scaled_dot_product_attention') and torch.cuda.is_available():
+    logger.warning(
+        "torch version does not support flash attention. You will get faster" +
+        " inference speed by upgrade torch to newest nightly version."
+    )
+
+
+def grab_best_device(use_gpu=True):
+    if torch.cuda.device_count() > 0 and use_gpu:
+        device = "cuda"
+    elif torch.backends.mps.is_available() and use_gpu and GLOBAL_ENABLE_MPS:
+        device = "mps"
+    else:
+        device = "cpu"
+    return device
+
+
+def _get_ckpt_path(model_type, use_small=False):
+    key = model_type
+    if use_small or USE_SMALL_MODELS:
+        key += "_small"
+    return os.path.join(CACHE_DIR, REMOTE_MODEL_PATHS[key]["file_name"])
+
+"""
+def _download(from_hf_path, file_name, destfilename):
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    hf_hub_download(repo_id=from_hf_path, filename=file_name, local_dir=CACHE_DIR, local_dir_use_symlinks=False)
+    # Bug in original repo? Downloaded name differs from expected...
+    if not os.path.exists(destfilename):
+        localname = os.path.join(CACHE_DIR, file_name)
+        os.rename(localname, destfilename)
+"""
+def _download(from_hf_path, file_name):
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    hf_hub_download(repo_id=from_hf_path, filename=file_name, local_dir=CACHE_DIR)
+
+
+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
+
+
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+
+@contextlib.contextmanager
+def _inference_mode():
+    with InferenceContext(), torch.inference_mode(), torch.no_grad(), autocast():
+        yield
+
+
+def _clear_cuda_cache():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+def clean_models(model_key=None):
+    global models
+    model_keys = [model_key] if model_key is not None else models.keys()
+    for k in model_keys:
+        if k in models:
+            del models[k]
+    _clear_cuda_cache()
+    gc.collect()
+
+
+def _load_model(ckpt_path, device, use_small=False, model_type="text"):
+    if model_type == "text":
+        ConfigClass = GPTConfig
+        ModelClass = GPT
+    elif model_type == "coarse":
+        ConfigClass = GPTConfig
+        ModelClass = GPT
+    elif model_type == "fine":
+        ConfigClass = FineGPTConfig
+        ModelClass = FineGPT
+    else:
+        raise NotImplementedError()
+
+    # Force-remove Models to allow running on >12Gb GPU
+    # CF: Probably not needed anymore
+    #global models
+    #models.clear()
+    #gc.collect()
+    #torch.cuda.empty_cache()
+    # to here...
+
+    model_key = f"{model_type}_small" if use_small or USE_SMALL_MODELS else model_type
+    model_info = REMOTE_MODEL_PATHS[model_key]
+    if not os.path.exists(ckpt_path):
+        logger.info(f"{model_type} model not found, downloading into `{CACHE_DIR}`.")
+        ## added next two lines to make it super clear which model is being downloaded
+        remote_filename = hf_hub_url(model_info["repo_id"], model_info["file_name"])
+        print(f"Downloading {model_key} {model_info['repo_id']} remote model file {remote_filename} {model_info['file_name']} to {CACHE_DIR}")
+        _download(model_info["repo_id"], model_info["file_name"])
+        # add next line to make it super clear which model is being loaded
+    print(f"Loading {model_key} model from {ckpt_path} to {device}") # added
+    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 = ConfigClass(**checkpoint["model_args"])
+    model = ModelClass(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 loaded: {round(n_params/1e6,1)}M params, {round(val_loss,3)} loss")
+    model.eval()
+    model.to(device)
+    del checkpoint, state_dict
+    _clear_cuda_cache()
+    if model_type == "text":
+        tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased")
+        return {
+            "model": model,
+            "tokenizer": tokenizer,
+        }
+    return model
+
+
+def _load_codec_model(device):
+    model = EncodecModel.encodec_model_24khz()
+    model.set_target_bandwidth(6.0)
+    model.eval()
+    model.to(device)
+    _clear_cuda_cache()
+    return model
+
+
+def load_model(use_gpu=True, use_small=False, force_reload=False, model_type="text"):
+    _load_model_f = funcy.partial(_load_model, model_type=model_type, use_small=use_small)
+    if model_type not in ("text", "coarse", "fine"):
+        raise NotImplementedError()
+    global models
+    global models_devices
+    device = grab_best_device(use_gpu=use_gpu)
+    model_key = f"{model_type}"
+    if OFFLOAD_CPU:
+        models_devices[model_key] = device
+        device = "cpu"
+    if model_key not in models or force_reload:
+        ckpt_path = _get_ckpt_path(model_type, use_small=use_small)
+        clean_models(model_key=model_key)
+        model = _load_model_f(ckpt_path, device)
+        models[model_key] = model
+    if model_type == "text":
+        models[model_key]["model"].to(device)
+    else:
+        models[model_key].to(device)
+    return models[model_key]
+
+
+def load_codec_model(use_gpu=True, force_reload=False):
+    global models
+    global models_devices
+    device = grab_best_device(use_gpu=use_gpu)
+    if device == "mps":
+        # encodec doesn't support mps
+        device = "cpu"
+    model_key = "codec"
+    if OFFLOAD_CPU:
+        models_devices[model_key] = device
+        device = "cpu"
+    if model_key not in models or force_reload:
+        clean_models(model_key=model_key)
+        model = _load_codec_model(device)
+        models[model_key] = model
+    models[model_key].to(device)
+    return models[model_key]
+
+
+def preload_models(
+    text_use_gpu=True,
+    text_use_small=False,
+    coarse_use_gpu=True,
+    coarse_use_small=False,
+    fine_use_gpu=True,
+    fine_use_small=False,
+    codec_use_gpu=True,
+    force_reload=False
+):
+    """Load all the necessary models for the pipeline."""
+    if grab_best_device() == "cpu" and (
+        text_use_gpu or coarse_use_gpu or fine_use_gpu or codec_use_gpu
+    ):
+        logger.warning("No GPU being used. Careful, inference might be very slow!")
+    _ = load_model(
+        model_type="text", use_gpu=text_use_gpu, use_small=text_use_small, force_reload=force_reload
+    )
+    _ = load_model(
+        model_type="coarse",
+        use_gpu=coarse_use_gpu,
+        use_small=coarse_use_small,
+        force_reload=force_reload,
+    )
+    _ = load_model(
+        model_type="fine", use_gpu=fine_use_gpu, use_small=fine_use_small, force_reload=force_reload
+    )
+    _ = load_codec_model(use_gpu=codec_use_gpu, force_reload=force_reload)
+
+
+####
+# Generation Functionality
+####
+
+
+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()
+
+
+TEXT_ENCODING_OFFSET = 10_048
+SEMANTIC_PAD_TOKEN = 10_000
+TEXT_PAD_TOKEN = 129_595
+SEMANTIC_INFER_TOKEN = 129_599
+
+
+def _load_history_prompt(history_prompt_input):
+    if isinstance(history_prompt_input, str) and history_prompt_input.endswith(".npz"):
+        history_prompt = np.load(history_prompt_input)
+    elif isinstance(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(CUR_PATH, "assets", "prompts", f"{history_prompt_input}.npz")
+        )
+    elif isinstance(history_prompt_input, dict):
+        assert("semantic_prompt" in history_prompt_input)
+        assert("coarse_prompt" in history_prompt_input)
+        assert("fine_prompt" in history_prompt_input)
+        history_prompt = history_prompt_input
+    else:
+        raise ValueError("history prompt format unrecognized")
+    return history_prompt
+
+
+def generate_text_semantic(
+    text,
+    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,
+    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)
+        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
+    # load models if not yet exist
+    global models
+    global models_devices
+    if "text" not in models:
+        preload_models()
+    model_container = models["text"]
+    model = model_container["model"]
+    tokenizer = model_container["tokenizer"]
+    encoded_text = np.array(_tokenize(tokenizer, text)) + 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=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(
+                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)
+            # multinomial bugged on mps: shuttle to cpu if necessary
+            inf_device = probs.device
+            if probs.device.type == "mps":
+                probs = probs.to("cpu")
+            item_next = torch.multinomial(probs, num_samples=1)
+            probs = probs.to(inf_device)
+            item_next = item_next.to(inf_device)
+            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(100 - 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(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 :]
+    if OFFLOAD_CPU:
+        model.to("cpu")
+    assert all(0 <= out) and all(out < SEMANTIC_VOCAB_SIZE)
+    _clear_cuda_cache()
+    return out
+
+
+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
+
+
+COARSE_SEMANTIC_PAD_TOKEN = 12_048
+COARSE_INFER_TOKEN = 12_050
+
+
+def generate_coarse(
+    x_semantic,
+    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,
+    use_kv_caching=False,
+):
+    """Generate coarse audio codes from semantic tokens."""
+# CF: Uncommented because it breaks swap voice more than once
+#    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)
+        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)
+    # load models if not yet exist
+    global models
+    global models_devices
+    if "coarse" not in models:
+        preload_models()
+    model = models["coarse"]
+    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)
+                # multinomial bugged on mps: shuttle to cpu if necessary
+                inf_device = probs.device
+                if probs.device.type == "mps":
+                    probs = probs.to("cpu")
+                item_next = torch.multinomial(probs, num_samples=1)
+                probs = probs.to(inf_device)
+                item_next = item_next.to(inf_device)
+                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(
+    x_coarse_gen,
+    history_prompt=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)
+        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]
+    # load models if not yet exist
+    global models
+    global models_devices
+    if "fine" not in models:
+        preload_models()
+    model = models["fine"]
+    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)
+                    # multinomial bugged on mps: shuttle to cpu if necessary
+                    inf_device = probs.device
+                    if probs.device.type == "mps":
+                        probs = probs.to("cpu")
+                    codebook_preds = torch.hstack(
+                        [
+                            torch.multinomial(probs[nnn], num_samples=1).to(inf_device)
+                            for nnn 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, 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(fine_tokens):
+    """Turn quantized audio codes into audio array using encodec."""
+    # load models if not yet exist
+    global models
+    global models_devices
+    if "codec" not in models:
+        preload_models()
+    model = models["codec"]
+    if OFFLOAD_CPU:
+        model.to(models_devices["codec"])
+    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
+    if OFFLOAD_CPU:
+        model.to("cpu")
+    return audio_arr

bark/hubert/customtokenizer.py

@@ -0,0 +1,195 @@
+"""
+Custom tokenizer model.
+Author: https://www.github.com/gitmylo/
+License: MIT
+"""
+
+import json
+import os.path
+from zipfile import ZipFile
+
+import numpy
+import torch
+from torch import nn, optim
+from torch.serialization import MAP_LOCATION
+from tqdm.auto import tqdm
+
+
+class CustomTokenizer(nn.Module):
+    def __init__(self, hidden_size=1024, input_size=768, output_size=10000, version=0):
+        super(CustomTokenizer, self).__init__()
+        next_size = input_size
+        if version == 0:
+            self.lstm = nn.LSTM(input_size, hidden_size, 2, batch_first=True)
+            next_size = hidden_size
+        if version == 1:
+            self.lstm = nn.LSTM(input_size, hidden_size, 2, batch_first=True)
+            self.intermediate = nn.Linear(hidden_size, 4096)
+            next_size = 4096
+
+        self.fc = nn.Linear(next_size, output_size)
+        self.softmax = nn.LogSoftmax(dim=1)
+        self.optimizer: optim.Optimizer = None
+        self.lossfunc = nn.CrossEntropyLoss()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        self.version = version
+
+    def forward(self, x):
+        x, _ = self.lstm(x)
+        if self.version == 1:
+            x = self.intermediate(x)
+        x = self.fc(x)
+        x = self.softmax(x)
+        return x
+
+    @torch.no_grad()
+    def get_token(self, x):
+        """
+        Used to get the token for the first
+        :param x: An array with shape (N, input_size) where N is a whole number greater or equal to 1, and input_size is the input size used when creating the model.
+        :return: An array with shape (N,) where N is the same as N from the input. Every number in the array is a whole number in range 0...output_size - 1 where output_size is the output size used when creating the model.
+        """
+        return torch.argmax(self(x), dim=1)
+
+    def prepare_training(self):
+        self.optimizer = optim.Adam(self.parameters(), 0.001)
+
+    def train_step(self, x_train, y_train, log_loss=False):
+        # y_train = y_train[:-1]
+        # y_train = y_train[1:]
+
+        optimizer = self.optimizer
+        lossfunc = self.lossfunc
+        # Zero the gradients
+        self.zero_grad()
+
+        # Forward pass
+        y_pred = self(x_train)
+
+        y_train_len = len(y_train)
+        y_pred_len = y_pred.shape[0]
+
+        if y_train_len > y_pred_len:
+            diff = y_train_len - y_pred_len
+            y_train = y_train[diff:]
+        elif y_train_len < y_pred_len:
+            diff = y_pred_len - y_train_len
+            y_pred = y_pred[:-diff, :]
+
+        y_train_hot = torch.zeros(len(y_train), self.output_size)
+        y_train_hot[range(len(y_train)), y_train] = 1
+        y_train_hot = y_train_hot.to('cuda')
+
+        # Calculate the loss
+        loss = lossfunc(y_pred, y_train_hot)
+
+        # Print loss
+        if log_loss:
+            print('Loss', loss.item())
+
+        # Backward pass
+        loss.backward()
+
+        # Update the weights
+        optimizer.step()
+
+    def save(self, path):
+        info_path = '.'.join(os.path.basename(path).split('.')[:-1]) + '/.info'
+        torch.save(self.state_dict(), path)
+        data_from_model = Data(self.input_size, self.hidden_size, self.output_size, self.version)
+        with ZipFile(path, 'a') as model_zip:
+            model_zip.writestr(info_path, data_from_model.save())
+            model_zip.close()
+
+    @staticmethod
+    def load_from_checkpoint(path, map_location: MAP_LOCATION = None):
+        old = True
+        with ZipFile(path) as model_zip:
+            filesMatch = [file for file in model_zip.namelist() if file.endswith('/.info')]
+            file = filesMatch[0] if filesMatch else None
+            if file:
+                old = False
+                print(f"Loading Custom Hubert Tokenizer {path}")
+                data_from_model = Data.load(model_zip.read(file).decode('utf-8'))
+            model_zip.close()
+        if old:
+            model = CustomTokenizer()
+        else:
+            model = CustomTokenizer(data_from_model.hidden_size, data_from_model.input_size, data_from_model.output_size, data_from_model.version)
+        model.load_state_dict(torch.load(path))
+        if map_location:
+            model = model.to(map_location)
+        return model
+
+
+
+class Data:
+    input_size: int
+    hidden_size: int
+    output_size: int
+    version: int
+
+    def __init__(self, input_size=768, hidden_size=1024, output_size=10000, version=0):
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        self.version = version
+
+    @staticmethod
+    def load(string):
+        data = json.loads(string)
+        return Data(data['input_size'], data['hidden_size'], data['output_size'], data['version'])
+
+    def save(self):
+        data = {
+            'input_size': self.input_size,
+            'hidden_size': self.hidden_size,
+            'output_size': self.output_size,
+            'version': self.version,
+        }
+        return json.dumps(data)
+
+
+def auto_train(data_path, save_path='model.pth', load_model: str | None = None, save_epochs=1, max_epochs=14):
+    data_x, data_y = [], []
+
+    if load_model and os.path.isfile(load_model):
+        print('Loading model from', load_model)
+        model_training = CustomTokenizer.load_from_checkpoint(load_model, 'cuda')
+    else:
+        print('Creating new model.')
+        model_training = CustomTokenizer(version=1).to('cuda')  # Settings for the model to run without lstm
+    save_path = os.path.join(data_path, save_path)
+    base_save_path = '.'.join(save_path.split('.')[:-1])
+
+    sem_string = '_semantic.npy'
+    feat_string = '_semantic_features.npy'
+
+    ready = os.path.join(data_path, 'ready')
+    for input_file in os.listdir(ready):
+        full_path = os.path.join(ready, input_file)
+        if input_file.endswith(sem_string):
+            data_y.append(numpy.load(full_path))
+        elif input_file.endswith(feat_string):
+            data_x.append(numpy.load(full_path))
+    model_training.prepare_training()
+
+    epoch = 1
+    with tqdm(total=((len(data_x) * len(data_y)) / 50) * save_epochs) as pbar1:
+        while epoch <= max_epochs:
+            for i in range(save_epochs):
+                j = 0
+                for x, y in zip(data_x, data_y):
+                    model_training.train_step(torch.tensor(x).to('cuda'), torch.tensor(y).to('cuda'), j % 50 == 0)  # Print loss every 50 steps
+                    j += 1
+                    pbar1.update()
+
+            save_p = save_path
+            save_p_2 = f'{base_save_path}_epoch_{epoch}.pth'
+            model_training.save(save_p)
+            model_training.save(save_p_2)
+            print(f'Epoch {epoch} completed')
+            epoch += 1
+        print(f'Done training for {max_epochs} epochs!')

bark/hubert/hubert_manager.py

@@ -0,0 +1,48 @@
+import os.path
+import shutil
+import urllib.request
+
+import huggingface_hub
+
+
+class HuBERTManager:
+
+
+    @staticmethod
+    def make_sure_hubert_installed(download_url: str = 'https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt', file_name: str = 'hubert.pt'):
+        install_dir = os.path.join('models', 'hubert')
+        if not os.path.isdir(install_dir):
+            os.makedirs(install_dir, exist_ok=True)
+        install_file = os.path.join(install_dir, file_name)
+        if not os.path.isfile(install_file):
+            print(f'Downloading HuBERT base model from {download_url}')
+            urllib.request.urlretrieve(download_url, install_file)
+            print('Downloaded HuBERT')
+        return install_file
+
+
+    @staticmethod
+    def make_sure_tokenizer_installed(model: str = 'quantifier_hubert_base_ls960_14.pth', repo: str = 'GitMylo/bark-voice-cloning', tokenizer_lang: str = 'en'):
+        local_file = tokenizer_lang + '_tokenizer.pth'
+        install_dir = os.path.join('models', 'hubert')
+        if not os.path.isdir(install_dir):
+            os.makedirs(install_dir, exist_ok=True)
+        install_file = os.path.join(install_dir, local_file)
+        if not os.path.isfile(install_file):
+            # refactor to use lists
+            if tokenizer_lang == 'en':
+                repo = 'GitMylo/bark-voice-cloning'
+                model = 'quantifier_hubert_base_ls960_14.pth'
+            elif tokenizer_lang == 'de':
+                repo = 'CountFloyd/bark-voice-cloning-german-HuBERT-quantizer'
+                model = 'german-HuBERT-quantizer_14_epoch.pth'
+            elif tokenizer_lang == 'pl':
+                repo = 'Hobis/bark-voice-cloning-polish-HuBERT-quantizer'
+                model = 'polish-HuBERT-quantizer_8_epoch.pth'
+            else:
+                raise 'Unknown Tokenizer Language!'
+            print(f'{local_file} not found. Downloading HuBERT custom tokenizer')
+            huggingface_hub.hf_hub_download(repo, model, local_dir=install_dir, local_dir_use_symlinks=False)
+            shutil.move(os.path.join(install_dir, model), install_file)
+            print('Downloaded tokenizer')
+        return install_file

bark/hubert/pre_kmeans_hubert.py

@@ -0,0 +1,107 @@
+"""
+Modified HuBERT model without kmeans.
+Original author: https://github.com/lucidrains/
+Modified by: https://www.github.com/gitmylo/
+License: MIT
+"""
+
+# Modified code from https://github.com/lucidrains/audiolm-pytorch/blob/main/audiolm_pytorch/hubert_kmeans.py
+
+from pathlib import Path
+
+import torch
+from torch import nn
+from einops import pack, unpack
+
+import fairseq
+
+from torchaudio.functional import resample
+
+from audiolm_pytorch.utils import curtail_to_multiple
+
+import logging
+logging.root.setLevel(logging.ERROR)
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+class CustomHubert(nn.Module):
+    """
+    checkpoint and kmeans can be downloaded at https://github.com/facebookresearch/fairseq/tree/main/examples/hubert
+    or you can train your own
+    """
+
+    def __init__(
+        self,
+        checkpoint_path,
+        target_sample_hz=16000,
+        seq_len_multiple_of=None,
+        output_layer=9,
+        device=None
+    ):
+        super().__init__()
+        self.target_sample_hz = target_sample_hz
+        self.seq_len_multiple_of = seq_len_multiple_of
+        self.output_layer = output_layer
+
+        if device is not None:
+            self.to(device)
+
+        model_path = Path(checkpoint_path)
+
+        assert model_path.exists(), f'path {checkpoint_path} does not exist'
+
+        print(f"Loading Hubert {checkpoint_path}")
+        checkpoint = torch.load(checkpoint_path)
+        load_model_input = {checkpoint_path: checkpoint}
+        model, *_ = fairseq.checkpoint_utils.load_model_ensemble_and_task(load_model_input)
+
+        if device is not None:
+            model[0].to(device)
+
+        self.model = model[0]
+        self.model.eval()
+
+    @property
+    def groups(self):
+        return 1
+
+    @torch.no_grad()
+    def forward(
+        self,
+        wav_input,
+        flatten=True,
+        input_sample_hz=None
+    ):
+        device = wav_input.device
+
+        if exists(input_sample_hz):
+            wav_input = resample(wav_input, input_sample_hz, self.target_sample_hz)
+
+        if exists(self.seq_len_multiple_of):
+            wav_input = curtail_to_multiple(wav_input, self.seq_len_multiple_of)
+
+        embed = self.model(
+            wav_input,
+            features_only=True,
+            mask=False,  # thanks to @maitycyrus for noticing that mask is defaulted to True in the fairseq code
+            output_layer=self.output_layer
+        )
+
+        embed, packed_shape = pack([embed['x']], '* d')
+
+        # codebook_indices = self.kmeans.predict(embed.cpu().detach().numpy())
+
+        codebook_indices = torch.from_numpy(embed.cpu().detach().numpy()).to(device)  # .long()
+
+        if flatten:
+            return codebook_indices
+
+        codebook_indices, = unpack(codebook_indices, packed_shape, '*')
+        return codebook_indices

bark/model.py

@@ -0,0 +1,218 @@
+"""
+Much of this code is adapted from Andrej Karpathy's NanoGPT
+(https://github.com/karpathy/nanoGPT)
+"""
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class LayerNorm(nn.Module):
+    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch nightly and still a bit scary
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            # print("WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0")
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                        .view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x, past_kv=None, use_cache=False):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        if past_kv is not None:
+            past_key = past_kv[0]
+            past_value = past_kv[1]
+            k = torch.cat((past_key, k), dim=-2)
+            v = torch.cat((past_value, v), dim=-2)
+
+        FULL_T = k.shape[-2]
+
+        if use_cache is True:
+            present = (k, v)
+        else:
+            present = None
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            if past_kv is not None:
+                # When `past_kv` is provided, we're doing incremental decoding and `q.shape[2] == 1`: q only contains
+                # the query for the last token. scaled_dot_product_attention interprets this as the first token in the
+                # sequence, so if is_causal=True it will mask out all attention from it. This is not what we want, so 
+                # to work around this we set is_causal=False.
+                is_causal = False
+            else:
+                is_causal = True
+
+            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout, is_causal=is_causal)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:,:,FULL_T-T:FULL_T,:FULL_T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return (y, present)
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+        self.gelu = nn.GELU()
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+        self.layer_idx = layer_idx
+
+    def forward(self, x, past_kv=None, use_cache=False):
+        attn_output, prev_kvs = self.attn(self.ln_1(x), past_kv=past_kv, use_cache=use_cache)
+        x = x + attn_output
+        x = x + self.mlp(self.ln_2(x))
+        return (x, prev_kvs)
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    input_vocab_size: int = 10_048
+    output_vocab_size: int = 10_048
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+class GPT(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.input_vocab_size is not None
+        assert config.output_vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.input_vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config, idx) for idx in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.output_vocab_size, bias=False)
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wte.weight.numel()
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def forward(self, idx, merge_context=False, past_kv=None, position_ids=None, use_cache=False):
+        device = idx.device
+        b, t = idx.size()
+        if past_kv is not None:
+            assert t == 1
+            tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        else:
+            if merge_context:
+                assert(idx.shape[1] >= 256+256+1)
+                t = idx.shape[1] - 256
+            else:
+                assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+
+            # forward the GPT model itself
+            if merge_context:
+                tok_emb = torch.cat([
+                    self.transformer.wte(idx[:,:256]) + self.transformer.wte(idx[:,256:256+256]),
+                    self.transformer.wte(idx[:,256+256:])
+                ], dim=1)
+            else:
+                tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+
+        if past_kv is None:
+            past_length = 0
+            past_kv = tuple([None] * len(self.transformer.h))
+        else:
+            past_length = past_kv[0][0].size(-2)
+
+        if position_ids is None:
+            position_ids = torch.arange(past_length, t + past_length, dtype=torch.long, device=device)
+            position_ids = position_ids.unsqueeze(0) # shape (1, t)
+            assert position_ids.shape == (1, t)
+
+        pos_emb = self.transformer.wpe(position_ids) # position embeddings of shape (1, t, n_embd)
+
+        x = self.transformer.drop(tok_emb + pos_emb)
+
+        new_kv = () if use_cache else None
+
+        for i, (block, past_layer_kv) in enumerate(zip(self.transformer.h, past_kv)):
+            x, kv = block(x, past_kv=past_layer_kv, use_cache=use_cache)
+
+            if use_cache:
+                new_kv = new_kv + (kv,)
+
+        x = self.transformer.ln_f(x)
+
+        # inference-time mini-optimization: only forward the lm_head on the very last position
+        logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+
+        return (logits, new_kv)

bark/model_fine.py

@@ -0,0 +1,149 @@
+"""
+Much of this code is adapted from Andrej Karpathy's NanoGPT
+(https://github.com/karpathy/nanoGPT)
+"""
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .model import GPT, GPTConfig, MLP
+
+
+class NonCausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch nightly and still a bit scary
+        self.flash = (
+            hasattr(torch.nn.functional, "scaled_dot_product_attention") and self.dropout == 0.0
+        )
+
+    def forward(self, x):
+        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=self.dropout, is_causal=False
+            )
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = (
+            y.transpose(1, 2).contiguous().view(B, T, C)
+        )  # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class FineBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = NonCausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class FineGPT(GPT):
+    def __init__(self, config):
+        super().__init__(config)
+        del self.lm_head
+        self.config = config
+        self.n_codes_total = config.n_codes_total
+        self.transformer = nn.ModuleDict(
+            dict(
+                wtes=nn.ModuleList(
+                    [
+                        nn.Embedding(config.input_vocab_size, config.n_embd)
+                        for _ in range(config.n_codes_total)
+                    ]
+                ),
+                wpe=nn.Embedding(config.block_size, config.n_embd),
+                drop=nn.Dropout(config.dropout),
+                h=nn.ModuleList([FineBlock(config) for _ in range(config.n_layer)]),
+                ln_f=nn.LayerNorm(config.n_embd),
+            )
+        )
+        self.lm_heads = nn.ModuleList(
+            [
+                nn.Linear(config.n_embd, config.output_vocab_size, bias=False)
+                for _ in range(config.n_codes_given, self.n_codes_total)
+            ]
+        )
+        for i in range(self.n_codes_total - config.n_codes_given):
+            self.transformer.wtes[i + 1].weight = self.lm_heads[i].weight
+
+    def forward(self, pred_idx, idx):
+        device = idx.device
+        b, t, codes = idx.size()
+        assert (
+            t <= self.config.block_size
+        ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        assert pred_idx > 0, "cannot predict 0th codebook"
+        assert codes == self.n_codes_total, (b, t, codes)
+        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)  # shape (1, t)
+
+        # forward the GPT model itself
+        tok_embs = [
+            wte(idx[:, :, i]).unsqueeze(-1) for i, wte in enumerate(self.transformer.wtes)
+        ]  # token embeddings of shape (b, t, n_embd)
+        tok_emb = torch.cat(tok_embs, dim=-1)
+        pos_emb = self.transformer.wpe(pos)  # position embeddings of shape (1, t, n_embd)
+        x = tok_emb[:, :, :, : pred_idx + 1].sum(dim=-1)
+        x = self.transformer.drop(x + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+        logits = self.lm_heads[pred_idx - self.config.n_codes_given](x)
+        return logits
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            for wte in self.transformer.wtes:
+                n_params -= wte.weight.numel()
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+
+@dataclass
+class FineGPTConfig(GPTConfig):
+    n_codes_total: int = 8
+    n_codes_given: int = 1

bark/settings.py

@@ -0,0 +1,7 @@
+import os
+
+def initenv(args):
+    os.environ['SUNO_USE_SMALL_MODELS'] = str("-smallmodels" in args)
+    os.environ['BARK_FORCE_CPU'] = str("-forcecpu" in args)
+    os.environ['SUNO_ENABLE_MPS'] = str("-enablemps" in args)
+    os.environ['SUNO_OFFLOAD_CPU'] = str("-offloadcpu" in args)

cloning/clonevoice.py

@@ -0,0 +1,68 @@
+from bark.generation import load_codec_model, generate_text_semantic, grab_best_device
+from encodec.utils import convert_audio
+from bark.hubert.hubert_manager import HuBERTManager
+from bark.hubert.pre_kmeans_hubert import CustomHubert
+from bark.hubert.customtokenizer import CustomTokenizer
+
+import torchaudio
+import torch
+import os
+import gradio
+
+
+def clone_voice(audio_filepath, dest_filename, progress=gradio.Progress(track_tqdm=True)):
+    # if len(text) < 1:
+    #    raise gradio.Error('No transcription text entered!')
+
+    use_gpu = False # not os.environ.get("BARK_FORCE_CPU", False)
+    progress(0, desc="Loading Codec")
+    model = load_codec_model(use_gpu=use_gpu)
+    
+    # From https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer
+    hubert_manager = HuBERTManager()
+    hubert_manager.make_sure_hubert_installed()
+    hubert_manager.make_sure_tokenizer_installed()
+
+    # From https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer 
+    # Load HuBERT for semantic tokens
+
+    # Load the HuBERT model
+    device = grab_best_device(use_gpu)
+    hubert_model = CustomHubert(checkpoint_path='./models/hubert/hubert.pt').to(device)
+
+    # Load the CustomTokenizer model
+    tokenizer = CustomTokenizer.load_from_checkpoint('./models/hubert/en_tokenizer.pth').to(device)  # change to the correct path 
+
+    progress(0.25, desc="Converting WAV")
+
+    # Load and pre-process the audio waveform
+    wav, sr = torchaudio.load(audio_filepath)
+    if wav.shape[0] == 2:  # Stereo to mono if needed
+        wav = wav.mean(0, keepdim=True)
+
+    wav = convert_audio(wav, sr, model.sample_rate, model.channels)
+    wav = wav.to(device)
+    progress(0.5, desc="Extracting codes")
+
+    semantic_vectors = hubert_model.forward(wav, input_sample_hz=model.sample_rate)
+    semantic_tokens = tokenizer.get_token(semantic_vectors)
+    
+    # Extract discrete codes from EnCodec
+    with torch.no_grad():
+        encoded_frames = model.encode(wav.unsqueeze(0))
+    codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1).squeeze()  # [n_q, T]
+
+    # get seconds of audio
+    # seconds = wav.shape[-1] / model.sample_rate
+    # generate semantic tokens
+    # semantic_tokens = generate_text_semantic(text, max_gen_duration_s=seconds, top_k=50, top_p=.95, temp=0.7)
+
+    # move codes to cpu
+    codes = codes.cpu().numpy()
+    # move semantic tokens to cpu
+    semantic_tokens = semantic_tokens.cpu().numpy()
+
+    import numpy as np
+    output_path = dest_filename + '.npz'
+    np.savez(output_path, fine_prompt=codes, coarse_prompt=codes[:2, :], semantic_prompt=semantic_tokens)
+    return ["Finished", output_path]

config.yaml

@@ -0,0 +1,8 @@
+input_text_desired_length: 110
+input_text_max_length: 170
+selected_theme: JohnSmith9982/small_and_pretty
+server_name: ''
+server_port: 0
+server_share: false
+silence_between_sentences: 250
+silence_between_speakers: 500

packages.txt

@@ -0,0 +1 @@
+ffmpeg

pyproject.toml

@@ -0,0 +1,69 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "bark-ui-enhanced"
+version = "0.7.0"
+description = "Bark text to audio model with addition features and a Web UI"
+readme = "README.md"
+requires-python = ">=3.8"
+authors =  [
+    {name = "Suno Inc (original Bark)", email = "[email protected]"},
+    {name = "Count Floyd"},
+]
+# MIT License
+license = {file = "LICENSE"}
+
+dependencies = [
+    "boto3",
+    "encodec",
+    "funcy",
+    "huggingface-hub>=0.14.1",
+    "numpy",
+    "scipy",
+    "tokenizers",
+    "torch",
+    "tqdm",
+    "transformers",
+]
+
+[project.urls]
+source = "https://github.com/C0untFloyd/bark-gui"
+
+[project.optional-dependencies]
+dev = [
+    "bandit",
+    "black",
+    "codecov",
+    "flake8",
+    "hypothesis>=6.14,<7",
+    "isort>=5.0.0,<6",
+    "jupyter",
+    "mypy",
+    "nbconvert",
+    "nbformat",
+    "pydocstyle",
+    "pylint",
+    "pytest",
+    "pytest-cov",
+]
+
+[tool.setuptools]
+packages = ["bark"]
+
+[tool.setuptools.package-data]
+bark = ["assets/prompts/*.npz", "assets/prompts/v2/*.npz"]
+
+
+[tool.black]
+line-length = 100
+target-version = ["py37", "py38", "py39", "py310"]
+include = '\.pyi?$'
+
+[tool.isort]
+profile = "black"
+line_length = 100
+skip_gitignore = true
+known_first_party = ["df", "libdf", "libdfdata"]
+

requirements.txt

@@ -0,0 +1,18 @@
+fairseq; platform_system != "Windows"
+fairseq@https://github.com/Sharrnah/fairseq/releases/download/v0.12.4/fairseq-0.12.4-cp310-cp310-win_amd64.whl; platform_system == "Windows"
+audiolm-pytorch
+gradio
+funcy
+linkify
+mutagen
+pytorch_seed
+pyyaml
+sentencepiece
+soundfile; platform_system == "Windows"
+sox; platform_system != "Windows"
+transformers
+torch==1.13.0
+torchaudio==0.13
+deepfilternet
+matplotlib
+Pillow

samples/dkitchen.wav

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

samples/nriver.wav

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

samples/scafe.wav

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