Support loading/saving of speaker .npz files

__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

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)

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

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]
+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]

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"]
 

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
+  * HuBert model loader