Merge pull request #255 from bmalaski/alltalkbeta_enhancements

Alltalkbeta enhancements - Finetuning

.gitignore

@@ -1,8 +1,173 @@
-voices/*.*
-models/*.*
-outputs/*.*
-finetune/*.*
+voices/*
+models/*
+outputs/*
+finetune/*
+alltalk_environment/*
 system/config/at_github_sha.json
 system/tts_engines/tts_engines.json
 confignew.json
 diagnostics.log
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+
+.idea/*

finetune.py

@@ -1,3 +1,4 @@
+import math
 import os
 import gc
 import re
@@ -6,6 +7,8 @@ import time
 import glob
 import json
 import site
+import datetime
+
 import torch
 import signal
 import random
@@ -25,16 +28,22 @@ from pathlib import Path
 from faster_whisper import WhisperModel  
 from TTS.tts.configs.xtts_config import XttsConfig
 from TTS.tts.models.xtts import Xtts
+
+from metrics_logger import MetricsLogger
 # Use a local Tokenizer to resolve Japanese support
 # from TTS.tts.layers.xtts.tokenizer import multilingual_cleaners
 from system.ft_tokenizer.tokenizer import multilingual_cleaners
 import importlib.metadata as metadata
 from packaging import version
+from tokenizers import ByteLevelBPETokenizer
+from tokenizers.pre_tokenizers import Whitespace
 
 # STARTUP VARIABLES 
 this_dir = Path(__file__).parent.resolve()
 audio_folder = this_dir / "finetune" / "put-voice-samples-in-here"
-out_path = this_dir / "finetune" / "tmp-trn"
+default_path = this_dir / "finetune" / "tmp-trn"
+out_path = default_path
+
 progress = 0
 theme = gr.themes.Default()
 refresh_symbol = '🔄'
@@ -310,13 +319,17 @@ def create_temporary_file(folder, suffix=".wav"):
 
 
 
-def format_audio_list(target_language, whisper_model, out_path, max_sample_length, eval_split_number, speaker_name_input, gradio_progress=gr.Progress()):
+def format_audio_list(target_language, whisper_model, max_sample_length, eval_split_number, speaker_name_input, create_bpe_tokenizer, gradio_progress=gr.Progress()):
     global validate_train_metadata_path, validate_eval_metadata_path, validate_audio_folder, validate_whisper_model, validate_target_language
     buffer = 0.3
     max_duration = float(max_sample_length)  # Ensure max_duration is a float
     eval_percentage = eval_split_number / 100.0
     speaker_name = speaker_name_input
     audio_total_size = 0
+    global out_path
+    if speaker_name and speaker_name != 'personsname':
+        out_path = this_dir / "finetune" / speaker_name
+
     os.makedirs(out_path, exist_ok=True)
     temp_folder = os.path.join(out_path, "temp")
     os.makedirs(temp_folder, exist_ok=True)
@@ -367,8 +380,13 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt
         print(f"[FINETUNE] No audio files found in {original_samples_folder}. Skipping processing.")
         return None, None, 0
 
+    whisper_words = []
+    audio_steps = (0, len(original_audio_files))
+    gradio_progress_duration = 0
+    gradio_progress(audio_steps, desc="Processing Audio Files", unit="files")
+
     for audio_path in tqdm(original_audio_files):  # No start argument for tqdm
-        gradio_progress((2, 4), desc="Processing Audio Files")
+        start = datetime.datetime.now()
         audio_file_name_without_ext, _ = os.path.splitext(os.path.basename(audio_path))
         temp_audio_path = os.path.join(temp_folder, f"{audio_file_name_without_ext}.wav")
 
@@ -415,6 +433,9 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt
             words_list.extend(words)
 
         for word_idx, word in enumerate(words_list):
+            if create_bpe_tokenizer:
+                whisper_words.append(word.word)
+
             if first_word:
                 sentence_start = word.start
                 if word_idx == 0:
@@ -481,6 +502,14 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt
 
         os.remove(temp_audio_path)
 
+        end = datetime.datetime.now()
+        gradio_progress_duration += (end - start).total_seconds()
+        audio_steps = (audio_steps[0] + 1, audio_steps[1])
+        additional_data_points_needed = audio_steps[1] - audio_steps[0]
+        avg_duration = gradio_progress_duration / audio_steps[0]
+        gradio_estimated_duration = (avg_duration * additional_data_points_needed)
+        gradio_progress(audio_steps, desc=f"Processing. Estimated Completion: {c_logger.format_duration(gradio_estimated_duration)}", unit="files")
+
     # Check if the WAV files folder contains files after processing
     audio_files = [os.path.join(audio_folder, file) for file in os.listdir(audio_folder) if file.endswith('.wav')]
 
@@ -532,6 +561,17 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt
     final_training_set.sort_values('audio_file').to_csv(train_metadata_path, sep='|', index=False)
     final_eval_set.sort_values('audio_file').to_csv(eval_metadata_path, sep='|', index=False)
 
+    if create_bpe_tokenizer:
+        print(f"Training BPE Tokenizer")
+        tokenizer = ByteLevelBPETokenizer(str(this_dir / base_path / "xttsv2_2.0.3" / "vocab.json"))
+        tokenizer.pre_tokenizer = Whitespace()
+        tokenizer.add_tokens(["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[STOP]", "[SAPCE]"])
+        tokenizer.train_from_iterator(whisper_words, vocab_size=30000, show_progress=True, min_frequency=2, special_tokens=[
+            "[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[STOP]", "[SPACE]",
+        ])
+        # Save the tokenizer model
+        tokenizer.save(path=str(out_path / "bpe_tokenizer-vocab.json"), pretty=True)
+
     gradio_progress((4,4), "Finalizing")
 
     del asr_model, final_eval_set, final_training_set, new_data_df, existing_metadata
@@ -648,7 +688,8 @@ def load_and_display_mismatches():
 #### STEP 2 BITS #####
 ######################
 
-from trainer import Trainer, TrainerArgs
+from trainer import TrainerArgs, Trainer
+from trainer_alltalk.trainer import Trainer
 
 from TTS.config.shared_configs import BaseDatasetConfig
 from TTS.tts.datasets import load_tts_samples
@@ -663,10 +704,10 @@ def basemodel_or_finetunedmodel_choice(value):
     elif value == "Existing finetuned model":
         basemodel_or_finetunedmodel = False
 
-def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, output_path, max_audio_length=255995, progress=gr.Progress()):  
+def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, max_audio_length=255995, progress=gr.Progress()):
     pfc_check_fail()
 
-    if "No Models Available" in eval_csv:
+    if "No Models Available" in model_to_train:
         print(f"[FINETUNE] \033[91mError\033[0m: Cannot train. You have not selected a model. Please download a model.")
         print(f"[FINETUNE] \033[91mError\033[0m: into a sub folder within the correct models folder.")
         return   
@@ -678,7 +719,7 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv,
     model_path = this_dir / "models" / "xtts" / model_to_train
 
     # Set here the path that the checkpoints will be saved. Default: ./training/
-    OUT_PATH = os.path.join(output_path, "training")
+    OUT_PATH = os.path.join(out_path, "training")
     print("[FINETUNE] \033[94mStarting Step 2\033[0m - Fine-tuning the XTTS Encoder")
     print(f"[FINETUNE] \033[94mLanguage: \033[92m{language} \033[94mEpochs: \033[92m{num_epochs} \033[94mBatch size: \033[92m{batch_size}\033[0m \033[94mGrad accumulation steps: \033[92m{grad_acumm}\033[0m")
     print(f"[FINETUNE] \033[94mTraining   : \033[92m{train_csv}\033[0m")
@@ -731,8 +772,30 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv,
     MEL_NORM_FILE = model_path / "mel_stats.pth"
     SPEAKERS_FILE = model_path / "speakers_xtts.pth"
 
+    training_assets = None
+    if (out_path / "bpe_tokenizer-vocab.json").exists():
+        print("[FINETUNE] Using custom BPE tokenizer")
+        training_assets = {
+            'Tokenizer': str(out_path / "bpe_tokenizer-vocab.json")
+        }
+
+    continue_path = None
+    if continue_run:
+        folders = glob.glob(os.path.join(OUT_PATH, '*/'))
+        if folders:
+            last_run = max(folders, key=os.path.getmtime)
+            if last_run:
+                checkpoints = glob.glob(os.path.join(last_run, "best_model_*.pth"))
+                if checkpoints:
+                    latest_checkpoint = max(checkpoints, key=os.path.getmtime)
+                    if latest_checkpoint:
+                        XTTS_CHECKPOINT = None
+                        continue_path = last_run
+                        print(f"[FINETUNE] Continuing previous fine tuning {latest_checkpoint}")
+
+
     # Copy the supporting files
