v4.1 - hash and metadata cleanup (rsxdalv#23)

* remove file names from metadata

* bugfix shadowed variables

* rename constants to uppercase

* split out bark css

* convert choices to classes

* improve setting class names

* split settings into a new file

* add sourcery to gitignore

* define FinalGenParams class

* improve types and log_generation API

* fix get_long_gen_history_prompt bug

* fix types, reformat to black

* add missing requirements

* modify hash function to improve consistency

* update readme

Author: rsxdalv

.gitignore

@@ -22,3 +22,4 @@ node_modules/
 # Editors
 .vscode/
 .idea/
+.sourcery.yaml

README.md

@@ -35,6 +35,15 @@ This code requires the following dependencies:
 
 
 ## Changelog
+June 4:
+* Update to v4.1 - improved hash function, code improvements
+
+June 3:
+* Update to v4 - new output structure, improved history view, codebase reorganization, improved metadata, output extensions support
+
+May __:
+* Update to v3 - voice clone demo
+
 May 17:
 * Update to v2 - generate results as they appear, preview long prompt generations piece by piece, enable up to 9 outputs, UI tweaks
 

requirements.txt

@@ -1,2 +1,6 @@
 # pip install nodejs-bin==16.15.1a4
 ffmpeg-python # Apache 2.0
+gradio==3.33.1
+python-dotenv==1.0.0
+soundfile==0.12.1 # torchaudio platform windows
+# run_cmd("pip install sox") # torchaudio platform linux

server.py

@@ -21,28 +21,29 @@
 
 
 def load_models(
+    text_use_gpu,
+    text_use_small,
+    coarse_use_gpu,
+    coarse_use_small,
+    fine_use_gpu,
+    fine_use_small,
+    codec_use_gpu,
+):
+    save_config_bark(
         text_use_gpu,
         text_use_small,
         coarse_use_gpu,
         coarse_use_small,
         fine_use_gpu,
         fine_use_small,
-        codec_use_gpu
-):
-    save_config_bark(text_use_gpu,
-                     text_use_small,
-                     coarse_use_gpu,
-                     coarse_use_small,
-                     fine_use_gpu,
-                     fine_use_small,
-                     codec_use_gpu)
+        codec_use_gpu,
+    )
     # download and load all models
     # TODO: try catch for memory errors
     model_manager.reload_models(config)
     return gr.Button.update(value="Reload models", interactive=True)
 
 
-
 def reload_config_and_restart_ui():
     os._exit(0)
     # print("Reloading config and restarting UI...")
@@ -53,8 +54,11 @@ def reload_config_and_restart_ui():
     # demo.launch(**gradio_interface_options)
 
 
-gradio_interface_options = config[
-    "gradio_interface_options"] if "gradio_interface_options" in config else default_config
+gradio_interface_options = (
+    config["gradio_interface_options"]
+    if "gradio_interface_options" in config
+    else default_config
+)
 
 with gr.Blocks(css=full_css) as demo:
     gr.Markdown("# TTS Generation WebUI (Bark & Tortoise)")
@@ -68,7 +72,9 @@ def reload_config_and_restart_ui():
         voices_tab(register_use_as_history_button)
 
         settings_tab_bark(config, save_config_bark, load_models)
-        settings_tab_gradio(save_config_gradio, reload_config_and_restart_ui, gradio_interface_options)
+        settings_tab_gradio(
+            save_config_gradio, reload_config_and_restart_ui, gradio_interface_options
+        )
 
 
 def print_pretty_options(options):

src/bark/BarkModelManager.py

