diff --git a/__init__.py b/__init__.py index ad99a34..be0626a 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,10 @@ -from .lib.general.model_downloader import add_model_folder_path +import os.path + +import folder_paths +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 # Node imports go here from .nodes import nodes diff --git a/lib/general/model_downloader.py b/lib/general/model_downloader.py index a915207..f59d466 100644 --- a/lib/general/model_downloader.py +++ b/lib/general/model_downloader.py @@ -6,8 +6,10 @@ 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 add_model_folder_path(type_name: str, path = None): + if path is None: + path = os.path.join(folder_paths.models_dir, type_name) + folder_paths.add_model_folder_path(type_name, path) def get_model_path_or_none(type_name: str, model_name: str) -> str | None: return folder_paths.get_full_path(type_name, model_name) diff --git a/nodes/bark/bark_generation.py b/nodes/bark/bark_generation.py index 723f6a5..d412291 100644 --- a/nodes/bark/bark_generation.py +++ b/nodes/bark/bark_generation.py @@ -401,7 +401,6 @@ class GenerateFine(BaseNode): # ), f"{isinstance(coarse_codebooks, np.ndarray)}, {len(coarse_codebooks.shape)}, }" if fine_history is not None: x_fine_history = fine_history - print(f"fine history: {x_fine_history.shape}") assert ( isinstance(x_fine_history, np.ndarray) and len(x_fine_history.shape) == 2 diff --git a/nodes/bark/bark_loader.py b/nodes/bark/bark_loader.py index d99a2f9..974277f 100644 --- a/nodes/bark/bark_loader.py +++ b/nodes/bark/bark_loader.py @@ -1,15 +1,18 @@ import dataclasses import json +import os from time import sleep from typing import Literal from dataclasses import dataclass +import numpy as np import safetensors.torch import torch from transformers import BertTokenizer import comfy.utils +import folder_paths from comfy.model_base import BaseModel from comfy.model_patcher import ModelPatcher from ...lib.bark.model import GPTConfig, GPT @@ -112,8 +115,55 @@ class BarkLoader(BaseNode): model.load() bar.update(1) - # TODO: Ensure the models are properly registered so they can be cleared automatically when not needed. - return tuple(load_models) -nodes = [BarkLoader] \ No newline at end of file +class LoadSpeakerNPZ(BaseNode): + name = "load_speaker_npz" + display_name = "Load bark speaker .npz" + + RETURN_NAMES = ("text", "coarse", "fine") + RETURN_TYPES = ("BarkSemanticTokens", "EncodecCodeBooks", "EncodecCodeBooks") + CATEGORY = category("bark") + FUNCTION = "load" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "file": (folder_paths.get_filename_list("bark_speakers"),) + } + } + + def load(self, file: str): + speaker_path = folder_paths.get_full_path_or_raise("bark_speakers", file) + speaker = np.load(speaker_path) + return speaker["semantic_prompt"], speaker["coarse_prompt"], speaker["fine_prompt"] + +class SaveSpeakerNPZ(BaseNode): + name = "save_speaker_npz" + display_name = "Save bark speaker .npz" + + RETURN_NAMES = () + RETURN_TYPES = () + CATEGORY = category("bark") + FUNCTION = "save" + + OUTPUT_NODE = True + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "name": ("STRING", {"default": "speaker"}), + "semantic_prompt": ("BarkSemanticTokens",), + "coarse_prompt": ("EncodecCodeBooks",), + "fine_prompt": ("EncodecCodeBooks",) + } + } + + def save(self, name: str, semantic_prompt, coarse_prompt, fine_prompt): + speaker_path = os.path.join(folder_paths.models_dir, "bark", "speakers", name) + 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 diff --git a/pyproject.toml b/pyproject.toml index 2db3db2..4ba12b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "audio-nodes" description = "Various nodes related to audio." -version = "1.0.2" +version = "1.0.3" license = {file = "LICENSE"} dependencies = ["funcy", "encodec"] diff --git a/readme.md b/readme.md index 3b48bed..0a6b77e 100644 --- a/readme.md +++ b/readme.md @@ -9,6 +9,7 @@ Current nodes: 1. Semantic 2. Coarse 3. Fine + * Load/save speaker npz, put speaker .npz files in `models/bark/speakers` * Encodec * Load models * Encode audio @@ -16,11 +17,8 @@ Current nodes: # TODO * Bark - * Load/save speaker npz - * Use model folder (`models/bark/speakers`) for loading/saving voices * Hubert quantizer for creating npz from audio https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer * Coarse-ify codebooks for creating custom npz files * Direct way to support long-form generations * Other - * HuBert model loader - * Concat audio \ No newline at end of file + * HuBert model loader \ No newline at end of file