-    destination_dir = this_dir / "finetune" / "tmp-trn" / "chkptandnorm"
+    destination_dir = out_path / "chkptandnorm"
     destination_dir.mkdir(parents=True, exist_ok=True)
     shutil.copy2(DVAE_CHECKPOINT, destination_dir / DVAE_CHECKPOINT.name)
     shutil.copy2(MEL_NORM_FILE, destination_dir / MEL_NORM_FILE.name)
@@ -763,6 +826,99 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv,
     if language == "ja":
         number_of_workers = 0
 
+    lr_scheduler = None
+    lr_scheduler_params = {}
+
+    if learning_rate_scheduler and learning_rate_scheduler != "None":
+        lr_gamma_mapping = {
+            1e-6: 0.9,
+            5e-6: 0.8,
+            1e-5: 0.3,
+            5e-5: 0.3,
+            1e-4: 0.3,
+            5e-4: 0.1,
+            1e-3: 0.1
+        }
+        lr_scheduler = learning_rate_scheduler
+        if lr_scheduler == "StepLR":
+            lr_scheduler_params = {'step_size': 30, 'gamma': 0.1, 'last_epoch': -1}
+        elif lr_scheduler == "MultiStepLR":
+            exponent = 3 - int(math.log2(num_epochs) / 2)
+            base = 2
+            num_milestones = min(num_epochs, int(math.pow(base, exponent)))
+            milestone_interval = num_epochs // (num_milestones + 1)
+            milestones = [milestone_interval * (i + 1) for i in range(num_milestones)]
+            lr_scheduler_params = {'milestones': milestones, 'gamma': lr_gamma_mapping[learning_rate], 'last_epoch': -1}
+        elif lr_scheduler == "ExponentialLR":
+            lr_scheduler_params = {'gamma': 0.5, 'last_epoch': -1}
+        elif lr_scheduler == "CosineAnnealingLR":
+            lr_scheduler_params = {'T_max': num_epochs, 'eta_min': 1e-6, 'last_epoch': -1}
+        elif lr_scheduler == "ReduceLROnPlateau":
+            lr_scheduler_params = {'mode': 'min', 'factor': 0.8, 'patience': 1, 'threshold': 0.0001,
+                                   'threshold_mode': 'rel', 'cooldown': 0, 'min_lr': 1e-8, 'eps': 1e-08,}
+        elif lr_scheduler == "CyclicLR":
+            lr_scheduler_params = {'base_lr': learning_rate, 'max_lr': 0.1, 'step_size_up': 2000, 'step_size_down': None,
+                                   'mode': 'triangular', 'gamma': 1.0, 'scale_fn': None, 'scale_mode': 'cycle',
+                                   'cycle_momentum': True, 'base_momentum': 0.8, 'max_momentum': 0.9, 'last_epoch': -1}
+        elif lr_scheduler == "OneCycleLR":
+            #TODO: Determine total steps / epoch for this
+            lr_scheduler_params = {'max_lr': learning_rate, 'total_steps': None, 'epochs_up': None, 'steps_per_epoch': None,
+                                   'anneal_strategy': 'cos', 'cycle_momentum': True, 'base_momentum': 0.85,
+                                   'max_momentum': 0.95, 'div_factor': 25.0, 'final_div_factor': 10000.0,
+                                   'last_epoch': -1}
+        elif lr_scheduler == "CosineAnnealingWarmRestarts":
+            if num_epochs < 4:
+                error_message = "For Cosine Annealing Warm Restarts, epochs must be at least 4. Please set a minimum of 4 epochs."
+                progress(1.0, desc=f"Error: {error_message}")
+                raise ValueError(error_message)
+            #Set 4 learning rate restarts
+            lr_scheduler_params = {'T_0': int(num_epochs / 4), 'T_mult': 1, 'eta_min': 1e-6, 'last_epoch': -1}
+
+    optimizer_params = None
+    optimizer_wd_only_on_weight = OPTIMIZER_WD_ONLY_ON_WEIGHTS
+
+    if optimizer == "AdamW":
+        optimizer_params = {
+            "betas": [0.9, 0.96],
+            "eps": 1e-8,
+            "weight_decay": 1e-2
+        }
+    elif optimizer == "RMSprop":
+        optimizer_params = {
+            "alpha": 0.99,
+            "eps": 1e-8,
+            "weight_decay": 1e-4
+        }
+    elif optimizer == "SGD":
+        optimizer_params = {
+            "momentum": 0.9,
+            "weight_decay": 1e-4
+        }
+    elif optimizer == "Adam":
+        optimizer_params = {
+            "betas": [0.9, 0.999],
+            "eps": 1e-8,
+            "weight_decay": 1e-4
+        }
+    elif optimizer == "Adagrad":
+        optimizer_params = {
+            "lr_decay": 0,
+            "weight_decay": 1e-4,
+            "eps": 1e-10
+        }
+    elif optimizer == "RAdam":
+        optimizer_params = {
+            "betas": [0.9, 0.999],
+            "eps": 1e-8,
+            "weight_decay": 1e-2
+        }
+    elif optimizer == "stepwisegraduallr" or optimizer == "noamlr":
+        optimizer_params = { }
+        # These are Hardcoded learning rate schedulers
+
+
+    print(f"[FINETUNE] Learning Scheduler {lr_scheduler}, params {lr_scheduler_params}")
+
     # training parameters config
     config = GPTTrainerConfig(
         epochs=num_epochs,
@@ -788,13 +944,13 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv,
         # target_loss="loss",
         print_eval=False,
         # Optimizer values like tortoise, pytorch implementation with modifications to not apply WD to non-weight parameters.
-        optimizer="AdamW",
+        optimizer=optimizer,
         optimizer_wd_only_on_weights=OPTIMIZER_WD_ONLY_ON_WEIGHTS,
-        optimizer_params={"betas": [0.9, 0.96], "eps": 1e-8, "weight_decay": 1e-2},
+        optimizer_params=optimizer_params,
         lr=learning_rate,  # learning rate
-        lr_scheduler="MultiStepLR",
+        lr_scheduler=lr_scheduler,
         # it was adjusted accordly for the new step scheme
-        lr_scheduler_params={"milestones": [50000 * 18, 150000 * 18, 300000 * 18], "gamma": 0.5, "last_epoch": -1},
+        lr_scheduler_params=lr_scheduler_params,
         test_sentences=[],
     )
     progress(0, desc="Model is currently training. See console for more information")
@@ -808,11 +964,15 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv,
         eval_split_size=config.eval_split_size,
     )
 
+    global c_logger
+    c_logger = MetricsLogger()
+
     # init the trainer
     trainer = Trainer(
         TrainerArgs(
             restore_path=None,  # xtts checkpoint is restored via xtts_checkpoint key so no need of restore it using Trainer restore_path parameter
             skip_train_epoch=False,
+            continue_path=continue_path,
             start_with_eval=START_WITH_EVAL,
             grad_accum_steps=GRAD_ACUMM_STEPS,
         ),
@@ -821,7 +981,17 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv,
         model=model,
         train_samples=train_samples,
         eval_samples=eval_samples,
+        training_assets=training_assets,
+        c_logger=c_logger,
+        warmup=warm_up,
+
     )
+
+    if(disable_shared_memory):
+        # Limit training to GPU memory instead of shared memory
+        print("[FINETUNE] Limiting GPU memory to 95% to prevent spillover")
+        torch.cuda.set_per_process_memory_fraction(0.95)
+
     trainer.fit()
 
     # get the longest text audio file to use as speaker reference
@@ -884,7 +1054,13 @@ def run_tts(lang, tts_text, speaker_audio_file):
     if XTTS_MODEL is None or not speaker_audio_file:
         return "You need to run the previous step to load the model !!", None, None
     speaker_audio_file = str(speaker_audio_file)
