From 5eb7ebca911c57d2c497f3d5adfda78a851e778b Mon Sep 17 00:00:00 2001
From: Mylo <36931363+gitmylo@users.noreply.github.com>
Date: Tue, 18 Feb 2025 17:20:02 +0100
Subject: [PATCH] Scaffolding
---
.idea/ComfyUI-audio-nodes.iml | 53 ++-
.idea/vcs.xml | 2 +-
__init__.py | 9 +
lib/__init__.py | 0
lib/bark/__init__.py | 2 +
lib/bark/api.py | 125 +++++
lib/bark/generation.py | 820 ++++++++++++++++++++++++++++++++
lib/bark/model.py | 218 +++++++++
lib/bark/model_fine.py | 149 ++++++
lib/bark/readme.md | 4 +
lib/general/__init__.py | 0
lib/general/model_downloader.py | 33 ++
nodes/__init__.py | 4 +
nodes/bark/__init__.py | 3 +
nodes/bark/bark_loader.py | 33 ++
nodes/basenode.py | 31 ++
nodes/general/__init__.py | 3 +
nodes/general/encodec.py | 26 +
requirements.txt | 2 +
util/__init__.py | 0
util/categories.py | 5 +
21 files changed, 1518 insertions(+), 4 deletions(-)
create mode 100644 __init__.py
create mode 100644 lib/__init__.py
create mode 100644 lib/bark/__init__.py
create mode 100644 lib/bark/api.py
create mode 100644 lib/bark/generation.py
create mode 100644 lib/bark/model.py
create mode 100644 lib/bark/model_fine.py
create mode 100644 lib/bark/readme.md
create mode 100644 lib/general/__init__.py
create mode 100644 lib/general/model_downloader.py
create mode 100644 nodes/__init__.py
create mode 100644 nodes/bark/__init__.py
create mode 100644 nodes/bark/bark_loader.py
create mode 100644 nodes/basenode.py
create mode 100644 nodes/general/__init__.py
create mode 100644 nodes/general/encodec.py
create mode 100644 requirements.txt
create mode 100644 util/__init__.py
create mode 100644 util/categories.py
diff --git a/.idea/ComfyUI-audio-nodes.iml b/.idea/ComfyUI-audio-nodes.iml
index 0af965f..ef1450c 100644
--- a/.idea/ComfyUI-audio-nodes.iml
+++ b/.idea/ComfyUI-audio-nodes.iml
@@ -1,10 +1,57 @@
-
-
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
-
+
\ No newline at end of file
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
index 94a25f7..35eb1dd 100644
--- a/.idea/vcs.xml
+++ b/.idea/vcs.xml
@@ -1,6 +1,6 @@
-
+
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..ad99a34
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,9 @@
+from .lib.general.model_downloader import add_model_folder_path
+
+add_model_folder_path("bark") # Bark models are downloaded here
+
+# Node imports go here
+from .nodes import nodes
+
+NODE_CLASS_MAPPINGS = {node.name: node for node in nodes}
+NODE_DISPLAY_NAME_MAPPINGS = {node.name: node.display_name for node in nodes}
\ No newline at end of file
diff --git a/lib/__init__.py b/lib/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/lib/bark/__init__.py b/lib/bark/__init__.py
new file mode 100644
index 0000000..e0b17c8
--- /dev/null
+++ b/lib/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
diff --git a/lib/bark/api.py b/lib/bark/api.py
new file mode 100644
index 0000000..7b646b7
--- /dev/null
+++ b/lib/bark/api.py
@@ -0,0 +1,125 @@
+from typing import Dict, Optional, Union
+
+import numpy as np
+
+from .generation import codec_decode, generate_coarse, generate_fine, generate_text_semantic
+
+
+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
diff --git a/lib/bark/generation.py b/lib/bark/generation.py
new file mode 100644
index 0000000..54f9870
--- /dev/null
+++ b/lib/bark/generation.py
@@ -0,0 +1,820 @@
+import contextlib
+import gc
+import os
+import re
+
+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
+
+from .model import GPTConfig, GPT
+from .model_fine import FineGPT, FineGPTConfig
+
+if (
+ torch.cuda.is_available() and
+ hasattr(torch.cuda, "amp") and
+ hasattr(torch.cuda.amp, "autocast") and
+ hasattr(torch.cuda, "is_bf16_supported") and
+ torch.cuda.is_bf16_supported()
+):
+ autocast = funcy.partial(torch.cuda.amp.autocast, dtype=torch.bfloat16)
+else:
+ @contextlib.contextmanager
+ def autocast():
+ yield
+
+
+# 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")
+
+
+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):
+ 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 list(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()
+ 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}`.")
+ _download(model_info["repo_id"], model_info["file_name"])
+ 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=n_tot_steps)
+ pbar_state = 0
+ tot_generated_duration_s = 0
+ kv_cache = None
+ for n in range(n_tot_steps):
+ if use_kv_caching and kv_cache is not None:
+ x_input = x[:, [-1]]
+ else:
+ x_input = x
+ logits, kv_cache = model(
+ x_input, merge_context=True, use_cache=use_kv_caching, past_kv=kv_cache
+ )
+ relevant_logits = logits[0, 0, :SEMANTIC_VOCAB_SIZE]
+ if allow_early_stop:
+ relevant_logits = torch.hstack(
+ (relevant_logits, logits[0, 0, [SEMANTIC_PAD_TOKEN]]) # eos
+ )
+ if top_p is not None:
+ # faster to convert to numpy
+ original_device = relevant_logits.device
+ relevant_logits = relevant_logits.detach().cpu().type(torch.float32).numpy()
+ sorted_indices = np.argsort(relevant_logits)[::-1]
+ sorted_logits = relevant_logits[sorted_indices]
+ cumulative_probs = np.cumsum(softmax(sorted_logits))
+ sorted_indices_to_remove = cumulative_probs > top_p
+ sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].copy()
+ sorted_indices_to_remove[0] = False
+ relevant_logits[sorted_indices[sorted_indices_to_remove]] = -np.inf
+ relevant_logits = torch.from_numpy(relevant_logits)
+ relevant_logits = relevant_logits.to(original_device)
+ if top_k is not None:
+ v, _ = torch.topk(relevant_logits, min(top_k, relevant_logits.size(-1)))
+ relevant_logits[relevant_logits < v[-1]] = -float("Inf")
+ probs = F.softmax(relevant_logits / temp, dim=-1)
+ item_next = torch.multinomial(probs, num_samples=1).to(torch.int32)
+ if allow_early_stop and (
+ item_next == SEMANTIC_VOCAB_SIZE
+ or (min_eos_p is not None and probs[-1] >= min_eos_p)
+ ):
+ # eos found, so break
+ pbar.update(n - pbar_state)
+ break
+ x = torch.cat((x, item_next[None]), dim=1)
+ tot_generated_duration_s += 1 / SEMANTIC_RATE_HZ
+ if max_gen_duration_s is not None and tot_generated_duration_s > max_gen_duration_s:
+ pbar.update(n - pbar_state)
+ break
+ if n == n_tot_steps - 1:
+ pbar.update(n - pbar_state)
+ break
+ del logits, relevant_logits, probs, item_next
+
+ if n > pbar_state:
+ if n > pbar.total:
+ pbar.total = n
+ pbar.update(n - pbar_state)
+ pbar_state = n
+ pbar.total = n
+ pbar.refresh()
+ pbar.close()
+ out = x.detach().cpu().numpy().squeeze()[256 + 256 + 1 :]
+ 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."""
+ 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)
+ item_next = torch.multinomial(probs, num_samples=1).to(torch.int32)
+ item_next += logit_start_idx
+ x_coarse_in = torch.cat((x_coarse_in, item_next[None]), dim=1)
+ x_in = torch.cat((x_in, item_next[None]), dim=1)
+ del logits, relevant_logits, probs, item_next
+ n_step += 1
+ del x_in
+ del x_semantic_in
+ if OFFLOAD_CPU:
+ model.to("cpu")
+ gen_coarse_arr = x_coarse_in.detach().cpu().numpy().squeeze()[len(x_coarse_history) :]
+ del x_coarse_in
+ assert len(gen_coarse_arr) == n_steps
+ gen_coarse_audio_arr = gen_coarse_arr.reshape(-1, N_COARSE_CODEBOOKS).T - SEMANTIC_VOCAB_SIZE
+ for n in range(1, N_COARSE_CODEBOOKS):
+ gen_coarse_audio_arr[n, :] -= n * CODEBOOK_SIZE
+ _clear_cuda_cache()
+ return gen_coarse_audio_arr
+
+
+def generate_fine(
+ 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)
+ codebook_preds = torch.multinomial(
+ probs[rel_start_fill_idx:1024], num_samples=1
+ ).reshape(-1)
+ codebook_preds = codebook_preds.to(torch.int32)
+ in_buffer[0, rel_start_fill_idx:, nn] = codebook_preds
+ del logits, codebook_preds
+ # transfer over info into model_in and convert to numpy
+ for nn in range(n_coarse, N_FINE_CODEBOOKS):
+ in_arr[
+ start_fill_idx : start_fill_idx + (1024 - rel_start_fill_idx), nn
+ ] = in_buffer[0, rel_start_fill_idx:, nn]
+ del in_buffer
+ gen_fine_arr = in_arr.detach().cpu().numpy().squeeze().T
+ del in_arr
+ if OFFLOAD_CPU:
+ model.to("cpu")
+ gen_fine_arr = gen_fine_arr[:, n_history:]
+ if n_remove_from_end > 0:
+ gen_fine_arr = gen_fine_arr[:, :-n_remove_from_end]
+ assert gen_fine_arr.shape[-1] == x_coarse_gen.shape[-1]
+ _clear_cuda_cache()
+ return gen_fine_arr
+
+
+def codec_decode(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
diff --git a/lib/bark/model.py b/lib/bark/model.py
new file mode 100644
index 0000000..457b49e
--- /dev/null
+++ b/lib/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)
diff --git a/lib/bark/model_fine.py b/lib/bark/model_fine.py
new file mode 100644
index 0000000..e51d3e6
--- /dev/null
+++ b/lib/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 >= 2.0
+ self.flash = (
+ hasattr(torch.nn.functional, "scaled_dot_product_attention")
+ )
+
+ 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
diff --git a/lib/bark/readme.md b/lib/bark/readme.md
new file mode 100644
index 0000000..516b0c9
--- /dev/null
+++ b/lib/bark/readme.md
@@ -0,0 +1,4 @@
+# Bark
+Suno AI's Bark source code can be found at https://github.com/suno-ai/bark/
+
+Bark is licensed under an MIT license.
\ No newline at end of file
diff --git a/lib/general/__init__.py b/lib/general/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/lib/general/model_downloader.py b/lib/general/model_downloader.py
new file mode 100644
index 0000000..7a87347
--- /dev/null
+++ b/lib/general/model_downloader.py
@@ -0,0 +1,33 @@
+import os
+
+from huggingface_hub import hf_hub_download
+
+import folder_paths
+from comfy.utils import load_torch_file
+
+
+def add_model_folder_path(type_name: str):
+ folder_paths.add_model_folder_path(type_name, os.path.join(folder_paths.models_dir, type_name))
+
+def get_model_path_or_none(type_name: str, model_name: str) -> str | None:
+ return folder_paths.get_full_path(type_name, model_name)
+
+def get_target_model_path(type_name: str, model_name: str) -> str:
+ """
+ Should output something like comfyui/models/type_name/model_name
+ :param type_name: Model type name
+ :param model_name: Model file name
+ :return: The full path
+ """
+ return os.path.join(folder_paths.models_dir, type_name, model_name)
+
+def download_model_file(type_name: str, model_name: str, repo: str, file: str) -> str:
+ """
+ Download a single model file.
+ :param type_name: The model type name to use, also a subdirectory
+ :param model_name: The model name to save as
+ :param repo: The hugginface repository to download the file from.
+ :param file: The file to download from the repository.
+ :return: The path to the downloaded model (prefers local if present)
+ """
+ hf_hub_download(repo, file, local_dir=get_target_model_path(type_name, model_name))
diff --git a/nodes/__init__.py b/nodes/__init__.py
new file mode 100644
index 0000000..14fba45
--- /dev/null
+++ b/nodes/__init__.py
@@ -0,0 +1,4 @@
+from .bark import nodes as bark_nodes
+from .general import nodes as general_nodes
+
+nodes = bark_nodes + general_nodes
\ No newline at end of file
diff --git a/nodes/bark/__init__.py b/nodes/bark/__init__.py
new file mode 100644
index 0000000..2cc69c9
--- /dev/null
+++ b/nodes/bark/__init__.py
@@ -0,0 +1,3 @@
+from .bark_loader import nodes as loader_nodes
+
+nodes = loader_nodes
\ No newline at end of file
diff --git a/nodes/bark/bark_loader.py b/nodes/bark/bark_loader.py
new file mode 100644
index 0000000..1abe296
--- /dev/null
+++ b/nodes/bark/bark_loader.py
@@ -0,0 +1,33 @@
+from comfy_execution.graph_utils import GraphBuilder
+from ...nodes.basenode import BaseNode
+from ...util.categories import category
+
+class BarkLoader(BaseNode):
+ name = "bark_loader"
+ display_name = "Load bark model(s)"
+
+ RETURN_NAMES = ("text", "coarse", "fine")
+ RETURN_TYPES = ("BarkTextModel", "BarkCoarseModel", "BarkFineModel")
+ CATEGORY = category("bark")
+ FUNCTION = "load"
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "text": (["large", "small", "none"],),
+ "coarse": (["large", "small", "none"],),
+ "fine": (["large", "small", "none"],),
+ "half": ("BOOLEAN",),
+ "offload_text": ("BOOLEAN",),
+ "offload_coarse": ("BOOLEAN",),
+ "offload_fine": ("BOOLEAN",)
+ }
+ }
+
+ def load(self, text, coarse, fine, half, offload_text, offload_coarse, offload_fine):
+ # Using https://huggingface.co/GitMylo/bark-safetensors since it's recommended to use safetensors
+
+ return None, None, None
+
+nodes = [BarkLoader]
\ No newline at end of file
diff --git a/nodes/basenode.py b/nodes/basenode.py
new file mode 100644
index 0000000..c00d0fb
--- /dev/null
+++ b/nodes/basenode.py
@@ -0,0 +1,31 @@
+from abc import abstractmethod
+
+from ..util.categories import category
+
+
+class BaseNode:
+ name = "Node"
+ display_name = name
+
+ @classmethod
+ @abstractmethod
+ def INPUT_TYPES(cls):
+ """
+ see https://docs.comfy.org/custom-nodes/backend/datatypes
+
+ return {
+ "required": {},
+ "optional": {},
+ "hidden": {}
+ }
+ """
+ return {
+ "required": {}, # (type, options)
+ "optional": {},
+ "hidden": {}
+ } # required, optional, hidden
+
+ RETURN_TYPES = ()
+ RETURN_NAMES = ()
+ FUNCTION = ""
+ CATEGORY = category("misc")
\ No newline at end of file
diff --git a/nodes/general/__init__.py b/nodes/general/__init__.py
new file mode 100644
index 0000000..29e7ca9
--- /dev/null
+++ b/nodes/general/__init__.py
@@ -0,0 +1,3 @@
+from .encodec import nodes as encodec_nodes
+
+nodes = encodec_nodes
\ No newline at end of file
diff --git a/nodes/general/encodec.py b/nodes/general/encodec.py
new file mode 100644
index 0000000..be8e55e
--- /dev/null
+++ b/nodes/general/encodec.py
@@ -0,0 +1,26 @@
+from ...util.categories import category
+from ...nodes.basenode import BaseNode
+
+class EncodecLoader(BaseNode):
+ name = "encodec_loader"
+ display_name = "Load encodec model"
+
+ RETURN_NAMES = ("model",)
+ RETURN_TYPES = ("EncodecModel",)
+ CATEGORY = category("general")
+ FUNCTION = "load"
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "offload": ("BOOLEAN",)
+ }
+ }
+
+ def load(self, offload):
+ # Using https://huggingface.co/GitMylo/bark-safetensors since it's recommended to use safetensors anyway
+ return (None,)
+
+
+nodes = [EncodecLoader]
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..7586a63
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,2 @@
+funcy
+encodec
\ No newline at end of file
diff --git a/util/__init__.py b/util/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/util/categories.py b/util/categories.py
new file mode 100644
index 0000000..8486065
--- /dev/null
+++ b/util/categories.py
@@ -0,0 +1,5 @@
+def format_name(name: str, emoji: str = "🔉") -> str:
+ return emoji + " " + name
+
+def category(name: str) -> str:
+ return format_name("AudioNodes/" + name)