diff --git a/models/.template.yaml b/models/.template.yaml
index 825ac0c..77e5113 100755
--- a/models/.template.yaml
+++ b/models/.template.yaml
@@ -11,7 +11,7 @@ use_tb_logger: true
datasets:
train:
name: ${dataset_name}
- n_workers: 8
+ n_workers: ${workers}
batch_size: ${batch_size}
mode: paired_voice_audio
path: ${dataset_path}
diff --git a/src/utils.py b/src/utils.py
index 1f6a707..89dd1d1 100755
--- a/src/utils.py
+++ b/src/utils.py
@@ -37,6 +37,8 @@ from tortoise.utils.text import split_and_recombine_text
from tortoise.utils.device import get_device_name, set_device_name
MODELS['dvae.pth'] = "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/3704aea61678e7e468a06d8eea121dba368a798e/.models/dvae.pth"
+WHISPER_MODELS = ["tiny", "base", "small", "medium", "large"]
+WHISPER_SPECIALIZED_MODELS = ["tiny.en", "base.en", "small.en", "medium.en"]
args = None
tts = None
@@ -663,6 +665,7 @@ class TrainingState():
# rip out iteration info
if not self.training_started:
if line.find('Start training from epoch') >= 0:
+ self.it_time_start = time.time()
self.epoch_time_start = time.time()
self.training_started = True # could just leverage the above variable, but this is python, and there's no point in these aggressive microoptimizations
should_return = True
@@ -703,11 +706,12 @@ class TrainingState():
self.it_time_delta = self.it_time_end-self.it_time_start
self.it_time_start = time.time()
self.it_taken = self.it_taken + 1
- try:
- rate = f'{"{:.3f}".format(self.it_time_delta)}s/it' if self.it_time_delta >= 1 else f'{"{:.3f}".format(1/self.it_time_delta)}it/s'
- self.it_rate = rate
- except Exception as e:
- pass
+ if self.it_time_delta:
+ try:
+ rate = f'{"{:.3f}".format(self.it_time_delta)}s/it' if self.it_time_delta >= 1 else f'{"{:.3f}".format(1/self.it_time_delta)}it/s'
+ self.it_rate = rate
+ except Exception as e:
+ pass
metric_step = [f"{self.epoch}/{self.epochs}", f"{self.it}/{self.its}", f"{step}/{steps}"]
metric_step = ", ".join(metric_step)
@@ -733,9 +737,23 @@ class TrainingState():
metric_loss = []
if len(self.losses) > 0:
metric_loss.append(f'Loss: {"{:3f}".format(self.losses[-1]["value"])}')
+
+ if len(self.losses) >= 2:
+ delta_loss = self.losses[-2]["value"] - self.losses[-1]["value"]
+ delta_step = self.losses[-2]["step"] - self.losses[-1]["step"]
+
+ inst_deriv = delta_loss / delta_step
+ est_loss = delta_loss + (self.its - self.it) * inst_deriv
+ metric_loss.append(f'Est. Final Loss: {"{:3f}".format(est_loss)}')
+
+ print(delta_loss, delta_step, inst_deriv, est_loss)
+
+
metric_loss = ", ".join(metric_loss)
- message = f'[{metric_step}] [{metric_rate}] [{metric_loss}] [ETA: {eta_hhmmss}]'
+
+
+ message = f'[{metric_step}] [{metric_rate}] [ETA: {eta_hhmmss}] [{metric_loss}]'
if lapsed:
self.epoch = self.epoch + 1
@@ -764,6 +782,13 @@ class TrainingState():
self.buffer.append(f'[{"{:.3f}".format(percent*100)}%] {message}')
if line.find('INFO: [epoch:') >= 0:
+ # to-do, actually validate this works, and probably kill training when it's found, the model's dead by this point
+ if ': nan' in line:
+ should_return = True
+
+ print("! NAN DETECTED !")
+ self.buffer.append("! NAN DETECTED !")
+
# easily rip out our stats...
match = re.findall(r'\b([a-z_0-9]+?)\b: +?([0-9]\.[0-9]+?e[+-]\d+|[\d,]+)\b', line)
if match and len(match) > 0:
@@ -824,13 +849,13 @@ def run_training(config_path, verbose=False, gpus=1, keep_x_past_datasets=0, pro
if result:
yield result
+ if progress is not None and message:
+ progress(percent, message)
+
if training_state:
training_state.process.stdout.close()
return_code = training_state.process.wait()
training_state = None
-
- #if return_code:
- # raise subprocess.CalledProcessError(return_code, cmd)
def get_training_losses():
global training_state
@@ -866,6 +891,9 @@ def reconnect_training(verbose=False, progress=gr.Progress(track_tqdm=True)):
if result:
yield result
+ if progress is not None and message:
+ progress(percent, message)
+
def stop_training():
global training_state
if training_state is None:
@@ -910,10 +938,10 @@ def convert_to_halfp():
def whisper_transcribe( file, language=None ):
# shouldn't happen, but it's for safety
if not whisper_model:
- load_whisper_model(language=language if language else b'en')
+ load_whisper_model(language=language)
if not args.whisper_cpp:
- return whisper_model.transcribe(file, language=language if language else "English")
+ return whisper_model.transcribe(file, language=language)
res = whisper_model.transcribe(file)
segments = whisper_model.extract_text_and_timestamps( res )
@@ -945,11 +973,8 @@ def prepare_dataset( files, outdir, language=None, progress=None ):
transcription = []
for file in enumerate_progress(files, desc="Iterating through voice files", progress=progress):
- print(f"Transcribing file: {file}")
-
- result = whisper_transcribe(file, language=language) # whisper_model.transcribe(file, language=language if language else "English")
+ result = whisper_transcribe(file, language=language)
results[os.path.basename(file)] = result
-
print(f"Transcribed file: {file}, {len(result['segments'])} found.")
waveform, sampling_rate = torchaudio.load(file)
@@ -988,7 +1013,7 @@ EPOCH_SCHEDULE = [ 9, 18, 25, 33 ]
def schedule_learning_rate( iterations, schedule=EPOCH_SCHEDULE ):
return [int(iterations * d) for d in schedule]
-def optimize_training_settings( epochs, learning_rate, text_ce_lr_weight, learning_rate_schedule, batch_size, gradient_accumulation_size, print_rate, save_rate, resume_path, half_p, bnb, source_model, voice ):
+def optimize_training_settings( epochs, learning_rate, text_ce_lr_weight, learning_rate_schedule, batch_size, gradient_accumulation_size, print_rate, save_rate, resume_path, half_p, bnb, workers, source_model, voice ):
name = f"{voice}-finetune"
dataset_name = f"{voice}-train"
dataset_path = f"./training/{voice}/train.txt"
@@ -1065,7 +1090,7 @@ def optimize_training_settings( epochs, learning_rate, text_ce_lr_weight, learni
messages
)
-def save_training_settings( iterations=None, learning_rate=None, text_ce_lr_weight=None, learning_rate_schedule=None, batch_size=None, gradient_accumulation_size=None, print_rate=None, save_rate=None, name=None, dataset_name=None, dataset_path=None, validation_name=None, validation_path=None, output_name=None, resume_path=None, half_p=None, bnb=None, source_model=None ):
+def save_training_settings( iterations=None, learning_rate=None, text_ce_lr_weight=None, learning_rate_schedule=None, batch_size=None, gradient_accumulation_size=None, print_rate=None, save_rate=None, name=None, dataset_name=None, dataset_path=None, validation_name=None, validation_path=None, output_name=None, resume_path=None, half_p=None, bnb=None, workers=None, source_model=None ):
if not source_model:
source_model = f"./models/tortoise/autoregressive{'_half' if half_p else ''}.pth"
@@ -1090,6 +1115,8 @@ def save_training_settings( iterations=None, learning_rate=None, text_ce_lr_weig
'float16': 'true' if half_p else 'false',
'bitsandbytes': 'true' if bnb else 'false',
+
+ 'workers': workers if workers else 2,
}
if resume_path:
@@ -1581,9 +1608,9 @@ def unload_tts():
global tts
if tts:
- print("Unloading TTS")
del tts
tts = None
+ print("Unloaded TTS")
do_gc()
def reload_tts( model=None ):
@@ -1656,55 +1683,44 @@ def unload_voicefixer():
global voicefixer
if voicefixer:
- print("Unloading Voicefixer")
del voicefixer
voicefixer = None
print("Unloaded Voicefixer")
do_gc()
-def load_whisper_model(name=None, progress=None, language=b'en'):
+def load_whisper_model(language=None, model_name=None, progress=None):
global whisper_model
- if not name:
- name = args.whisper_model
+ if not model_name:
+ model_name = args.whisper_model
else:
- args.whisper_model = name
+ args.whisper_model = model_name
save_args_settings()
- notify_progress(f"Loading Whisper model: {args.whisper_model}", progress)
+ if language and f'{model_name}.{language}' in WHISPER_SPECIALIZED_MODELS:
+ model_name = f'{model_name}.{language}'
+ print(f"Loading specialized model for language: {language}")
+
+ notify_progress(f"Loading Whisper model: {model_name}", progress)
if args.whisper_cpp:
from whispercpp import Whisper
- whisper_model = Whisper(name, models_dir='./models/', language=language.encode('ascii'))
+ if not language:
+ language = 'auto'
+
+ whisper_model = Whisper(model_name, models_dir='./models/', language=language.encode('ascii'))
else:
import whisper
- whisper_model = whisper.load_model(args.whisper_model)
+ whisper_model = whisper.load_model(model_name)
+
print("Loaded Whisper model")
def unload_whisper():
global whisper_model
if whisper_model:
- print("Unloading Whisper")
del whisper_model
whisper_model = None
print("Unloaded Whisper")
- do_gc()
-
-"""
-def update_whisper_model(name, progress=None):
- if not name:
- return
-
- args.whisper_model = name
- save_args_settings()
-
- global whisper_model
- if whisper_model:
- unload_whisper()
- load_whisper_model(name)
- else:
- args.whisper_model = name
- save_args_settings()
-"""
\ No newline at end of file
+ do_gc()
\ No newline at end of file
diff --git a/src/webui.py b/src/webui.py
index 89d9462..3ca76d0 100755
--- a/src/webui.py
+++ b/src/webui.py
@@ -268,6 +268,8 @@ def import_training_settings_proxy( voice ):
if "ext" in config and "bitsandbytes" in config["ext"]:
bnb = config["ext"]["bitsandbytes"]
+ workers = config['datasets']['train']['n_workers']
+
messages = "\n".join(messages)
return (
@@ -282,12 +284,13 @@ def import_training_settings_proxy( voice ):
resume_path,
half_p,
bnb,
+ workers,
source_model,
messages
)
-def save_training_settings_proxy( epochs, learning_rate, text_ce_lr_weight, learning_rate_schedule, batch_size, gradient_accumulation_size, print_rate, save_rate, resume_path, half_p, bnb, source_model, voice ):
+def save_training_settings_proxy( epochs, learning_rate, text_ce_lr_weight, learning_rate_schedule, batch_size, gradient_accumulation_size, print_rate, save_rate, resume_path, half_p, bnb, workers, source_model, voice ):
name = f"{voice}-finetune"
dataset_name = f"{voice}-train"
dataset_path = f"./training/{voice}/train.txt"
@@ -330,6 +333,7 @@ def save_training_settings_proxy( epochs, learning_rate, text_ce_lr_weight, lear
resume_path=resume_path,
half_p=half_p,
bnb=bnb,
+ workers=workers,
source_model=source_model,
))
return "\n".join(messages)
@@ -466,7 +470,7 @@ def setup_gradio():
with gr.Column():
dataset_settings = [
gr.Dropdown( choices=voice_list, label="Dataset Source", type="value", value=voice_list[0] if len(voice_list) > 0 else "" ),
- gr.Textbox(label="Language", placeholder="English")
+ gr.Textbox(label="Language", value="en")
]
prepare_dataset_button = gr.Button(value="Prepare")
with gr.Column():
@@ -499,11 +503,16 @@ def setup_gradio():
training_settings = training_settings + [
gr.Textbox(label="Resume State Path", placeholder="./training/${voice}-finetune/training_state/${last_state}.state"),
]
- training_halfp = gr.Checkbox(label="Half Precision", value=args.training_default_halfp)
- training_bnb = gr.Checkbox(label="BitsAndBytes", value=args.training_default_bnb)
+
+ with gr.Row():
+ training_halfp = gr.Checkbox(label="Half Precision", value=args.training_default_halfp)
+ training_bnb = gr.Checkbox(label="BitsAndBytes", value=args.training_default_bnb)
+
+ training_workers = gr.Number(label="Worker Processes", value=2, precision=0)
+
source_model = gr.Dropdown( choices=autoregressive_models, label="Source Model", type="value", value=autoregressive_models[0] )
dataset_list_dropdown = gr.Dropdown( choices=dataset_list, label="Dataset", type="value", value=dataset_list[0] if len(dataset_list) else "" )
- training_settings = training_settings + [ training_halfp, training_bnb, source_model, dataset_list_dropdown ]
+ training_settings = training_settings + [ training_halfp, training_bnb, training_workers, source_model, dataset_list_dropdown ]
with gr.Row():
refresh_dataset_list = gr.Button(value="Refresh Dataset List")
@@ -572,7 +581,7 @@ def setup_gradio():
autoregressive_model_dropdown = gr.Dropdown(choices=autoregressive_models, label="Autoregressive Model", value=args.autoregressive_model if args.autoregressive_model else autoregressive_models[0])
- 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)
+ whisper_model_dropdown = gr.Dropdown(WHISPER_MODELS, label="Whisper Model", value=args.whisper_model)
use_whisper_cpp = gr.Checkbox(label="Use Whisper.cpp", value=args.whisper_cpp)
exec_inputs = exec_inputs + [ autoregressive_model_dropdown, whisper_model_dropdown, use_whisper_cpp, training_halfp, training_bnb ]
@@ -797,7 +806,7 @@ def setup_gradio():
)
import_dataset_button.click(import_training_settings_proxy,
inputs=dataset_list_dropdown,
- outputs=training_settings[:11] + [save_yaml_output] #console_output
+ outputs=training_settings[:13] + [save_yaml_output] #console_output
)
save_yaml_button.click(save_training_settings_proxy,
inputs=training_settings,