Spaces:
Build error
Build error
File size: 7,205 Bytes
1207342 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import os
import shutil
import torch
import traceback
from pathlib import Path
from multiprocessing import cpu_count
from functions.core_functions1 import clear_gpu_cache
from functions.logging_utils import remove_log_file
from functions.slice_utils import open_slice, close_slice
from utils.formatter import format_audio_list
from utils.gpt_train import train_gpt
def get_audio_files_from_folder(folder_path):
audio_files = []
for root, dirs, files in os.walk(folder_path):
for file in files:
if file.endswith(".wav") or file.endswith(".mp3") or file.endswith(".flac") or file.endswith(".m4a") or file.endswith(".webm"):
audio_files.append(os.path.join(root, file))
return audio_files
def preprocess_dataset(audio_path, audio_folder, language, whisper_model, out_path, train_csv, eval_csv, progress):
out_path = os.path.join(out_path, "dataset")
os.makedirs(out_path, exist_ok=True)
if audio_path is not None and audio_path != []:
try:
train_meta, eval_meta, audio_total_size = format_audio_list(audio_path, whisper_model=whisper_model, target_language=language, out_path=out_path, gradio_progress=progress)
except:
traceback.print_exc()
error = traceback.format_exc()
return f"The data processing was interrupted due to an error! Please check the console to verify the full error message! \n Error summary: {error}", "", ""
elif audio_folder is not None:
audio_files = get_audio_files_from_folder(audio_folder)
try:
train_meta, eval_meta, audio_total_size = format_audio_list(audio_files, whisper_model=whisper_model, target_language=language, out_path=out_path, gradio_progress=progress)
except:
traceback.print_exc()
error = traceback.format_exc()
return f"The data processing was interrupted due to an error! Please check the console to verify the full error message! \n Error summary: {error}", "", ""
else:
return "You should provide either audio files or a folder containing audio files!", "", ""
if audio_total_size < 120:
message = "The sum of the duration of the audios that you provided should be at least 2 minutes!"
print(message)
return message, "", ""
print("Dataset Processed!")
return "Dataset Processed!", train_meta, eval_meta
def train_model(custom_model, version, language, train_csv, eval_csv, num_epochs, batch_size, grad_accum, output_path, max_audio_length):
run_dir = Path(output_path) / "run"
if run_dir.exists():
os.remove(run_dir)
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!", "", "", "", "", ""
try:
max_audio_length = int(max_audio_length * 22050)
speaker_xtts_path, config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(custom_model, version, language, num_epochs, batch_size, grad_accum, train_csv, eval_csv, output_path=output_path, max_audio_length=max_audio_length)
except:
traceback.print_exc()
error = traceback.format_exc()
return f"The training was interrupted due to an error! Please check the console to check the full error message! \n Error summary: {error}", "", "", "", "", ""
ready_dir = Path(output_path) / "ready"
ft_xtts_checkpoint = os.path.join(exp_path, "best_model.pth")
shutil.copy(ft_xtts_checkpoint, ready_dir / "unoptimize_model.pth")
ft_xtts_checkpoint = os.path.join(ready_dir, "unoptimize_model.pth")
speaker_reference_path = Path(speaker_wav)
speaker_reference_new_path = ready_dir / "reference.wav"
shutil.copy(speaker_reference_path, speaker_reference_new_path)
print("Model training done!")
return "Model training done!", config_path, vocab_file, ft_xtts_checkpoint, speaker_xtts_path, speaker_reference_new_path
def optimize_model(out_path, clear_train_data):
out_path = Path(out_path)
ready_dir = out_path / "ready"
run_dir = out_path / "run"
dataset_dir = out_path / "dataset"
if clear_train_data in {"run", "all"} and run_dir.exists():
try:
shutil.rmtree(run_dir)
except PermissionError as e:
print(f"An error occurred while deleting {run_dir}: {e}")
if clear_train_data in {"dataset", "all"} and dataset_dir.exists():
try:
shutil.rmtree(dataset_dir)
except PermissionError as e:
print(f"An error occurred while deleting {dataset_dir}: {e}")
model_path = ready_dir / "unoptimize_model.pth"
if not model_path.is_file():
return "Unoptimized model not found in ready folder", ""
checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
del checkpoint["optimizer"]
for key in list(checkpoint["model"].keys()):
if "dvae" in key:
del checkpoint["model"][key]
os.remove(model_path)
optimized_model_file_name = "model.pth"
optimized_model = ready_dir / optimized_model_file_name
torch.save(checkpoint, optimized_model)
ft_xtts_checkpoint = str(optimized_model)
return f"Model optimized and saved at {ft_xtts_checkpoint}!", ft_xtts_checkpoint
def load_params(out_path):
path_output = Path(out_path)
dataset_path = path_output / "dataset"
if not dataset_path.exists():
return "The output folder does not exist!", "", "", ""
eval_train = dataset_path / "metadata_train.csv"
eval_csv = dataset_path / "metadata_eval.csv"
lang_file_path = dataset_path / "lang.txt"
current_language = None
if os.path.exists(lang_file_path):
with open(lang_file_path, 'r', encoding='utf-8') as existing_lang_file:
current_language = existing_lang_file.read().strip()
print(current_language)
return "The data has been updated", eval_train, eval_csv, current_language
def load_params_tts(out_path, version):
path_output = Path(out_path)
ready_dir = path_output / "ready"
xtts_config_path = ready_dir / "config.json"
xtts_vocab_path = ready_dir / "vocab.json"
xtts_checkpoint_path = ready_dir / "model.pth"
xtts_speaker_path = ready_dir / "speaker.pth"
speaker_reference_path = ready_dir / "reference.wav"
missing_files = []
if not xtts_config_path.exists():
missing_files.append(str(xtts_config_path))
if not xtts_vocab_path.exists():
missing_files.append(str(xtts_vocab_path))
if not xtts_checkpoint_path.exists():
missing_files.append(str(xtts_checkpoint_path))
if not xtts_speaker_path.exists():
missing_files.append(str(xtts_speaker_path))
if not speaker_reference_path.exists():
missing_files.append(str(speaker_reference_path))
if missing_files:
return f"The following files are missing from the ready folder: {', '.join(missing_files)}", "", "", "", "", ""
print("Loaded parameters for TTS.")
return "Loaded parameters for TTS.", str(xtts_checkpoint_path), str(xtts_config_path), str(xtts_vocab_path), str(xtts_speaker_path), str(speaker_reference_path) |