diff --git a/__init__.py b/__init__.py index be0626a..1376e30 100644 --- a/__init__.py +++ b/__init__.py @@ -5,6 +5,7 @@ from .lib.general.model_downloader import add_model_folder_path add_model_folder_path("bark") # Bark models are downloaded here add_model_folder_path("bark_speakers", os.path.join(folder_paths.models_dir, "bark", "speakers")) # Bark speaker npz files +add_model_folder_path("bark_quantizers", os.path.join(folder_paths.models_dir, "bark", "quantizers")) # Bark Hubert Quantizers # Node imports go here from .nodes import nodes diff --git a/lib/bark/hubert/__init__.py b/lib/bark/hubert/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/lib/bark/hubert/customtokenizer.py b/lib/bark/hubert/customtokenizer.py new file mode 100644 index 0000000..96e40a6 --- /dev/null +++ b/lib/bark/hubert/customtokenizer.py @@ -0,0 +1,200 @@ +""" +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 + + +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 + 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, map_location=map_location)) + 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): + 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') + 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) + try: + prefix = input_file.split("_")[0] + number = int(prefix) + except ValueError as e: + raise e + if input_file.endswith(sem_string): + data_y[number] = numpy.load(full_path) + elif input_file.endswith(feat_string): + data_x[number] = numpy.load(full_path) + + model_training.prepare_training() + epoch = 1 + + while 1: + for i in range(save_epochs): + j = 0 + for i in range(max(len(data_x), len(data_y))): + x = data_x.get(i) + y = data_y.get(i) + if x is None or y is None: + print(f'The training data does not match. key={i}') + continue + model_training.train_step(torch.tensor(x).to('cuda'), torch.tensor(y).to('cuda'), j % 50 == 0) # Print loss every 50 steps + j += 1 + 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 diff --git a/lib/bark/hubert/hubert_manager.py b/lib/bark/hubert/hubert_manager.py new file mode 100644 index 0000000..49fd389 --- /dev/null +++ b/lib/bark/hubert/hubert_manager.py @@ -0,0 +1,21 @@ +import os.path +import shutil + +import huggingface_hub + +import folder_paths + + +class HuBERTManager: + @staticmethod + def make_sure_tokenizer_installed(model: str = 'quantifier_hubert_base_ls960_14.pth', repo: str = 'GitMylo/bark-voice-cloning', local_file: str = 'tokenizer.pth'): + install_dir = os.path.join(folder_paths.models_dir, "bark", "quantizers", local_file) + 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): + print('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 diff --git a/lib/bark/hubert/pre_kmeans_hubert.py b/lib/bark/hubert/pre_kmeans_hubert.py new file mode 100644 index 0000000..0f15dce --- /dev/null +++ b/lib/bark/hubert/pre_kmeans_hubert.py @@ -0,0 +1,115 @@ +""" +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 + +# Modified again in 2025 to use transformers instead of fairseq + +from pathlib import Path + +import torch +import transformers +from torch import nn +from einops import pack, unpack + +# import fairseq + +from torchaudio.functional import resample +from transformers import HubertModel + + +def exists(val): + return val is not None + + +def default(val, d): + return val if exists(val) else d + +# Taken from https://github.com/lucidrains/audiolm-pytorch/blob/main/audiolm_pytorch/utils.py to not add new dependencies +def round_down_nearest_multiple(num, divisor): + return num // divisor * divisor + +def curtail_to_multiple(t, mult, from_left = False): + data_len = t.shape[-1] + rounded_seq_len = round_down_nearest_multiple(data_len, mult) + seq_slice = slice(None, rounded_seq_len) if not from_left else slice(-rounded_seq_len, None) + return t[..., seq_slice] + + +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 + ): + super().__init__() + self.target_sample_hz = target_sample_hz + self.seq_len_multiple_of = seq_len_multiple_of + self.output_layer = output_layer + + # model_path = Path(checkpoint_path) + + # assert model_path.exists(), f'path {checkpoint_path} does not exist' + + # checkpoint = torch.load(checkpoint_path, map_location=device) + # load_model_input = {checkpoint_path: checkpoint} + # model, *_ = fairseq.checkpoint_utils.load_model_ensemble_and_task(load_model_input) + + # TODO: Use the ComfyUI models directory (sorry for now) + self.model: HubertModel = HubertModel.from_pretrained("facebook/hubert-base-ls960") + self.model.eval() + + def to(self, *args, **kwargs): + self.model.to(*args, **kwargs) + + @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.forward( + wav_input, + output_hidden_states=True + # 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 + ).hidden_states[self.output_layer] # TODO: Ensure this is the proper way to get the output layer + + # embed, packed_shape = pack([embed['x']], '* d') + embed, packed_shape = pack([embed], '* 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 diff --git a/nodes/bark/bark_generation.py b/nodes/bark/bark_generation.py index d412291..2f22ff0 100644 --- a/nodes/bark/bark_generation.py +++ b/nodes/bark/bark_generation.py @@ -6,11 +6,12 @@ import numpy import numpy as np import torch import torch.nn.functional as F +from encodec.utils import convert_audio from scipy.special import softmax from comfy.utils import ProgressBar from ...lib.bark.generation import autocast, InferenceContext -from .bark_loader import BarkModelRules +from .bark_loader import BarkModelRules, BarkHubertModelRules, BarkHubertQuantizerModelRules from ...util.categories import category from ...nodes.basenode import BaseNode @@ -491,4 +492,84 @@ class GenerateFine(BaseNode): return (gen_fine_arr,) -nodes = [PromptEncode, GenerateSemantic, GenerateCoarse, GenerateFine] \ No newline at end of file +class HubertEncode(BaseNode): + name = "bark_hubert_encode" + display_name = "HuBERT vectorize audio" + + RETURN_NAMES = ("vectors",) + RETURN_TYPES = ("HuBERTVectors",) + CATEGORY = category("bark/cloning") + FUNCTION = "encode" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "hubert_model": ("BarkHuBERTModel",), + "audio": ("AUDIO",) + } + } + + def encode(self, hubert_model: BarkHubertModelRules, audio): + device = hubert_model.patcher.load_device + hubert_model.patcher.load(device) + model = hubert_model.model + + wav: torch.Tensor = audio["waveform"] + sr: int = audio["sample_rate"] + wav = convert_audio(wav, sr, model.target_sample_hz, 1) + wav = wav.to(device)[0] # shape was [1, 1, 1, 179840] + + vecs = model.forward(wav, input_sample_hz=sr) + + return (vecs,) + +class HubertQuantize(BaseNode): + name = "bark_hubert_quantize" + display_name = "quantize HuBERT vectors to bark semantic tokens" + + RETURN_NAMES = ("semantic_tokens",) + RETURN_TYPES = ("BarkSemanticTokens",) + CATEGORY = category("bark/cloning") + FUNCTION = "quantize" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "quantizer_model": ("BarkHuBERTQuantizerModel",), + "vectors": ("HuBERTVectors",) + } + } + + def quantize(self, quantizer_model: BarkHubertQuantizerModelRules, vectors: torch.Tensor): + device = quantizer_model.patcher.load_device + quantizer_model.patcher.load(device) + model = quantizer_model.model + + vectors = vectors.to(device) + tokens = model.get_token(vectors) + + return (tokens,) + +class EncodecCourseify(BaseNode): + name = "bark_encodec_coarseify" + display_name = "Course-ify encodec codebooks" + + RETURN_NAMES = ("coarse_codebooks",) + RETURN_TYPES = ("EncodecCodeBooks",) + CATEGORY = category("bark") + FUNCTION = "encode" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "fine_codebooks": ("EncodecCodeBooks",) + } + } + + def encode(self, fine_codebooks): + return (fine_codebooks[:2, :],) + +nodes = [PromptEncode, GenerateSemantic, GenerateCoarse, GenerateFine, HubertEncode, HubertQuantize, EncodecCourseify] \ No newline at end of file diff --git a/nodes/bark/bark_loader.py b/nodes/bark/bark_loader.py index 974277f..1669818 100644 --- a/nodes/bark/bark_loader.py +++ b/nodes/bark/bark_loader.py @@ -9,12 +9,15 @@ from dataclasses import dataclass import numpy as np import safetensors.torch import torch -from transformers import BertTokenizer +from transformers import BertTokenizer, HubertModel import comfy.utils import folder_paths from comfy.model_base import BaseModel from comfy.model_patcher import ModelPatcher +from lib.bark.hubert.hubert_manager import HuBERTManager +from ...lib.bark.hubert.customtokenizer import CustomTokenizer +from ...lib.bark.hubert.pre_kmeans_hubert import CustomHubert from ...lib.bark.model import GPTConfig, GPT from ...lib.bark.model_fine import FineGPTConfig, FineGPT from ...lib.general.model_downloader import download_model_file @@ -166,4 +169,79 @@ class SaveSpeakerNPZ(BaseNode): np.savez(speaker_path, semantic_prompt=semantic_prompt, coarse_prompt=coarse_prompt, fine_prompt=fine_prompt) return () -nodes = [BarkLoader, LoadSpeakerNPZ, SaveSpeakerNPZ] \ No newline at end of file +@dataclass +class BarkHubertModelRules: + cpu: bool + + model: CustomHubert = None + patcher: ModelPatcher = None + + def load(self): + main_device = mm.get_torch_device() if not self.cpu else torch.device("cpu") + offload_device = mm.unet_offload_device() + + self.model = CustomHubert() + self.patcher = ModelPatcher(self.model, main_device, offload_device) + +class LoadHubertBark(BaseNode): + name = "load_bark_hubert" + display_name = "Load HuBERT model (bark)" + + RETURN_NAMES = ("model",) + RETURN_TYPES = ("BarkHuBERTModel",) + CATEGORY = category("bark/cloning") + FUNCTION = "load" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # "cpu": ("BOOLEAN",) + } + } + + def load(self, cpu: bool = True): + model = BarkHubertModelRules(cpu) + model.load() + return (model,) + +@dataclass +class BarkHubertQuantizerModelRules: + cpu: bool + model: str + + model: CustomTokenizer = None + patcher: ModelPatcher = None + + def load(self): + main_device = mm.get_torch_device() if not self.cpu else torch.device("cpu") + offload_device = mm.unet_offload_device() + + model_file = HuBERTManager.make_sure_tokenizer_installed(self.model) + self.model = CustomTokenizer.load_from_checkpoint(model_file) + self.patcher = ModelPatcher(self.model, main_device, offload_device) + +class LoadHubertQuantizerBark(BaseNode): + name = "load_bark_hubert_quantizer" + display_name = "Load HuBERT quantizer model (bark)" + + RETURN_NAMES = ("model",) + RETURN_TYPES = ("BarkHuBERTQuantizerModel",) + CATEGORY = category("bark/cloning") + FUNCTION = "load" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # "cpu": ("BOOLEAN",) + "model": (["quantifier_hubert_base_ls960.pth", "quantifier_hubert_base_ls960_14.pth", "quantifier_V1_hubert_base_ls960_23.pth"],) # TODO: Load from online/local registry + } + } + + def load(self, model, cpu: bool = True): + model = BarkHubertQuantizerModelRules(cpu, model) + model.load() + return (model,) + +nodes = [BarkLoader, LoadSpeakerNPZ, SaveSpeakerNPZ, LoadHubertBark, LoadHubertQuantizerBark] \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 4ba12b8..3c2dd91 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "audio-nodes" description = "Various nodes related to audio." -version = "1.0.3" +version = "1.0.4" license = {file = "LICENSE"} dependencies = ["funcy", "encodec"] diff --git a/readme.md b/readme.md index 0a6b77e..d871666 100644 --- a/readme.md +++ b/readme.md @@ -10,6 +10,8 @@ Current nodes: 2. Coarse 3. Fine * Load/save speaker npz, put speaker .npz files in `models/bark/speakers` + * HuBert model loader + * Hubert quantizer for creating npz from audio https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer * Encodec * Load models * Encode audio @@ -17,8 +19,6 @@ Current nodes: # TODO * Bark - * Hubert quantizer for creating npz from audio https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer + * Add an option to select the final layer for the hubert quantizer (similar to clip skip, but for hubert, so let's call it hubert skip) * Coarse-ify codebooks for creating custom npz files - * Direct way to support long-form generations -* Other - * HuBert model loader \ No newline at end of file + * Direct way to support long-form generations \ No newline at end of file