From 2615cafd75bc7ad397f9fea2d0055b125af0ffaf Mon Sep 17 00:00:00 2001 From: mrq Date: Sat, 18 Feb 2023 14:10:26 +0000 Subject: [PATCH] added dropdown to select autoregressive model for TTS, fixed a bug where the settings saveer constantly fires I hate gradio so much why are dropdown.change broken to contiuously fire and send an empty array --- models/finetunes/.gitkeep | 0 src/utils.py | 57 ++++++++++++++++++++++++++++++++++----- src/webui.py | 32 +++++++++++++++++----- 3 files changed, 75 insertions(+), 14 deletions(-) create mode 100755 models/finetunes/.gitkeep diff --git a/models/finetunes/.gitkeep b/models/finetunes/.gitkeep new file mode 100755 index 0000000..e69de29 diff --git a/src/utils.py b/src/utils.py index 4188e11..166f333 100755 --- a/src/utils.py +++ b/src/utils.py @@ -62,6 +62,7 @@ def setup_args(): 'defer-tts-load': False, 'device-override': None, 'whisper-model': "base", + 'autoregressive-model': None, 'concurrency-count': 2, 'output-sample-rate': 44100, 'output-volume': 1, @@ -87,6 +88,7 @@ def setup_args(): parser.add_argument("--defer-tts-load", default=default_arguments['defer-tts-load'], action='store_true', help="Defers loading TTS model") parser.add_argument("--device-override", default=default_arguments['device-override'], help="A device string to override pass through Torch") parser.add_argument("--whisper-model", default=default_arguments['whisper-model'], help="Specifies which whisper model to use for transcription.") + parser.add_argument("--autoregressive-model", default=default_arguments['autoregressive-model'], help="Specifies which autoregressive model to use for sampling.") parser.add_argument("--sample-batch-size", default=default_arguments['sample-batch-size'], type=int, help="Sets how many batches to use during the autoregressive samples pass") parser.add_argument("--concurrency-count", type=int, default=default_arguments['concurrency-count'], help="How many Gradio events to process at once") parser.add_argument("--output-sample-rate", type=int, default=default_arguments['output-sample-rate'], help="Sample rate to resample the output to (from 24KHz)") @@ -151,10 +153,8 @@ def generate( global args global tts - try: - tts - except NameError: - raise Exception("TTS is still initializing...") + if not tts: + raise Exception("TTS is uninitialized or still initializing...") if voice != "microphone": voices = [voice] @@ -493,7 +493,7 @@ def setup_tortoise(restart=False): tts = None print("Initializating TorToiSe...") - tts = TextToSpeech(minor_optimizations=not args.low_vram) + tts = TextToSpeech(minor_optimizations=not args.low_vram, autoregressive_model_path=args.autoregressive_model) get_model_path('dvae.pth') print("TorToiSe initialized, ready for generation.") return tts @@ -720,7 +720,47 @@ def get_voice_list(dir=get_voice_dir()): os.makedirs(dir, exist_ok=True) return sorted([d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d)) and len(os.listdir(os.path.join(dir, d))) > 0 ]) + ["microphone", "random"] -def export_exec_settings( listen, share, check_for_updates, models_from_local_only, low_vram, embed_output_metadata, latents_lean_and_mean, voice_fixer, voice_fixer_use_cuda, force_cpu_for_conditioning_latents, defer_tts_load, device_override, sample_batch_size, concurrency_count, output_sample_rate, output_volume, whisper_model ): +def get_autoregressive_models(dir="./models/finetuned/"): + os.makedirs(dir, exist_ok=True) + return [get_model_path('autoregressive.pth')] + sorted([d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d)) and len(os.listdir(os.path.join(dir, d))) > 0 ]) + + +def update_autoregressive_model(path_name): + + global tts + if not tts: + raise Exception("TTS is uninitialized or still initializing...") + + print(f"Loading model: {path_name}") + if hasattr(tts, 'load_autoregressive_model') and tts.load_autoregressive_model(path_name): + args.autoregressive_model = path_name + save_args_settings() + # polyfill in case a user did NOT update the packages + else: + from tortoise.models.autoregressive import UnifiedVoice + + previous_path = tts.autoregressive_model_path + tts.autoregressive_model_path = path_name if path_name and os.path.exists(path_name) else get_model_path('autoregressive.pth') + + del tts.autoregressive + tts.autoregressive = UnifiedVoice(max_mel_tokens=604, max_text_tokens=402, max_conditioning_inputs=2, layers=30, + model_dim=1024, + heads=16, number_text_tokens=255, start_text_token=255, checkpointing=False, + train_solo_embeddings=False).cpu().eval() + tts.autoregressive.load_state_dict(torch.load(tts.autoregressive_model_path)) + tts.autoregressive.post_init_gpt2_config(kv_cache=tts.use_kv_cache) + if tts.preloaded_tensors: + tts.autoregressive = tts.autoregressive.to(tts.device) + + if previous_path != tts.autoregressive_model_path: + args.autoregressive_model = path_name + save_args_settings() + + print(f"Loaded model: {tts.autoregressive_model_path}") + + return path_name + +def update_args( listen, share, check_for_updates, models_from_local_only, low_vram, embed_output_metadata, latents_lean_and_mean, voice_fixer, voice_fixer_use_cuda, force_cpu_for_conditioning_latents, defer_tts_load, device_override, sample_batch_size, concurrency_count, output_sample_rate, output_volume ): global args args.listen = listen @@ -731,7 +771,6 @@ def export_exec_settings( listen, share, check_for_updates, models_from_local_on args.force_cpu_for_conditioning_latents = force_cpu_for_conditioning_latents args.defer_tts_load = defer_tts_load args.device_override = device_override - args.whisper_model = whisper_model args.sample_batch_size = sample_batch_size args.embed_output_metadata = embed_output_metadata args.latents_lean_and_mean = latents_lean_and_mean @@ -741,6 +780,9 @@ def export_exec_settings( listen, share, check_for_updates, models_from_local_on args.output_sample_rate = output_sample_rate args.output_volume = output_volume + save_args_settings() + +def save_args_settings(): settings = { 'listen': None if args.listen else args.listen, 'share': args.share, @@ -751,6 +793,7 @@ def export_exec_settings( listen, share, check_for_updates, models_from_local_on 'defer-tts-load': args.defer_tts_load, 'device-override': args.device_override, 'whisper-model': args.whisper_model, + 'autoregressive-model': args.autoregressive_model, 'sample-batch-size': args.sample_batch_size, 'embed-output-metadata': args.embed_output_metadata, 'latents-lean-and-mean': args.latents_lean_and_mean, diff --git a/src/webui.py b/src/webui.py index c56e38c..0ce4759 100755 --- a/src/webui.py +++ b/src/webui.py @@ -90,10 +90,8 @@ def compute_latents(voice, voice_latents_chunks, progress=gr.Progress(track_tqdm global tts global args - try: - tts - except NameError: - raise gr.Error("TTS is still initializing...") + if not tts: + raise Exception("TTS is uninitialized or still initializing...") voice_samples, conditioning_latents = load_voice(voice, load_latents=False) @@ -213,6 +211,14 @@ def update_voices(): def history_copy_settings( voice, file ): return import_generate_settings( f"./results/{voice}/{file}" ) +def update_model_settings( autoregressive_model, whisper_model ): + if args.autoregressive_model != autoregressive_model: + update_autoregressive_model(autoregressive_model) + + args.whisper_model = whisper_model + + save_args_settings() + def setup_gradio(): global args global ui @@ -370,13 +376,17 @@ def setup_gradio(): gr.Number(label="Concurrency Count", precision=0, value=args.concurrency_count), gr.Number(label="Ouptut Sample Rate", precision=0, value=args.output_sample_rate), gr.Slider(label="Ouptut Volume", minimum=0, maximum=2, value=args.output_volume), - gr.Dropdown(label="Whisper Model", value=args.whisper_model, choices=["tiny", "tiny.en", "base", "base.en", "small", "small.en", "medium", "medium.en", "large"]), ] + + autoregressive_model_dropdown = gr.Dropdown(get_autoregressive_models(), label="Autoregressive Model", value=args.autoregressive_model) + whisper_model_dropdown = gr.Dropdown(["tiny", "tiny.en", "base", "base.en", "small", "small.en", "medium", "medium.en", "large"], label="Whisper Model", value=args.whisper_model) + save_settings_button = gr.Button(value="Save Settings") + gr.Button(value="Check for Updates").click(check_for_updates) - gr.Button(value="Reload TTS").click(reload_tts) + gr.Button(value="(Re)Load TTS").click(reload_tts) for i in exec_inputs: - i.change( fn=export_exec_settings, inputs=exec_inputs ) + i.change( fn=update_args, inputs=exec_inputs ) # console_output = gr.TextArea(label="Console Output", interactive=False, max_lines=8) @@ -533,6 +543,14 @@ def setup_gradio(): outputs=save_yaml_output #console_output ) + save_settings_button.click(update_model_settings, + inputs=[ + autoregressive_model_dropdown, + whisper_model_dropdown, + ], + outputs=None + ) + if os.path.isfile('./config/generate.json'): ui.load(import_generate_settings, inputs=None, outputs=input_settings)