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)