-    gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(audio_path=speaker_audio_file, gpt_cond_len=XTTS_MODEL.config.gpt_cond_len, max_ref_length=XTTS_MODEL.config.max_ref_len, sound_norm_refs=XTTS_MODEL.config.sound_norm_refs)
+    if os.path.isdir(speaker_audio_file):
+        wavs_files = glob.glob(os.path.join(speaker_audio_file, "*.wav"))
+        speaker_audio_file = wavs_files[0]
+    else:
+        wavs_files = [speaker_audio_file]
+
+    gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(audio_path=wavs_files, gpt_cond_len=XTTS_MODEL.config.gpt_cond_len, max_ref_length=XTTS_MODEL.config.max_ref_len, sound_norm_refs=XTTS_MODEL.config.sound_norm_refs)
     out = XTTS_MODEL.inference(
         text=tts_text,
         language=lang,
@@ -904,33 +1080,45 @@ def run_tts(lang, tts_text, speaker_audio_file):
 
     return "Speech generated !", out_path, speaker_audio_file
 
-def get_available_voices(minimum_size_kb=1200):
+def get_available_voices(minimum_size_kb=1200, speaker_name=None):
+    directory = out_path
+    if speaker_name and speaker_name != 'personsname':
+        directory = this_dir / "finetune" / speaker_name
+
     voice_files = [
-        voice for voice in Path(f"{this_dir}/finetune/tmp-trn/wavs").glob("*.wav")
+        voice for voice in Path(f"{directory}/wavs").glob("*.wav")
         if voice.stat().st_size > minimum_size_kb * 1200  # Convert KB to bytes
     ]
     return sorted([str(file) for file in voice_files])  # Return full path as string
 
-def find_best_models(directory):
+def find_best_models(directory, speaker_name=None):
+    if speaker_name and speaker_name != 'personsname':
+        directory = this_dir / "finetune" / speaker_name
+
     """Find files named 'best_model.pth' in the given directory."""
     return [str(file) for file in Path(directory).rglob("best_model.pth")]
 
-def find_models(directory, extension):
+def find_models(directory, extension,  speaker_name=None):
+    if speaker_name and speaker_name != 'personsname':
+        directory = this_dir / "finetune" / speaker_name
+
     """Find files with a specific extension in the given directory."""
     return [str(file) for file in Path(directory).rglob(f"*.{extension}")]
 
-def find_jsons(directory, filename):
+def find_jsons(directory, filename, speaker_name=None):
+    if speaker_name and speaker_name != 'personsname':
+        directory = this_dir / "finetune" / speaker_name
+
     """Find files with a specific filename in the given directory."""
     return [str(file) for file in Path(directory).rglob(filename)]
 
-# Your main directory
-main_directory = Path(this_dir) / "finetune" / "tmp-trn"
+
 # XTTS checkpoint files (best_model.pth)
-xtts_checkpoint_files = find_best_models(main_directory)
+xtts_checkpoint_files = find_best_models(out_path)
 # XTTS config files (config.json)
-xtts_config_files = find_jsons(main_directory, "config.json")
+xtts_config_files = find_jsons(out_path, "config.json")
 # XTTS vocab files (vocab.json)
-xtts_vocab_files = find_jsons(main_directory, "vocab.json")
+xtts_vocab_files = find_jsons(out_path, "vocab.json")
 
 ##########################
 #### STEP 4 AND OTHER ####
@@ -976,11 +1164,11 @@ def compact_custom_model(xtts_checkpoint_copy, folder_path, overwrite_existing):
     files_to_copy = ["vocab.json", "config.json"]
     copy_files(folder_path_new, target_dir, files_to_copy)
     # Copy second set of files
-    chkptandnorm_path = this_dir / "finetune" / "tmp-trn" / "chkptandnorm"
+    chkptandnorm_path = out_path / "chkptandnorm"
     files_to_copy2 = ["speakers_xtts.pth", "mel_stats.pth", "dvae.pth"]
     copy_files(chkptandnorm_path, target_dir, files_to_copy2)
     # Copy large wav files
-    source_wavs_dir = this_dir / "finetune" / "tmp-trn" / "wavs"
+    source_wavs_dir = out_path / "wavs"
     target_wavs_dir = target_dir / "wavs"
     target_wavs_dir.mkdir(parents=True, exist_ok=True)
     for file_path in source_wavs_dir.iterdir():
@@ -991,7 +1179,7 @@ def compact_custom_model(xtts_checkpoint_copy, folder_path, overwrite_existing):
 
 def delete_training_data():
     # Define the folder to be deleted
-    folder_to_delete = Path(this_dir / "finetune" / "tmp-trn")
+    folder_to_delete = Path(out_path)
 
     # Check if the folder exists before deleting
     if folder_to_delete.exists():
@@ -1077,6 +1265,9 @@ logging.basicConfig(
     ]
 )
 
+c_logger = MetricsLogger()
+def load_metrics():
+    return c_logger.plot_metrics(), f"Running Time: {c_logger.format_duration(c_logger.total_duration)} - Estimated Completion: {c_logger.format_duration(c_logger.estimated_duration)}"
 def read_logs():
     sys.stdout.flush()
     with open(sys.stdout.log_file, "r") as f:
@@ -1088,10 +1279,10 @@ def cleanup_before_exit(signum, frame):
     sys.exit(0)
 
 def create_refresh_button(refresh_components, refresh_methods, elem_class, interactive=True):
-    def refresh():
+    def refresh(speaker_name):
         updates = {}
         for component, method in zip(refresh_components, refresh_methods):
-            args = method() if callable(method) else method
+            args = method(speaker_name=speaker_name) if callable(method) else method
             if args and 'choices' in args:
                 # Select the most recent file (last in the sorted list)
                 args['value'] = args['choices'][-1] if args['choices'] else ""