@@ -20,10 +20,11 @@ def reload_models(self, config):
         codec_use_gpu = model_config["codec_use_gpu"]
 
         print(
-            f'''\t- Text Generation:\t\t GPU: {"Yes" if text_use_gpu else "No"}, Small Model: {"Yes" if text_use_small else "No"}
+            f"""\t- Text Generation:\t\t GPU: {"Yes" if text_use_gpu else "No"}, Small Model: {"Yes" if text_use_small else "No"}
 \t- Coarse-to-Fine Inference:\t GPU: {"Yes" if coarse_use_gpu else "No"}, Small Model: {"Yes" if coarse_use_small else "No"}
 \t- Fine-tuning:\t\t\t GPU: {"Yes" if fine_use_gpu else "No"}, Small Model: {"Yes" if fine_use_small else "No"}
-\t- Codec:\t\t\t GPU: {"Yes" if codec_use_gpu else "No"}''')
+\t- Codec:\t\t\t GPU: {"Yes" if codec_use_gpu else "No"}"""
+        )
 
         preload_models(
             text_use_gpu=text_use_gpu,

src/bark/FinalGenParams.py

@@ -0,0 +1,9 @@
+from typing import Any, TypedDict
+
+
+class FinalGenParams(TypedDict):
+    text: str
+    history_prompt: str | Any
+    text_temp: float
+    waveform_temp: float
+    output_full: bool

src/bark/bark_css.py

@@ -0,0 +1,15 @@
+bark_css = """
+.generating {
+    pointer-events: none;
+}
+.tts-audio > .empty.small.unpadded_box {
+    padding: var(--size-2);
+    width: var(--size-full);
+    height: var(--size-14);
+    min-height: var(--size-14);
+}
+.tts-image > .empty.large.unpadded_box {
+    height: 150px;
+    min-height: 150px;
+}
+"""

src/bark/clone/tab_voice_clone_demo.py

@@ -1,9 +1,15 @@
 import tempfile
 from src.bark.npz_tools import save_npz
 from src.bark.FullGeneration import FullGeneration
-from models.bark_voice_cloning_hubert_quantizer.hubert.hubert_manager import HuBERTManager
-from models.bark_voice_cloning_hubert_quantizer.hubert.pre_kmeans_hubert import CustomHubert
-from models.bark_voice_cloning_hubert_quantizer.hubert.customtokenizer import CustomTokenizer
+from models.bark_voice_cloning_hubert_quantizer.hubert.hubert_manager import (
+    HuBERTManager,
+)
+from models.bark_voice_cloning_hubert_quantizer.hubert.pre_kmeans_hubert import (
+    CustomHubert,
+)
+from models.bark_voice_cloning_hubert_quantizer.hubert.customtokenizer import (
+    CustomTokenizer,
+)
 import torchaudio
 import torch
 from encodec.utils import convert_audio
@@ -46,7 +52,9 @@ def _load_tokenizer():
     tokenizer_path = HuBERTManager.make_sure_tokenizer_installed()
     global tokenizer
     if tokenizer is None:
-        tokenizer = CustomTokenizer.load_from_checkpoint('data/models/hubert/tokenizer.pth')
+        tokenizer = CustomTokenizer.load_from_checkpoint(
+            "data/models/hubert/tokenizer.pth"
+        )
         tokenizer.load_state_dict(torch.load(tokenizer_path))
     return tokenizer
 
@@ -76,7 +84,7 @@ def get_prompts(path_to_wav: str, use_gpu: bool):
 
 
 def get_encodec_prompts(path_to_wav: str, use_gpu=True):
-    device = 'cuda' if use_gpu else 'cpu'
+    device = "cuda" if use_gpu else "cpu"
     model: EncodecModel = load_codec_model(use_gpu=use_gpu)
     wav, sr = torchaudio.load(path_to_wav)
     wav = convert_audio(wav, sr, model.sample_rate, model.channels)
@@ -87,37 +95,41 @@ def get_encodec_prompts(path_to_wav: str, use_gpu=True):
     with torch.no_grad():
         encoded_frames = model.encode(wav)
 
-    fine_prompt: np.ndarray = torch.cat(
-        [encoded[0] for encoded in encoded_frames], dim=-1).squeeze().cpu().numpy()
+    fine_prompt: np.ndarray = (
+        torch.cat([encoded[0] for encoded in encoded_frames], dim=-1)
+        .squeeze()
+        .cpu()
+        .numpy()
+    )
     coarse_prompt = fine_prompt[:2, :]
     return fine_prompt, coarse_prompt
 
 
 def save_cloned_voice(
     full_generation: FullGeneration,
 ):
-    voice_name = f'test_clone_voice{str(np.random.randint(100000))}'
-    filename = f'voices/{voice_name}.npz'
+    voice_name = f"test_clone_voice{str(np.random.randint(100000))}"
+    filename = f"voices/{voice_name}.npz"
     save_npz(filename, full_generation)
     return filename
 
 
 def tab_voice_clone_demo():
     with gr.Tab("Bark Voice Clone Demo"):
-        gr.Markdown("""
+        gr.Markdown(
+            """
         Unethical use of this technology is prohibited.
         This demo is based on https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer repository.
-        """)
+        """
+        )
 
         # TODO: try with ffmpeg (except mp3)
         # file_input = gr.Audio(label="Input Audio", type="numpy", source="upload", interactive=True)
-        file_input = gr.File(label="Input Audio", file_types=[
-            ".wav"], interactive=True)
+        file_input = gr.File(label="Input Audio", file_types=[".wav"], interactive=True)
 
         use_gpu_checkbox = gr.Checkbox(label="Use GPU", value=True)
 
-        generate_voice_button = gr.Button(
-            value="Generate Voice", variant="primary")
+        generate_voice_button = gr.Button(value="Generate Voice", variant="primary")
 
         def generate_voice(wav_file_obj: tempfile._TemporaryFileWrapper, use_gpu: bool):
             if wav_file_obj is None:
@@ -128,8 +140,11 @@ def generate_voice(wav_file_obj: tempfile._TemporaryFileWrapper, use_gpu: bool):
             filename = save_cloned_voice(full_generation)
             return f"Saved: {filename}"
 
-        output = gr.Label(
-            "Output will appear here after input", type="auto")
+        output = gr.Label("Output will appear here after input", type="auto")
 
-        generate_voice_button.click(fn=generate_voice, inputs=[
-            file_input, use_gpu_checkbox], outputs=output, preprocess=True)
+        generate_voice_button.click(
+            fn=generate_voice,
+            inputs=[file_input, use_gpu_checkbox],
+            outputs=output,
+            preprocess=True,
+        )

src/bark/create_voice_string.py

@@ -1,10 +1,12 @@
 from models.bark.bark.generation import SUPPORTED_LANGS
-
+from typing import Union
 import os
 
 
-def create_voice_string(language: str, speaker_id: int, use_v2: bool) -> str:
-    history_prompt = f"{SUPPORTED_LANGS[language][1]}_speaker_{speaker_id}"
+def create_voice_string(
+    language: str, speaker_id: Union[int, str], use_v2: bool
+) -> str:
+    history_prompt = f"{SUPPORTED_LANGS[language][1]}_speaker_{speaker_id}"  # type: ignore
     if use_v2:
         history_prompt = os.path.join("v2", history_prompt)
     return history_prompt

src/bark/generate_and_save_metadata.py

@@ -1,14 +1,30 @@
 import json
 from models.bark.bark.generation import models
+from typing import Optional
 
 
-def generate_and_save_metadata(prompt: str, language: str, speaker_id: int, text_temp: float, waveform_temp: float,
-                               seed: int, filename: str, date: str, filename_png: str, filename_json: str,
-                               history_prompt_npz: str, filename_npz: str, history_prompt: str, history_hash: str):
+def generate_and_save_metadata(
+    prompt: str,
+    language: Optional[str],
+    speaker_id: Optional[int],
+    text_temp: float,
+    waveform_temp: float,
+    seed: int,
+    filename: str,
+    date: str,
+    filename_png: str,
+    filename_json: str,
+    history_prompt_npz: Optional[str],
+    filename_npz: str,
+    history_prompt: str,
+    history_hash: str,
+):
     is_big_semantic_model = models["text"]["model"].config.n_embd > 768
     is_big_coarse_model = models["coarse"].config.n_embd > 768
     is_big_fine_model = models["fine"].config.n_embd > 768
     metadata = {
+        "_version": "0.0.1",
+        "_hash_version": "0.0.2",
         # "id": generation_hash, # generation_hash is the same as history_hash but for current generation
         # "model_semantic_hash": model_semantic_hash,
         "is_big_semantic_model": is_big_semantic_model,
@@ -26,10 +42,6 @@ def generate_and_save_metadata(prompt: str, language: str, speaker_id: int, text
         "waveform_temp": waveform_temp,
         "date": date,
         "seed": str(seed),
-        "filename": filename,
-        "filename_png": filename_png,
-        "filename_json": filename_json,
-        "filename_npz": filename_npz,
         # "files": {
         #     "wav": filename,
         #     "png": filename_png,

src/bark/generate_choice_string.py

@@ -1,12 +1,10 @@
 from src.bark.create_voice_string import create_voice_string
 from src.bark.get_speaker_gender import get_speaker_gender
-
+from typing import Union
 import gradio as gr
 
 
-def generate_choice_string(use_v2: bool, language: str, speaker_id: str):
+def generate_choice_string(use_v2: bool, language: str, speaker_id: Union[int, str]):
     history_prompt = create_voice_string(language, speaker_id, use_v2)
     gender = get_speaker_gender(history_prompt)
-    return gr.Markdown.update(
-        value=f"Chosen voice: {history_prompt}, Gender: {gender}"
-    )
+    return gr.Markdown.update(value=f"Chosen voice: {history_prompt}, Gender: {gender}")

src/bark/generate_random_seed.py

@@ -2,4 +2,4 @@
 
 
 def generate_random_seed() -> int:
-    return np.random.default_rng().integers(1, 2 ** 32 - 1)
+    return np.random.default_rng().integers(1, 2**32 - 1)

src/bark/generation_settings.py

@@ -0,0 +1,34 @@
+class HistorySettings:
+    EMPTY = "Empty history"
+    VOICE = "or Use a voice:"
+    NPZ_FILE = "or Use old generation as history:"
+
+    choices = [
+        EMPTY,
+        VOICE,
+        NPZ_FILE,
+    ]
+
+
+class PromptSplitSettings:
+    NONE = "Short prompt (<15s)"
+    LINES = "Split prompt by lines"
+    LENGTH = "Split prompt by length"
+
+    choices = [
+        NONE,
+        LINES,
+        # LENGTH,
+    ]
+
+
+class LongPromptHistorySettings:
+    CONTINUE = "Use old generation as history"
+    CONSTANT = "or Use history prompt setting"
+    EMPTY = "or Clear history"
+
+    choices = [
+        CONTINUE,
+        CONSTANT,
+        EMPTY,
+    ]