decorator upgrades, fixes, UI improvements (#517)

rsxdalv · web-flow · commit ba626d16b68c · 2025-06-04T00:51:25.000+03:00
* add sample voices

* rank up chatterbox extension in the order

* change resemble enhance gitignore pattern

* import text splitter to webui

* improve Gradio random seed UI

* improve model management API

* improve generator compatibility with extensions

* add interactivity to model unloading button

* attempt Dockerfile pip fix

* add chunked generations to chatterbox React UI

* voice dropdown stub

* readme

* add save_wav accumulated
diff --git a/.gitignore b/.gitignore
@@ -36,7 +36,7 @@ extensions.external.json
 /data/models/parler_tts/
 /data/models/musicgen_audiogen/*/
 /data/models/gpt_sovits/
-/data/models/resemble_enhance/ds/
+/data/models/resemble_enhance/
 /data/models/ap_bwe/weights/
 /data/models/openvoice/
 /data/models/openvoice_v2/
diff --git a/Dockerfile b/Dockerfile
@@ -24,6 +24,7 @@ RUN npm --version
 ENV TORCH_VERSION=2.6.0
 
 ENV PATH="/root/.cargo/bin:$PATH"
+RUN pip install --no-cache-dir --upgrade pip setuptools wheel
 RUN pip install --no-cache-dir setuptools torch==$TORCH_VERSION torchvision torchaudio 
 
 # Set working directory
diff --git a/README.md b/README.md
@@ -86,10 +86,17 @@
 
 ## Changelog
 
+June 4:
+* Attempt dockerfile fix.
+* Add interactivity to model unloading button, improve Gradio random seed UI.
+* Add sample voices.
+
 June 1:
 * Add presets API.
 * Add API Preset config to React UI.
 
+## May 2025
+
 May 31:
 * Improve React UI Audio player.
 * Fix ROCm installation version.
diff --git a/extensions.json b/extensions.json
@@ -90,6 +90,20 @@
             "extension_website": "https://github.com/rsxdalv/extension_f5_tts",
             "extension_platform_version": "0.0.1"
         },
+        {
+            "package_name": "extension_chatterbox",
+            "name": "Chatterbox",
+            "requirements": "git+https://github.com/rsxdalv/extension_chatterbox@main",
+            "description": "Chatterbox, Resemble AI's first production-grade open source TTS model",
+            "extension_type": "interface",
+            "extension_class": "text-to-speech",
+            "author": "Resemble AI",
+            "extension_author": "rsxdalv",
+            "license": "MIT",
+            "website": "https://github.com/resemble-ai/chatterbox",
+            "extension_website": "https://github.com/rsxdalv/extension_chatterbox",
+            "extension_platform_version": "0.0.1"
+        },
         {
             "package_name": "extension_bark",
             "recommended": true,
@@ -656,20 +670,6 @@
             "website": "https://github.com/rsxdalv/extension_audiobook_generator",
             "extension_website": "https://github.com/rsxdalv/extension_audiobook_generator",
             "extension_platform_version": "0.0.1"
-        },
-        {
-            "package_name": "extension_chatterbox",
-            "name": "Chatterbox",
-            "requirements": "git+https://github.com/rsxdalv/extension_chatterbox@main",
-            "description": "Chatterbox, Resemble AI's first production-grade open source TTS model",
-            "extension_type": "interface",
-            "extension_class": "text-to-speech",
-            "author": "Resemble AI",
-            "extension_author": "rsxdalv",
-            "license": "MIT",
-            "website": "https://github.com/resemble-ai/chatterbox",
-            "extension_website": "https://github.com/rsxdalv/extension_chatterbox",
-            "extension_platform_version": "0.0.1"
         }
     ],
     "decorators": [
diff --git a/react-ui/src/components/ChatterboxInputs.tsx b/react-ui/src/components/ChatterboxInputs.tsx
@@ -15,6 +15,8 @@ import {
 import { UnloadModelButton } from "./component/ModelDropdown";
 import { AudioInput } from "./AudioComponents";
 import { SeedInput } from "./SeedInput";
+import { Switch } from "./ui/switch";
+import { Label } from "./ui/label";
 
 // Type alias for components that need basic params
 type ChatterboxBasicParams = Omit<ChatterboxParams, "audio_prompt_path"> & {
@@ -75,7 +77,7 @@ export const ChatterboxInputs = ({
             onChange={handleChange}
             label="CFG Weight/Pace"
             name="cfg_weight"
-            min="0.2"
+            min="0"
             max="1"
             step="0.05"
             decimals={2}
@@ -96,6 +98,13 @@ export const ChatterboxInputs = ({
           />
         </div>
 
+        {/* add voice dropdown */}
+        {/* <div className="pt-2">
+          <VoiceDropdown
+            params={chatterboxParams}
+            handleChange={handleChange}
+          />
+        </div> */}
         <div className="pt-2">
           <AudioInput
             url={chatterboxParams.audio_prompt_path?.path}