@@ -1103,12 +1294,33 @@ def create_refresh_button(refresh_components, refresh_methods, elem_class, inter
     refresh_button = gr.Button("Refresh Dropdowns", elem_classes=elem_class, interactive=interactive)
     refresh_button.click(
         fn=refresh,
-        inputs=[],
+        inputs=[speaker_name_input_testing],
         outputs=refresh_components
     )
 
     return refresh_button
 
+def create_refresh_button_next(refresh_components, refresh_methods, elem_class, interactive=True):
+    def refresh_export(speaker_name):
+        updates = {}
+        for component, method in zip(refresh_components, refresh_methods):
+            args = method(speaker_name=speaker_name) if callable(method) else method
+            if args and 'choices' in args:
+                # Select the most recent file (last in the sorted list)
+                args['value'] = args['choices'][-1] if args['choices'] else ""
+            for k, v in args.items():
+                setattr(component, k, v)
+            updates[component] = gr.update(**(args or {}))
+        return updates
+
+    refresh_button = gr.Button("Refresh Dropdowns", elem_classes=elem_class, interactive=interactive)
+    refresh_button.click(
+        fn=refresh_export,
+        inputs=[speaker_name_input_export],
+        outputs=refresh_components
+    )
+
+    return refresh_button
 
 
 pfc_markdown = f"""
@@ -1315,11 +1527,11 @@ if __name__ == "__main__":
 #######################
         with gr.Tab("📁 Step 1 - Generating the dataset"):
             with gr.Tab("Generate Dataset"):  
-                out_path = gr.Textbox(
-                    label="Output path (where data and checkpoints will be saved):",
-                    value=out_path,
-                    visible=False,
-                )
+                # out_path = gr.Textbox(
+                #     label="Output path (where data and checkpoints will be saved):",
+                #     value=out_path,
+                #     visible=False,
+                # )
                 # Define directories
                 this_dir = Path(__file__).parent.resolve()
                 voice_samples_dir = this_dir / "finetune" / "put-voice-samples-in-here"
@@ -1381,6 +1593,12 @@ if __name__ == "__main__":
                 delete_dataset_button.click(delete_existing_training_data, outputs=output_text)
                 
                 with gr.Row():
+                    speaker_name_input = gr.Textbox(
+                        label="Training Project Name",
+                        value="personsname",
+                        visible=True,
+                        scale=2,
+                    )
                     whisper_model = gr.Dropdown(
                         label="Whisper Model",
                         value="large-v2",
@@ -1418,25 +1636,24 @@ if __name__ == "__main__":
                         scale=1,
                     )
                     max_sample_length = gr.Dropdown(
-                        label="Maximum Audio Length Created by Whisper (seconds)",
+                        label="Max Audio Length (in seconds)",
                         value="30",
                         choices=["10", "15", "20", "25", "30", "35", "40", "45", "50", "55", "60", "65", "70", "75", "80", "85", "90"],
                         scale=2
                     )
                     eval_split_number = gr.Number(
-                        label="Evaluation data Split (the % to use for Evaluation data)",
+                        label="Evaluation Data Split",
                         value=15,  # Default value
                         minimum=5,  # Minimum value
                         maximum=95,  # Maximum value
                         step=1,  # Increment step
                         scale=2,
                     )
-                    speaker_name_input = gr.Textbox(
-                    label="The name of the speaker/person you are training",
-                    value="personsname",
-                    visible=True,
-                    scale=2,
-                )
+                    create_bpe_tokenizer = gr.Checkbox(
+                        label='BPE Tokenizer',
+                        value=False,
+                        info='Custom Tokenizer for training'
+                    )
                 progress_data = gr.Label(
                     label="Progress:"
                 )
@@ -1449,7 +1666,7 @@ if __name__ == "__main__":
 
                 prompt_compute_btn = gr.Button(value="Step 1 - Create dataset")
             
-                def preprocess_dataset(language, whisper_model, out_path, max_sample_length, eval_split_number, speaker_name_input, progress=gr.Progress()):
+                def preprocess_dataset(language, whisper_model, max_sample_length, eval_split_number, speaker_name_input, create_bpe_tokenizer, progress=gr.Progress()):
                     clear_gpu_cache()
                     test_for_audio_files = [file for file in os.listdir(audio_folder) if any(file.lower().endswith(ext) for ext in ['.wav', '.mp3', '.flac'])]
                     if not test_for_audio_files:
@@ -1457,7 +1674,7 @@ if __name__ == "__main__":
                     else:
                         try:
 
-                            train_meta, eval_meta, audio_total_size = format_audio_list(target_language=language, whisper_model=whisper_model, out_path=out_path, max_sample_length=max_sample_length, eval_split_number=eval_split_number, speaker_name_input=speaker_name_input, gradio_progress=progress)
+                            train_meta, eval_meta, audio_total_size = format_audio_list(target_language=language, whisper_model=whisper_model, max_sample_length=max_sample_length, eval_split_number=eval_split_number, speaker_name_input=speaker_name_input, create_bpe_tokenizer=create_bpe_tokenizer, gradio_progress=progress)
                         except:
                             traceback.print_exc()
                             error = traceback.format_exc()
@@ -1472,7 +1689,7 @@ if __name__ == "__main__":
                         return message, "", ""
 
                     print("[FINETUNE] Dataset Generated. Either run Dataset Validation or move to Step 2")
-                    return "Dataset Generated. Either run Dataset Validation or move to Step 2", train_meta, eval_meta
+                    return "Dataset Generated. Either run Dataset Validation or move to Step 2", train_meta, eval_meta, speaker_name_input
             
             def play_selected_audio(selected_file):
                 if selected_file and selected_file != "Please Generate Your Dataset":
@@ -1499,49 +1716,86 @@ if __name__ == "__main__":
             with gr.Tab("Generate Dataset Instructions"):
                 gr.Markdown(
                     f"""
-                    ### 📁 <u>Generating the dataset</u><br>
-                    ### 🟥 <u>Important Note - Windows - "UserWarning: huggingface_hub cache-system uses symlinks."</u>
-                    ◽ This error is caused by Huggingfaces download software downloading the Whisper model. If you get this error, please restart your Windows command prompt with "Run as Administrator" and restart finetuning.<br>
-                    ◽ This should only occur the 1st time it downloads the Whisper model. [Huggingface Reference here](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache#limitations)
-                    ### 🟥 <u>Important Note - Language support.</u>
-                    ◽ Although I have done my best to help automate this step, the Whisper model is not great with all languages and may fail to build your training data correctly. The Large-v3 may be better at certain languages.<br>
-                    ◽ You can find information about the Whisper model [here](https://github.com/openai/whisper?tab=readme-ov-file#available-models-and-languages) and you can find data about manually building training data [here](https://docs.coqui.ai/en/latest/formatting_your_dataset.html), as well as details below about the file structure this step performs.<br>
-                    ### 🟦 <u>What you need to do</u>
-                    ◽ Please read Coqui's guide on what makes a good dataset [here](https://docs.coqui.ai/en/latest/what_makes_a_good_dataset.html#what-makes-a-good-dataset)<br>
-                    ◽ Place your audio files in <span style="color: #3366ff;">{str(audio_folder)}</span>          
-                    ◽ Your audio samples can be in the format <span style="color: #3366ff;">mp3, wav,</span> or <span style="color: #3366ff;">flac.</span><br>
-                    ◽ You will need a minimum of <span style="color: #3366ff;">2 minutes</span> of audio in either one or multiple audio files. 5 to 10 minutes of audio would probably be better, allowing for more varied sample data to be generated.<br>
-                    ◽ Very small sample files cause errors, so I would recommend that the samples are at least 30 seconds and longer.<br>
-                    ◽ FYI Anecdotal evidence suggests that the Whisper 2 model may yield superior results in audio splitting and dataset creation.<br>
-                    ### 🟨 <u>What this step is doing</u><br>
-                    ◽ With step one, we are going to be stripping your audio file(s) into smaller files, using Whisper to find spoken words/sentences, compile that into excel sheets of training data, ready for Step 2.<br>
-                    ◽ Whisper is making a best effort to find spoken audio and break it down into smaller audio files. The content of these audio files is then transcribed into CSV fles (which you can edit in Excel or similar).<br>
-                    ◽ These files (audio and CSV) are used at the next step to train the model "this is what the audio sounds like and these are the words being spoken".<br>
-                    ◽ If you are using Audio that has multiple speakers/people speaking, you will HAVE to edit the generated CSV files and remove the **OTHER** person out of the training list.<br>
-                    ◽ If you wish to manually look at the CSV files before running the next step, you are welcome to do so and edit them as necessary. The greater the accuracy of the text the better the training will be.<br>
-                    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; ◽ <span style="color: #3366ff;">/alltalk_tts/finetune/tmp-trn/lang.txt</span> Contains a two digit langauage code e.g. `en`, `fr` etc.<br>
-                    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; ◽ <span style="color: #3366ff;">/alltalk_tts/finetune/tmp-trn/metadata_train.csv</span> Contains the list of wavs to text that will be used. This is what the model is trained with.<br>
-                    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; ◽ <span style="color: #3366ff;">/alltalk_tts/finetune/tmp-trn/metadata_eval.csv</span> Lists the evaluation wav to text, used to generate TTS while training and evaluate the quality of generated TTS to the original sample.<br>
-                    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; ◽ <span style="color: #3366ff;">/alltalk_tts/finetune/tmp-trn/wavs/</span> These are the wavs split out from your original training samples and match the contents of the CSV files.<br>
-                    ◽ Whilst you can choose multiple Whisper models, its best only to use the 1 model as each one is about 3GB in size and will download to your local huggingface cache on first-time use. <br>
-                    ◽ If and when you have completed training, should you wish to delete this 3GB Whisper model from your system, you are welcome to do so.<br>
-                    ### 🟩 <u>How long will this take?</u><br>
-                    ◽ First time, it needs to download the Whisper model which is 3GB. After that a few minutes on an average 3-4 year old system.<br>
+                    #### 🟦 <u>What You Need to Do</u>
+                    1. Read [Coqui's guide on creating a good dataset](https://docs.coqui.ai/en/latest/what_makes_a_good_dataset.html).<br>
+                    2. Place your audio files (MP3, WAV, or FLAC) in `{str(audio_folder)}` Or Upload them through the interface.<br>
+                    3. Provide at least 2 minutes of audio, preferably 5-10 minutes for better results.<br>
+                    4. Maximum Audio Length will ensure the audio samples do not exceed the set value in seconds. The AI model can only process so much audio per audio sample, per epoch.<br>
+                    5. Enter a unique Project Name if you wish to create a dedicated training folder. You will need to specify this name as you move through the interfaces and click the Refresh buttons where needed.<br>
+                    6. Select the Whisper model (large-v2 recommended for English use cases).<br>
+                    7. Choose your Dataset Language.<br>
+                    8. Adjust the Evaluation Split percentage (default 15% is suitable for most cases).<br>
+                    9. Optionally enable the BPE Tokenizer for custom tokenization, which can improve the model's ability to handle unique words, accents, or languages by creating a vocabulary tailored to your specific dataset. This is especially useful for non-standard speech patterns or languages with complex morphology.<br>
+
+                    #### 🟨 <u>What This Step Does</u>
+                    1. Uses Whisper to transcribe and segment your audio files.<br>
+                    2. Creates smaller audio clips and corresponding transcriptions.<br>
+                    3. Generates two CSV files: `metadata_train.csv` and `metadata_eval.csv`.<br>
+                    4. Saves processed audio in `/finetune/[project_name]/wavs/`.<br>
+                    5. Stores CSV files and metadata in `/finetune/[project_name]/`.<br>
+
+                    Key files generated:<br>
+                    ◽ `/finetune/[project_name]/lang.txt`: Contains the two-letter language code.<br>
+                    ◽ `/finetune/[project_name]/metadata_train.csv`: List of audio files and transcriptions for training.<br>
+                    ◽ `/finetune/[project_name]/metadata_eval.csv`: Evaluation data for assessing training quality.<br>
+
+                    #### <u>Important Considerations</u>
+                    ◽ For multi-speaker audio, manually edit CSV files to remove other speakers.<br>
+                    ◽ You can manually edit CSV files to improve transcription accuracy & use the `Dataset Validation` to help identify possible bad transcriptions.<br>
+                    ◽ The Whisper 2 model may provide better results in audio splitting and dataset creation for English languages, however please see the Whisper site for more details on other languages.<br>
+                    ◽ Each Whisper model is about 3GB. It's recommended to stick with one model to save disk space.<br>
                     """
                 )
+                
+            with gr.Tab("Generate Dataset Notes"):
+                gr.Markdown(
+                    f"""
+                    #### 🟥 <u>Important Notes</u>
+                    ◽ **Windows Users:** If you encounter a `UserWarning: huggingface_hub cache-system uses symlinks` error, restart your Windows command prompt with `Run as Administrator` and relaunch the finetuning process. This typically only occurs during the first Whisper model download.<br>
+                    ◽ **Language Support:** Whisper's performance varies across languages. The Large-v3 model may offer better results for certain languages. For problematic languages, consider manual dataset creation.<br>
+
+                    #### 🟩 <u>Processing Time</u>
+                    ◽ First run: Requires downloading the 3GB Whisper model.<br>
+                    ◽ Subsequent runs: A few minutes on an average 3-4 year old system.<br>
+
+                    After processing, use the 'Dataset Validation' tab to check transcription accuracy and make necessary corrections.<br>
+
+                    For more details on Whisper models or manual dataset creation, refer to:<br>
+                    ◽ [Whisper Models Information](https://github.com/openai/whisper?tab=readme-ov-file#available-models-and-languages)<br>
+                    ◽ [Manual Dataset Creation Guide](https://docs.coqui.ai/en/latest/formatting_your_dataset.html)
+                    """
+                )                
+                
+                
             
 #######################
 #### GRADIO STEP 2 ####
 #######################
         with gr.Tab("💻 Step 2 - Training"):
-            with gr.Tab("Training the model"):                    
+            with gr.Tab("Training the model"):
                 with gr.Row():
-                    model_to_train_choices = list(available_models.keys())
-                    model_to_train = gr.Dropdown(
-                        choices=model_to_train_choices,
-                        label="Select the Model to train",
-                        value=model_to_train_choices[0] if model_to_train_choices else None
+                    speaker_name_input_training = gr.Textbox(
+                        label="Project Name",
+                        value="personsname",
+                        visible=True,
+                        scale=1,
                     )
+                    with gr.Group():
+                        continue_run = gr.Checkbox(
+                            value=False,
+                            label="Continue Previous Project",
+                            scale=1,
+                        )
+                        disable_shared_memory = gr.Checkbox(
+                            value=False,
+                            label="Disable Shared Memory Use",
+                            scale=1,
+                        )
+                        warm_up = gr.Checkbox(
+                            value=False,
+                            label="Perform Small a Warmup Learning",
+                            scale=1,
+                        )
                     train_csv = gr.Textbox(
                         label="Train CSV:",
                         scale=2,
@@ -1550,6 +1804,15 @@ if __name__ == "__main__":
                         label="Eval CSV:",
                         scale=2,
                     )
+                with gr.Row():
+                    model_to_train_choices = list(available_models.keys())
+                    model_to_train = gr.Dropdown(
+                        choices=model_to_train_choices,
+                        label="Select the Model to train",
+                        value=model_to_train_choices[0] if model_to_train_choices else None,
+                        scale=2,
+                    )                  
+  
                     learning_rates = gr.Dropdown(
                         value=5e-6,
                         label="Learning Rate",
@@ -1564,7 +1827,42 @@ if __name__ == "__main__":
                         ],
                         type="value",
                         allow_custom_value=True,
-                        scale=0,
+                        scale=1,
+                    )
+                    learning_rate_scheduler = gr.Dropdown(
+                        value="CosineAnnealingWarmRestarts",
+                        label="Learning Rate Scheduler(s)",
+                        choices=[
+                            ("None", "None"),
+                            ("Cosine Annealing", "CosineAnnealingLR"),
+                            ("Cosine Annealing Warm Restarts", "CosineAnnealingWarmRestarts"),
+                            ("Cyclic", "CyclicLR"),
+                            ("Exponential", "ExponentialLR"),
+                            ("Multi Step", "MultiStepLR"),
+                            ("Reduce on Plateau", "ReduceLROnPlateau"),
+                            ("Step", "StepLR")
+                            #("OneCycleLR", "OneCycleLR"),
+                        ],
+                        type="value",
+                        allow_custom_value=False,
+                        scale=2,
+                    )
+                    optimizer = gr.Dropdown(
+                        value="AdamW",
+                        label="Optimizer",
+                        choices=[
+                            ("AdamW", "AdamW"),
+                            ("SGD with Momentum", "SGD"),
+                            ("RMSprop", "RMSprop"),
+                            ("Adagrad", "Adagrad"),
+                            ("Adam", "Adam"),
+                            ("Rectified Adam", "RAdam"),
+                            ("Step Wise", "stepwisegraduallr"),
+                            ("Noam", "noamlr")
+                        ],
+                        type="value",
+                        allow_custom_value=False,
+                        scale=2,
                     )
                 with gr.Row():
                     num_epochs =  gr.Slider(
@@ -1584,7 +1882,7 @@ if __name__ == "__main__":
                     grad_acumm = gr.Slider(
                         label="Grad accumulation steps:",
                         minimum=1,
-                        maximum=32,
+                        maximum=128,
                         step=1,
                         value=args.grad_acumm,
                     )
@@ -1599,42 +1897,65 @@ if __name__ == "__main__":
                 progress_train = gr.Label(
                     label="Progress:"
                 )
+
+                with gr.Row():
+                    train_time = gr.Label("Estimated Total Training Time", show_label=False, scale=2)
+                    train_btn = gr.Button(value="Step 2 - Run the training", scale=1)
+
+                with gr.Row():
+                    model_data = gr.Image(c_logger.plot_metrics(), show_label=False)
+
                 logs_tts_train = gr.Textbox(
                     label="Logs:",
                     interactive=False,
                     lines=10,
                 )
+                demo.load(load_metrics, None, [model_data, train_time], every=1)
                 demo.load(read_logs, None, logs_tts_train, every=1)
-                train_btn = gr.Button(value="Step 2 - Run the training")
 
-                def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, output_path, max_audio_length, progress=gr.Progress()):
+                def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, max_audio_length, speaker_name_input_training, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, progress=gr.Progress()):
                     clear_gpu_cache()
+                    global out_path
+                    if speaker_name_input_training and speaker_name_input_training != 'personsname':
+                        out_path = this_dir / "finetune" / speaker_name_input_training
+
                     if not train_csv or not eval_csv:
-                        return "You need to run the data processing step or manually set `Train CSV` and `Eval CSV` fields !", "", "", "", ""
+                        if (out_path / "metadata_eval.csv").exists() and (out_path / "metadata_train.csv").exists():
+                            train_csv = out_path / "metadata_train.csv"
+                            eval_csv = out_path / "metadata_eval.csv"
+                            print("Using existing metadata and training csv.")
+                        else:
+                            return "You need to run the data processing step or manually set `Train CSV` and `Eval CSV` fields !", "", "", "", ""
                     try:
                         # convert seconds to waveform frames
                         max_audio_length = int(max_audio_length * 22050)
                         learning_rate = float(learning_rates)  # Convert the learning rate value to a float
                         progress(0, "Initializing training...")
-                        config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, output_path=str(output_path), max_audio_length=max_audio_length, progress=gr.Progress())
-                    except:
-                        traceback.print_exc()
-                        error = traceback.format_exc()
-                        return f"The training was interrupted due an error !! Please check the console to check the full error message! \n Error summary: {error}", "", "", "", ""
+                        config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, max_audio_length=max_audio_length, progress=gr.Progress())
+                    
+                        # copy original files to avoid parameters changes issues
+                        shutil.copy(config_path, exp_path)
+                        shutil.copy(vocab_file, exp_path)
 
-                    # copy original files to avoid parameters changes issues
-                    shutil.copy(config_path, exp_path)
-                    shutil.copy(vocab_file, exp_path)
+                        ft_xtts_checkpoint = os.path.join(exp_path, "best_model.pth")
+                        print("[FINETUNE] Model training done. Move to Step 3")
+                        clear_gpu_cache()
+                        return "Model training done. Move to Step 3", config_path, vocab_file, ft_xtts_checkpoint, speaker_wav, speaker_name_input_training
 
-                    ft_xtts_checkpoint = os.path.join(exp_path, "best_model.pth")
-                    print("[FINETUNE] Model training done. Move to Step 3")
-                    clear_gpu_cache()
-                    return "Model training done. Move to Step 3", config_path, vocab_file, ft_xtts_checkpoint, speaker_wav
+                    except ValueError as ve:
+                        # This will catch the specific ValueError we raised in train_gpt
+                        error_message = str(ve)
+                        print(f"[FINETUNE] Error: {error_message}")
+                        return f"Training error: {error_message}", "", "", "", "", ""
+                    except Exception as e:
+                        # This will catch any other unexpected errors
+                        error_message = f"An unexpected error occurred: {str(e)}"
+                        print(f"[FINETUNE] Error: {error_message}")
+                        return f"Training error: {error_message}", "", "", "", "", ""
                 
             with gr.Tab("Training the model Instructions"):
                 gr.Markdown(
                     f"""
-                    ### 💻 <u>Training</u><br>
                     ### 🟥 <u>Important Note - Language support.</u>
                     ◽ If this step is failing/erroring you may wish to check your training data was created correctly (Detailed in Step 1), confirming that wav files have been generated and your `metadata_train.csv` and `metadata_eval.csv` files have been populated.<br>             
                     ◽ Ignore 'fatal: not a git repository (or any of the parent directories): .git' as this is a legacy training script issue.<br>  
@@ -1645,6 +1966,23 @@ if __name__ == "__main__":
                     &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;- What you are training (A human voice, A cartoon voice, A new language entirely etc)<br>
                     &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;- How much audio you have (you may want less or more eval or epochs)<br>
                     ◽ The key indicator of sufficient model training is whether it sounds right to you. Coqui suggests beginning with the base settings for training. If the resulting audio doesn't meet your expectations, additional training sessions may be necessary.<br>
+                    #### 🟪 <u>Continue Previous Project</u>
+                    This option allows you to resume training from where it left off if your previous session was interrupted (e.g., due to a power outage or system crash). When enabled:<br>
+                    ◽ The system will attempt to load the most recent checkpoint from your previous training session.<br>
+                    ◽ It will continue training from that point, preserving your progress and saving time.<br>
+                    ◽ This is particularly useful for long training sessions or when you want to extend training on a previously trained model.<br>
+                    #### 🟪 <u>Disable Shared Memory</u>
+                    This option affects how the system manages memory during training:<br>
+                    ◽ When enabled, it limits GPU memory usage to 95% of its capacity.<br>
+                    ◽ This prevents the training process from spilling over into your system's shared memory.<br>
+                    ◽ You're more likely to get a clear "out of memory" error if your GPU runs out of memory, rather than experiencing slower performance due to memory spillover.<br>
+                    ◽ This option is useful if you prefer to keep training confined to GPU memory for better performance and clearer error reporting.<br>
+
+                    Important Note for Windows Users with Low VRAM GPUs:<br>
+                    ◽ If you're using Windows and have a GPU with lower VRAM (e.g., 8GB or less), it's recommended to leave this option disabled.<br>
+                    ◽ Windows can utilize system RAM as extended VRAM, which can be crucial for completing the training process on lower VRAM GPUs.<br>
+                    ◽ Disabling this option allows Windows to manage memory allocation more flexibly, potentially preventing out-of-memory errors on GPUs with limited VRAM.
+                    
                     ### 🟨 <u>What this step is doing</u><br>
                     ◽ Very simply put, it's taking all our wav files generated in Step 1, along with our recorded speech in out excel documents and its training the model on that voice e.g. listen to this audio file and this is what is being said in it, so learn to reproduce it.<br>
                     ### 🟩 <u>How long will this take?</u><br>
@@ -1689,16 +2027,45 @@ if __name__ == "__main__":
                                 gr.Markdown(
                     f"""
                     ### 🟩 <u>Learning Rate</u><br>
-                    ◽ 1e-6: Very small learning rate, slow but stable learning.<br>
-                    ◽ 5e-6: Small learning rate, slow but stable learning.<br>
-                    ◽ 1e-5: Moderate learning rate, balanced between stability and convergence speed.<br>
-                    ◽ 5e-5: Higher learning rate, faster convergence but potential instability.<br>
-                    ◽ 1e-4: High learning rate, faster convergence but increased risk of instability.<br>
-                    ◽ 5e-4: Very high learning rate, fast convergence but higher risk of instability.<br>
-                    ◽ 1e-3: Extremely high learning rate, very fast convergence but high risk of instability.<br><br>
+                    ◽ **1e-6:** Very small learning rate, slow but stable learning.<br>
+                    ◽ **5e-6:** Small learning rate, slow but stable learning.<br>
+                    ◽ **1e-5:** Moderate learning rate, balanced between stability and convergence speed.<br>
+                    ◽ **5e-5:** Higher learning rate, faster convergence but potential instability.<br>
+                    ◽ **1e-4:** High learning rate, faster convergence but increased risk of instability.<br>
+                    ◽ **5e-4:** Very high learning rate, fast convergence but higher risk of instability.<br>
+                    ◽ **1e-3:** Extremely high learning rate, very fast convergence but high risk of instability.<br><br>
                     The optimal learning rate depends on the model architecture, dataset, and other hyperparameters.
                     """
                 )
+            with gr.Tab("Info - Learning Rate Schedulers"):
+                                gr.Markdown(
+                    f"""
+                    ### 🟩 <u>Learning Rate Schedulers</u><br>
+                    ◽ **None:** Learning rate remains unchanged throughout the training. <br>
+                    ◽ **Step:** Decreases the learning rate by a factor at specified intervals, offering a straightforward method for controlling learning rate decay.<br>
+                    ◽ **MultiStep:** Decreases the learning rate by a factor at specified milestones, useful for gradually reducing the learning rate during training.<br>
+                    ◽ **Exponential:** Decreases the learning rate exponentially over epochs, providing a simple decay mechanism for training.<br>
+                    ◽ **CosineAnnealing:** Adjusts the learning rate following a cosine annealing schedule, facilitating smoother optimization and potentially better convergence.<br>
+                    ◽ **Reduce On Plateau:** Reduces the learning rate when a monitored metric has stopped improving, helping to fine-tune learning rates during training.<br>
+                    ◽ **Cyclic:** Cycles the learning rate between two boundary values, allowing for dynamic learning rate schedules during training.<br>
+                    ◽ **One Cycle (disabled):** Sets the learning rate using a cyclical learning rate policy and varies momentum inversely, potentially accelerating training and improving performance.<br>
+                    ◽ **Cosine Annealing Warm Restarts:** Adjusts the learning rate following a cosine annealing schedule with warm restarts, offering a balance between exploration and exploitation during optimization.<br><br>
+                    The choice of learning rate scheduler depends on the specific characteristics of the model, dataset, and training objectives.
+                    """
+                )
+            with gr.Tab("Info - Optimizers"):
+                                gr.Markdown(
+                    f"""
+                    ### 🟩 <u>Optimizers</u><br>
+                    ◽ **SGD with Momentum:** Good for convergence, requires careful tuning of the learning rate and momentum.<br>
+                    ◽ **RMSprop:** Adaptively adjusts learning rates, good for varying gradient magnitudes.<br>
+                    ◽ **Adam:** Widely used, performs well with default parameters.<br>
+                    ◽ **Adagrad:** Good for sparse data, adapts learning rates based on historical gradients.<br>
+                    ◽ **AdamW:** Similar to Adam but with decoupled weight decay for better regularization.<br><br>
+                    When picking an optimizer, you need to experiment with different options to find the one that achieves the best performance and convergence for your task.<br><br>
+                    AdamW has been set to default as its a sensible pick that will handle most needs.
+                    """
+                )
             with gr.Tab("Info - Epochs"):
                                 gr.Markdown(
                     f"""
@@ -1778,6 +2145,12 @@ if __name__ == "__main__":
                                 allow_custom_value=True,  # Allow custom values
                                 scale=1,
                             )
+                            speaker_name_input_testing = gr.Textbox(
+                                label="Project Name (Refresh Dropdowns on change)",
+                                value="personsname",
+                                visible=True,
+                                scale=1,
+                            )
 
                         with gr.Row():
                             tts_language = gr.Dropdown(
@@ -1804,12 +2177,13 @@ if __name__ == "__main__":
                             )
                             # Create refresh button
                             refresh_button = create_refresh_button(
-                                [xtts_checkpoint, xtts_config, xtts_vocab, speaker_reference_audio],
+                                [xtts_checkpoint, xtts_config, xtts_vocab, speaker_reference_audio,
+                                 speaker_name_input_testing],
                                 [
-                                    lambda: {"choices": find_best_models(main_directory), "value": ""},
-                                    lambda: {"choices": find_jsons(main_directory, "config.json"), "value": ""},
-                                    lambda: {"choices": find_jsons(main_directory, "vocab.json"), "value": ""},
-                                    lambda: {"choices": get_available_voices(), "value": ""},
+                                    lambda speaker_name: {"choices": find_best_models(out_path, speaker_name=speaker_name), "value": ""},
+                                    lambda speaker_name: {"choices": find_jsons(out_path, "config.json", speaker_name=speaker_name), "value": ""},
+                                    lambda speaker_name: {"choices": find_jsons(out_path, "vocab.json", speaker_name=speaker_name),"value": ""},
+                                    lambda speaker_name: {"choices": get_available_voices(speaker_name=speaker_name),"value": ""},
                                 ],
                                 elem_class="refresh-button-class"
                             )
@@ -1878,11 +2252,17 @@ if __name__ == "__main__":
                     allow_custom_value=True,
                     scale=2,
                 )
+                speaker_name_input_export = gr.Textbox(
+                    label="Project Name (Refresh Dropdowns on change)",
+                    value="personsname",
+                    visible=True,
+                    scale=1,
+                )
                 # Create refresh button
-                refresh_button = create_refresh_button(
-                    [xtts_checkpoint_copy,],
+                refresh_button = create_refresh_button_next(
+                    [xtts_checkpoint_copy,speaker_name_input_export],
                     [
-                        lambda: {"choices": find_best_models(main_directory), "value": ""},
+                        lambda speaker_name: {"choices": find_best_models(out_path, speaker_name), "value": ""},
                     ],
                     elem_class="refresh-button-class")
             with gr.Row():
@@ -1901,15 +2281,16 @@ if __name__ == "__main__":
                 inputs=[
                     lang,
                     whisper_model,
-                    out_path,
                     max_sample_length,
                     eval_split_number,
                     speaker_name_input,
+                    create_bpe_tokenizer
                 ],
                 outputs=[
                     progress_data,
                     train_csv,
                     eval_csv,
+                    speaker_name_input_training
                 ],
             )
 
@@ -1925,10 +2306,15 @@ if __name__ == "__main__":
                     num_epochs,
                     batch_size,
                     grad_acumm,
-                    out_path,
                     max_audio_length,
+                    speaker_name_input_training,
+                    continue_run,
+                    disable_shared_memory,
+                    learning_rate_scheduler,
+                    optimizer,
+                    warm_up
                 ],
-                outputs=[progress_train, xtts_config, xtts_vocab, xtts_checkpoint, speaker_reference_audio],
+                outputs=[progress_train, xtts_config, xtts_vocab, xtts_checkpoint, speaker_reference_audio, speaker_name_input_testing],
             )
             
             load_btn.click(

metrics_logger.py

@@ -0,0 +1,161 @@
+import datetime
+import io
+import logging
+
+from matplotlib import pyplot as plt
+from trainer import ConsoleLogger
+from trainer.logging.console_logger import tcolors
+from trainer.utils.distributed import rank_zero_only
+
+logger = logging.getLogger("trainer")
+
+class MetricsLogger(ConsoleLogger):
+    def __init__(self, ):
+        super().__init__()
+        self.cur_epoch_start = None
+        self.cur_epoch_end = None
+        self.cur_epoch = None
+        self.epoch_durations = []
+        self.max_epoch = None
+        self.batch_steps = None
+        self.global_steps = None
+        self.loss = []
+        self.avg_loss = []
+        self.loss_mel_ce = []
+        self.avg_loss_mel_ce = []
+        self.loss_text_ce = []
+        self.avg_loss_text_ce = []
+        self.learning_rates = []
+        self.total_duration = None
+        self.estimated_duration = None
+        self.start = None
+
+    @staticmethod
+    def get_now():
+        return datetime.datetime.now()
+    @staticmethod
+    def format_duration(seconds):
+        if seconds:
+            # Calculate hours, minutes, and seconds
+            hours = seconds // 3600
+            minutes = (seconds % 3600) // 60
+            seconds = seconds % 60
+            return f"{int(hours)}:{int(minutes)}:{int(seconds)}"
+        else:
+            return None
+
+    @rank_zero_only
+    def print_epoch_start(self, epoch, max_epoch, output_path=None):
+        self.cur_epoch = epoch
+        self.cur_epoch_start = self.get_now()
+        self.max_epoch = max_epoch
+        super().print_epoch_start(epoch, max_epoch, output_path)
+
+    @rank_zero_only
+    def print_train_start(self):
+        super().print_train_start()
+
+    @rank_zero_only
+    def print_train_step(self, batch_steps, step, global_step, loss_dict, avg_loss_dict):
+        self.batch_steps = batch_steps
+        self.global_steps = global_step
+
+        for key, value in loss_dict.items():
+            if key == 'loss_mel_ce':
+                self.loss_mel_ce.append((global_step, value))
+            elif key == 'loss':
+                self.loss.append((global_step, value))
+            elif key == 'loss_text_ce':
+                self.loss_text_ce.append((global_step, value))
+            elif key == 'current_lr':
+                self.learning_rates.append((global_step, value))
+        super().print_train_step(batch_steps, step, global_step, loss_dict, avg_loss_dict)
+
+    # pylint: disable=unused-argument
+    @rank_zero_only
+    def print_train_epoch_end(self, global_step, epoch, epoch_time, print_dict):
+        super().print_epoch_end(global_step, epoch, epoch_time, print_dict)
+        if self.estimated_duration:
+            self.log_with_flush(
+                "\n{}{} >> ETA: {} {}".format(tcolors.UNDERLINE, tcolors.BOLD, self.format_duration(self.estimated_duration), tcolors.ENDC),
+            )
+
+    @rank_zero_only
+    def print_eval_start(self):
+        super().print_eval_start()
+
+    @rank_zero_only
+    def print_eval_step(self, step, loss_dict, avg_loss_dict):
+        super().print_eval_step(step, loss_dict, avg_loss_dict)
+
+    @rank_zero_only
+    def print_epoch_end(self, epoch, avg_loss_dict):
+        self.cur_epoch_end = self.get_now()
+        self.epoch_durations.append((self.cur_epoch_end - self.cur_epoch_start).total_seconds())
+        for key, value in avg_loss_dict.items():
+            if key == 'avg_loss_mel_ce':
+                self.avg_loss_mel_ce.append((epoch, value))
+            elif key == 'avg_loss':
+                self.avg_loss.append((epoch, value))
+            elif key == 'avg_loss_text_ce':
+                self.avg_loss_text_ce.append((epoch, value))
+
+        if len(self.epoch_durations) > 1 and self.max_epoch:
+            self.total_duration = sum(self.epoch_durations)
+            additional_data_points_needed = self.max_epoch - len(self.epoch_durations)
+            if additional_data_points_needed > 0:
+                weighted_avg_duration = sum((i + 1) * duration for i, duration in enumerate(self.epoch_durations)) / sum(range(1, len(self.epoch_durations) + 1))
+                self.estimated_duration = (weighted_avg_duration * additional_data_points_needed)
+
+        super().print_epoch_end(epoch, avg_loss_dict)
+
+    def plot_metrics(self, show_loss_mel_ce=False, show_loss_text_ce=False, show_loss=True):
+        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(24, 18))
+
+        ax2.set_title('Steps')
+        ax2.set_xlabel('Steps')
+        ax2.set_ylabel('Value')
+
+        ax1.set_title('Epochs')
+        ax1.set_xlabel('Epoch')
+        ax1.set_ylabel('Value')
+
+        ax3.set_title('Learning Rates')
+        ax3.set_xlabel('Steps')
+        ax3.set_ylabel('Value')
+
+        if self.loss and self.loss_text_ce and self.loss_mel_ce:
+            if show_loss:
+                ax2.plot(*zip(*self.loss), color='green', linestyle='-', label='Loss')
+                ax2.fill_between(*zip(*self.loss), color='lightgreen', alpha=0.1)
+            if show_loss_text_ce:
+                ax2.plot(*zip(*self.loss_text_ce), color='red', linestyle='-', label='Loss Text')
+                ax2.fill_between(*zip(*self.loss_text_ce), color='lightcoral', alpha=0.1)
+            if show_loss_mel_ce:
+                ax2.plot(*zip(*self.loss_mel_ce), color='blue', linestyle='-', label='Loss MEL')
+                ax2.fill_between(*zip(*self.loss_mel_ce), color='lightblue', alpha=0.1)
+            ax2.legend()
+
+        if self.learning_rates:
+            ax3.plot(*zip(*self.learning_rates), linestyle='-', color='blue', label='Learning Rate')
+            ax3.fill_between(*zip(*self.learning_rates), color='lightblue', alpha=0.1)
+            ax3.legend()
+
+        if self.avg_loss and self.avg_loss_text_ce and self.avg_loss_mel_ce:
+            if show_loss:
+                ax1.plot(*zip(*self.avg_loss), color='green', linestyle='-', label='Avg Loss')
+                ax1.fill_between(*zip(*self.avg_loss), color='lightgreen', alpha=0.1)
+            if show_loss_text_ce:
+                ax1.plot(*zip(*self.avg_loss_text_ce), color='red', linestyle='-', label='Avg Loss Text')
+                ax1.fill_between(*zip(*self.avg_loss_text_ce), color='lightcoral', alpha=0.1)
+            if show_loss_mel_ce:
+                ax1.plot(*zip(*self.avg_loss_mel_ce), color='blue', linestyle='-', label='Avg Loss MEL')
+                ax1.fill_between(*zip(*self.avg_loss_mel_ce), color='lightblue', alpha=0.1)
+            ax1.set_xlim(left=0)
+            ax1.legend()
+
+        plt.style.use('dark_background')
+        plt.tight_layout()
+        plt.savefig("training_metrics.png", format='png', bbox_inches='tight', pad_inches=0.1)
+        plt.close()
+        return "training_metrics.png"

system/tts_engines/xtts/model_engine.py

@@ -1,6 +1,7 @@
 ###############################################
 # DONT CHANGE # These are base imports needed #
 ###############################################
+import glob
 import os
 import sys
 import json
@@ -471,10 +472,15 @@ class tts_class:
         print(f"[{self.branding}Debug] Deciding if streaming or not") if self.debug_tts else None
         print(f"[{self.branding}Debug] self.current_model_loaded", self.current_model_loaded) if self.debug_tts else None
         self.current_model_loaded
+        if os.path.isdir(voice):
+            wavs_files = glob.glob(os.path.join(voice, "*.wav"))
+        else:
+            wavs_files = [f"{self.main_dir}/voices/{voice}"]
+
         if self.current_model_loaded.startswith ("xtts"):
             print(f"[{self.branding}Debug] Text arriving at engine {text}") if self.debug_tts else None
             gpt_cond_latent, speaker_embedding = self.model.get_conditioning_latents(
-                audio_path=[f"{self.main_dir}/voices/{voice}"],
+                audio_path=wavs_files,
                 gpt_cond_len=self.model.config.gpt_cond_len,
                 max_ref_length=self.model.config.max_ref_len,
                 sound_norm_refs=self.model.config.sound_norm_refs,
@@ -551,7 +557,7 @@ class tts_class:
             self.model.tts_to_file(
                 text=text,
                 file_path=output_file,
-                speaker_wav=[f"{self.main_dir}/voices/{voice}"],
+                speaker_wav=wavs_files,
                 language=language,
                 temperature=temperature,
                 length_penalty=self.model.config.length_penalty,

trainer_alltalk/WarmUpScheduler.py

@@ -0,0 +1,35 @@
+from torch.optim.lr_scheduler import LRScheduler
+
+class WarmUpScheduler(LRScheduler):
+    def __init__(self, optimizer, total_epochs, warmup_lr, target_lr, after_scheduler=None):
+        self.total_epochs = total_epochs
+        self.warmup_epochs = int(total_epochs * 0.05)
+        if self.warmup_epochs <= 0:
+            self.warmup_epochs = 1
+        self.warmup_lr = warmup_lr
+        self.target_lr = target_lr
+        self.after_scheduler = after_scheduler
+        print(f"[WARMUP] {self.warmup_epochs} epochs")
+        super(WarmUpScheduler, self).__init__(optimizer)
+
+    def get_lr(self):
+        if self.last_epoch < self.warmup_epochs:
+            progress = self.last_epoch / self.warmup_epochs
+            return [self.warmup_lr + progress * (self.target_lr - self.warmup_lr) for _ in self.base_lrs]
+        if self.after_scheduler:
+            return self.after_scheduler.get_last_lr()
+        return [base_lr for base_lr in self.base_lrs]
+
+    def step(self, epoch=None):
+        if epoch is not None and epoch < self.warmup_epochs:
+            self.last_epoch = epoch
+        elif self.after_scheduler and epoch is not None:
+            if epoch < self.warmup_epochs:
+                self.last_epoch = epoch
+            else:
+                self.after_scheduler.step(epoch - self.warmup_epochs)
+                self.last_epoch = epoch
+        else:
+            super(WarmUpScheduler, self).step(epoch)
+        if self.after_scheduler and self.last_epoch >= self.warmup_epochs:
+            self.after_scheduler.step(epoch - self.warmup_epochs if epoch is not None else None)

tts_server.py

@@ -1099,7 +1099,7 @@ def process_rvc_narrator(part_type, voice_gen, pitch, default_model_file, output
 class JSONInput(BaseModel):
     text_input: str = Field(..., max_length=int(api_defaults["api_max_characters"]), description=f"text_input needs to be {api_defaults['api_max_characters']} characters or less.")
     text_filtering: str = Field(..., pattern="^(none|standard|html)$", description="text_filtering needs to be 'none', 'standard' or 'html'.")
-    character_voice_gen: str = Field(..., pattern=r'^[\(\)a-zA-Z0-9\_\-./\s]+$', description="character_voice_gen needs to be the name of a valid voice for the loaded TTS engine.")
+    #character_voice_gen: str = Field(..., pattern=r'^[\(\)a-zA-Z0-9\_\-./\s]+$', description="character_voice_gen needs to be the name of a valid voice for the loaded TTS engine.")
     rvccharacter_voice_gen: str = Field(..., description="rvccharacter_voice_gen needs to be the name of a valid pth file in the 'folder\\file.pth' format or the word 'Disabled'.")
     rvccharacter_pitch: float = Field(..., description="RVC Character pitch needs to be a number between -24 and 24")
     narrator_enabled: str = Field(..., pattern="^(true|false|silent)$", description="narrator_enabled needs to be 'true', 'false' or 'silent'.")