@@ -110,6 +119,22 @@ export const ChatterboxInputs = ({
           />
         </div>
 
+        <div className="flex items-center space-x-2">
+          <Switch
+            id="chunked"
+            checked={chatterboxParams.chunked}
+            name="chunked"
+            onCheckedChange={(x) =>
+              handleChange({ target: { name: "chunked", value: x } })
+            }
+          />
+          <div>
+            <Label htmlFor="chunked">Chunked</Label>
+            <CardDescription>
+              Enable chunked generation for longer prompts.
+            </CardDescription>
+          </div>
+        </div>
         <div className="flex items-center gap-4 pt-2">
           <RadioWithLabel
             label="Device"
diff --git a/react-ui/src/tabs/ChatterboxParams.ts b/react-ui/src/tabs/ChatterboxParams.ts
@@ -27,6 +27,7 @@ export interface ChatterboxParams extends Seeded {
   device: string;
   dtype: string;
   model_name: string;
+  chunked: boolean;
 }
 
 export const defaultChatterboxParams: ChatterboxParams = {
@@ -40,6 +41,7 @@ export const defaultChatterboxParams: ChatterboxParams = {
   model_name: "just_a_placeholder",
   seed: -1,
   use_random_seed: true,
+  chunked: false,
 };
 
 export type ChatterboxResult = {
diff --git a/tts_webui/decorators/__init__.py b/tts_webui/decorators/__init__.py
@@ -3,7 +3,7 @@
 from .decorator_add_model_type import decorator_add_model_type, decorator_add_model_type_generator
 from .decorator_apply_torch_seed import decorator_apply_torch_seed, decorator_apply_torch_seed_generator
 from .decorator_log_generation import decorator_log_generation, decorator_log_generation_generator
-from .decorator_save_metadata import decorator_save_metadata
+from .decorator_save_metadata import decorator_save_metadata, decorator_save_metadata_generator
 from .decorator_save_musicgen_npz import decorator_save_musicgen_npz
 from .decorator_save_wav import decorator_save_wav, decorator_save_wav_generator
 from .gradio_dict_decorator import dictionarize, dictionarize_wraps
diff --git a/tts_webui/decorators/decorator_save_metadata.py b/tts_webui/decorators/decorator_save_metadata.py
@@ -3,31 +3,47 @@
 from tts_webui.utils.outputs.path import get_relative_output_path_ext
 
 
+def _save_metadata_to_result(result_dict, kwargs):
+    """Helper function to save metadata to a result dictionary."""
+    path = get_relative_output_path_ext(result_dict, ".json")
+    print("Saving metadata to", path)
+
+    metadata = {
+        "_version": "0.0.1",
+        "_hash_version": "0.0.2",
+        **kwargs,
+        "outputs": None,
+        "date": str(result_dict["date"]),
+        "hash": audio_array_to_sha256(result_dict["audio_out"][1]),
+        # **result_dict,
+    }
+
+    with open(path, "w") as outfile:
+        json.dump(
+            metadata,
+            outfile,
+            indent=2,
+            skipkeys=True,
+            default=lambda o: f"<<non-serializable: {type(o).__qualname__}>>",
+        )
+
+    result_dict["metadata"] = metadata
+    return result_dict
+
+
 def decorator_save_metadata(fn):
     def wrapper(*args, **kwargs):
         result_dict = fn(*args, **kwargs)
-        path = get_relative_output_path_ext(result_dict, ".json")
-        print("Saving metadata to", path)
-
-        metadata = {
-            "_version": "0.0.1",
-            "_hash_version": "0.0.2",
-            **kwargs,
-            "outputs": None,
-            "date": str(result_dict["date"]),
-            "hash": audio_array_to_sha256(result_dict["audio_out"][1]),
-            # **result_dict,
-        }
-        with open(path, "w") as outfile:
-            json.dump(
-                metadata,
-                outfile,
-                indent=2,
-                skipkeys=True,
-                default=lambda o: f"<<non-serializable: {type(o).__qualname__}>>",
-            )
-
-        result_dict["metadata"] = metadata
-        return result_dict
+        return _save_metadata_to_result(result_dict, kwargs)
+
+    return wrapper
+
+
+def decorator_save_metadata_generator(fn):
+    def wrapper(*args, **kwargs):
+        for result_dict in fn(*args, **kwargs):
+            if result_dict is None:
+                continue
+            yield _save_metadata_to_result(result_dict, kwargs)
 
     return wrapper
diff --git a/tts_webui/decorators/decorator_save_wav.py b/tts_webui/decorators/decorator_save_wav.py
@@ -28,3 +28,24 @@ def wrapper(*args, **kwargs):
             yield result_dict
 
     return wrapper
+
+
+def decorator_save_wav_generator_accumulated(fn):
+    def wrapper(*args, **kwargs):
+        accumulated_result_dict = None
+        for result_dict in fn(*args, **kwargs):
+            if result_dict is None:
+                continue
+            if accumulated_result_dict is None:
+                accumulated_result_dict = result_dict
+            else:
+                accumulated_result_dict["audio_out"][1] = np.concatenate(
+                    [
+                        accumulated_result_dict["audio_out"][1],
+                        result_dict["audio_out"][1],
+                    ]
+                )
+            yield result_dict
+        _save_wav(accumulated_result_dict)
+
+    return wrapper
diff --git a/tts_webui/extensions_loader/decorator_extensions.py b/tts_webui/extensions_loader/decorator_extensions.py
@@ -176,11 +176,25 @@ def wrapped(*args, **kwargs):
     return decorator
 
 
+def _create_decorator_generator(wrappers_list):
+    def decorator(fn0):
+        for wrapper in wrappers_list:
+            fn0 = wrapper(fn0)
+
+        @functools.wraps(fn0)
+        def wrapped(*args, **kwargs):
+            yield from fn0(*args, **kwargs)
+
+        return wrapped
+
+    return decorator
+
+
 # Define the four decorators using the helper function
 decorator_extension_outer = _create_decorator(OUTER_WRAPPERS)
 decorator_extension_inner = _create_decorator(INNER_WRAPPERS)
-decorator_extension_outer_generator = _create_decorator(OUTER_WRAPPERS_GEN)
-decorator_extension_inner_generator = _create_decorator(INNER_WRAPPERS_GEN)
+decorator_extension_outer_generator = _create_decorator_generator(OUTER_WRAPPERS_GEN)
+decorator_extension_inner_generator = _create_decorator_generator(INNER_WRAPPERS_GEN)
 
 if __name__ == "__main__":
     pass
diff --git a/tts_webui/utils/list_dir_models.py b/tts_webui/utils/list_dir_models.py
@@ -48,7 +48,18 @@ def model_select_ui(
 
 
 def unload_model_button(prefix: str):
-    return gr.Button(value="Unload Model").click(
+    button = gr.Button(value="Unload Model", variant="stop")
+
+    button.click(
+        fn=lambda: gr.Button(
+            value="Unloading...", variant="primary", interactive=False
+        ),
+        outputs=[button],
+    ).then(
         fn=lambda: unload_model(model_namespace=prefix),
         api_name=f"{prefix}_unload_model",
+    ).then(
+        fn=lambda: gr.Button(value="Unload Model", variant="stop", interactive=True),
+        outputs=[button],
     )
+    return button
diff --git a/tts_webui/utils/manage_model_state.py b/tts_webui/utils/manage_model_state.py
@@ -25,6 +25,9 @@ def is_model_loaded(self, model_name):
     def get_model_name(self):
         return self._model_name
 
+    def set_model_name(self, model_name):
+        self._model_name = model_name
+
 
 model_states = {}
 
@@ -95,3 +98,15 @@ def is_model_loaded(model_namespace):
         model_namespace in model_states
         and model_states[model_namespace].get_model() is not None
     )
+
+
+def rename_model(model_namespace, new_name):
+    if model_namespace in model_states:
+        model_states[model_namespace].set_model_name(new_name)
+
+
+def get_current_model(model_namespace):
+    if model_namespace in model_states:
+        return model_states[model_namespace].get_model()
+    else:
+        return None
diff --git a/tts_webui/utils/randomize_seed.py b/tts_webui/utils/randomize_seed.py
@@ -10,9 +10,9 @@ def randomize_seed(seed, randomize_seed):
 
 
 def randomize_seed_ui():
-    with gr.Row():
-        seed_input = gr.Textbox(label="Seed", value="-1")
-        randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)
+    with gr.Accordion("Seed", open=True), gr.Group(), gr.Row(equal_height=True):
+        seed_input = gr.Textbox(scale=9, container=False, show_label=False, show_copy_button=True, value="-1")
+        randomize_seed_checkbox = gr.Checkbox(scale=1, label="Randomize seed", value=True)
 
     return (
         seed_input,
diff --git a/tts_webui/utils/split_text_functions.py b/tts_webui/utils/split_text_functions.py
@@ -7,3 +7,15 @@ def split_by_lines(prompt: str):
 
 def split_by_length_simple(prompt: str):
     return [prompt[i : i + 200] for i in range(0, len(prompt), 200)]
+
+
+try:
+    from tortoise.utils.text import split_and_recombine_text
+except ImportError:
+
+    def split_and_recombine_text(
+        text: str, desired_length: int = 200, max_length: int = 300
+    ):
+        return split_by_length_simple(text)
+
+    # split_and_recombine_text = split_by_length_simple
diff --git a/voices/chatterbox/Alice.wav b/voices/chatterbox/Alice.wav
diff --git a/voices/chatterbox/Emmett.wav b/voices/chatterbox/Emmett.wav
diff --git a/voices/chatterbox/Sloane.wav b/voices/chatterbox/Sloane.wav
