| 1 | import gc |
| 2 | import json |
| 3 | import os |
| 4 | import platform |
| 5 | import queue |
| 6 | import random |
| 7 | import re |
| 8 | import shutil |
| 9 | import signal |
| 10 | import subprocess |
| 11 | import sys |
| 12 | import tempfile |
| 13 | import threading |
| 14 | import time |
| 15 | from glob import glob |
| 16 | from importlib.resources import files |
| 17 | |
| 18 | import click |
| 19 | import gradio as gr |
| 20 | import librosa |
| 21 | import numpy as np |
| 22 | import psutil |
| 23 | import torch |
| 24 | import torchaudio |
| 25 | from cached_path import cached_path |
| 26 | from datasets import Dataset as Dataset_ |
| 27 | from datasets.arrow_writer import ArrowWriter |
| 28 | from safetensors.torch import load_file, save_file |
| 29 | from scipy.io import wavfile |
| 30 | |
| 31 | from f5_tts.api import F5TTS |
| 32 | from f5_tts.infer.utils_infer import transcribe |
| 33 | from f5_tts.model.utils import convert_char_to_pinyin |
| 34 | |
| 35 | |
| 36 | def _safe_project_path(base: str, name: str) -> str: |
| 37 | """Return the resolved absolute path of base/name, raising ValueError if name |
| 38 | is absolute, contains a null byte, or resolves outside base.""" |
| 39 | if not name or os.path.isabs(name) or "\x00" in name: |
| 40 | raise ValueError(f"invalid project_name: {name!r}") |
| 41 | # Strip path separators and control characters to a plain filename component. |
| 42 | name = re.sub(r"[/\\]", "", name).strip() |
| 43 | if not name or name in (".", ".."): |
| 44 | raise ValueError(f"invalid project_name: {name!r}") |
| 45 | candidate = os.path.realpath(os.path.join(base, name)) |
| 46 | base_real = os.path.realpath(base) |
| 47 | if not (candidate + os.sep).startswith(base_real + os.sep): |
| 48 | raise ValueError(f"project_name escapes base directory: {name!r}") |
| 49 | return candidate |
| 50 | |
| 51 | |
| 52 | training_process = None |
| 53 | system = platform.system() |
| 54 | python_executable = sys.executable or "python" |
| 55 | tts_api = None |
| 56 | last_checkpoint = "" |
| 57 | last_device = "" |
| 58 | last_ema = None |
| 59 | |
| 60 | |
| 61 | path_data = str(files("f5_tts").joinpath("../../data")) |
| 62 | path_project_ckpts = str(files("f5_tts").joinpath("../../ckpts")) |
| 63 | file_train = str(files("f5_tts").joinpath("train/finetune_cli.py")) |
| 64 | |
| 65 | device = ( |
| 66 | "cuda" |
| 67 | if torch.cuda.is_available() |
| 68 | else "xpu" |
| 69 | if torch.xpu.is_available() |
| 70 | else "mps" |
| 71 | if torch.backends.mps.is_available() |
| 72 | else "cpu" |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | # Save settings from a JSON file |
| 77 | def save_settings( |
| 78 | project_name, |
| 79 | exp_name, |
| 80 | learning_rate, |
| 81 | batch_size_per_gpu, |
| 82 | batch_size_type, |
| 83 | max_samples, |
| 84 | grad_accumulation_steps, |
| 85 | max_grad_norm, |
| 86 | epochs, |
| 87 | num_warmup_updates, |
| 88 | save_per_updates, |
| 89 | keep_last_n_checkpoints, |
| 90 | last_per_updates, |
| 91 | finetune, |
| 92 | file_checkpoint_train, |
| 93 | tokenizer_type, |
| 94 | tokenizer_file, |
| 95 | mixed_precision, |
| 96 | logger, |
| 97 | ch_8bit_adam, |
| 98 | ): |
| 99 | path_project = _safe_project_path(path_project_ckpts, project_name) |
| 100 | os.makedirs(path_project, exist_ok=True) |
| 101 | file_setting = os.path.join(path_project, "setting.json") |
| 102 | |
| 103 | settings = { |
| 104 | "exp_name": exp_name, |
| 105 | "learning_rate": learning_rate, |
| 106 | "batch_size_per_gpu": batch_size_per_gpu, |
| 107 | "batch_size_type": batch_size_type, |
| 108 | "max_samples": max_samples, |
| 109 | "grad_accumulation_steps": grad_accumulation_steps, |
| 110 | "max_grad_norm": max_grad_norm, |
| 111 | "epochs": epochs, |
| 112 | "num_warmup_updates": num_warmup_updates, |
| 113 | "save_per_updates": save_per_updates, |
| 114 | "keep_last_n_checkpoints": keep_last_n_checkpoints, |
| 115 | "last_per_updates": last_per_updates, |
| 116 | "finetune": finetune, |
| 117 | "file_checkpoint_train": file_checkpoint_train, |
| 118 | "tokenizer_type": tokenizer_type, |
| 119 | "tokenizer_file": tokenizer_file, |
| 120 | "mixed_precision": mixed_precision, |
| 121 | "logger": logger, |
| 122 | "bnb_optimizer": ch_8bit_adam, |
| 123 | } |
| 124 | with open(file_setting, "w") as f: |
| 125 | json.dump(settings, f, indent=4) |
| 126 | return "Settings saved!" |
| 127 | |
| 128 | |
| 129 | # Load settings from a JSON file |
| 130 | def load_settings(project_name): |
| 131 | project_name = project_name.replace("_pinyin", "").replace("_char", "") |
| 132 | path_project = _safe_project_path(path_project_ckpts, project_name) |
| 133 | file_setting = os.path.join(path_project, "setting.json") |
| 134 | |
| 135 | # Default settings |
| 136 | default_settings = { |
| 137 | "exp_name": "F5TTS_v1_Base", |
| 138 | "learning_rate": 1e-5, |
| 139 | "batch_size_per_gpu": 3200, |
| 140 | "batch_size_type": "frame", |
| 141 | "max_samples": 64, |
| 142 | "grad_accumulation_steps": 1, |
| 143 | "max_grad_norm": 1.0, |
| 144 | "epochs": 100, |
| 145 | "num_warmup_updates": 100, |
| 146 | "save_per_updates": 500, |
| 147 | "keep_last_n_checkpoints": -1, |
| 148 | "last_per_updates": 100, |
| 149 | "finetune": True, |
| 150 | "file_checkpoint_train": "", |
| 151 | "tokenizer_type": "pinyin", |
| 152 | "tokenizer_file": "", |
| 153 | "mixed_precision": "fp16", |
| 154 | "logger": "none", |
| 155 | "bnb_optimizer": False, |
| 156 | } |
| 157 | if device == "mps": |
| 158 | default_settings["mixed_precision"] = "none" |
| 159 | |
| 160 | # Load settings from file if it exists |
| 161 | if os.path.isfile(file_setting): |
| 162 | with open(file_setting, "r") as f: |
| 163 | file_settings = json.load(f) |
| 164 | default_settings.update(file_settings) |
| 165 | |
| 166 | # Return as a tuple in the correct order |
| 167 | return ( |
| 168 | default_settings["exp_name"], |
| 169 | default_settings["learning_rate"], |
| 170 | default_settings["batch_size_per_gpu"], |
| 171 | default_settings["batch_size_type"], |
| 172 | default_settings["max_samples"], |
| 173 | default_settings["grad_accumulation_steps"], |
| 174 | default_settings["max_grad_norm"], |
| 175 | default_settings["epochs"], |
| 176 | default_settings["num_warmup_updates"], |
| 177 | default_settings["save_per_updates"], |
| 178 | default_settings["keep_last_n_checkpoints"], |
| 179 | default_settings["last_per_updates"], |
| 180 | default_settings["finetune"], |
| 181 | default_settings["file_checkpoint_train"], |
| 182 | default_settings["tokenizer_type"], |
| 183 | default_settings["tokenizer_file"], |
| 184 | default_settings["mixed_precision"], |
| 185 | default_settings["logger"], |
| 186 | default_settings["bnb_optimizer"], |
| 187 | ) |
| 188 | |
| 189 | |
| 190 | # Load metadata |
| 191 | def get_audio_duration(audio_path): |
| 192 | """Calculate the duration mono of an audio file.""" |
| 193 | audio, sample_rate = torchaudio.load(audio_path) |
| 194 | return audio.shape[1] / sample_rate |
| 195 | |
| 196 | |
| 197 | class Slicer: # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.py |
| 198 | def __init__( |
| 199 | self, |
| 200 | sr: int, |
| 201 | threshold: float = -40.0, |
| 202 | min_length: int = 20000, # 20 seconds |
| 203 | min_interval: int = 300, |
| 204 | hop_size: int = 20, |
| 205 | max_sil_kept: int = 2000, |
| 206 | ): |
| 207 | if not min_length >= min_interval >= hop_size: |
| 208 | raise ValueError("The following condition must be satisfied: min_length >= min_interval >= hop_size") |
| 209 | if not max_sil_kept >= hop_size: |
| 210 | raise ValueError("The following condition must be satisfied: max_sil_kept >= hop_size") |
| 211 | min_interval = sr * min_interval / 1000 |
| 212 | self.threshold = 10 ** (threshold / 20.0) |
| 213 | self.hop_size = round(sr * hop_size / 1000) |
| 214 | self.win_size = min(round(min_interval), 4 * self.hop_size) |
| 215 | self.min_length = round(sr * min_length / 1000 / self.hop_size) |
| 216 | self.min_interval = round(min_interval / self.hop_size) |
| 217 | self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size) |
| 218 | |
| 219 | def _apply_slice(self, waveform, begin, end): |
| 220 | if len(waveform.shape) > 1: |
| 221 | return waveform[:, begin * self.hop_size : min(waveform.shape[1], end * self.hop_size)] |
| 222 | else: |
| 223 | return waveform[begin * self.hop_size : min(waveform.shape[0], end * self.hop_size)] |
| 224 | |
| 225 | # @timeit |
| 226 | def slice(self, waveform): |
| 227 | if len(waveform.shape) > 1: |
| 228 | samples = waveform.mean(axis=0) |
| 229 | else: |
| 230 | samples = waveform |
| 231 | if samples.shape[0] <= self.min_length: |
| 232 | return [waveform] |
| 233 | rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0) |
| 234 | sil_tags = [] |
| 235 | silence_start = None |
| 236 | clip_start = 0 |
| 237 | for i, rms in enumerate(rms_list): |
| 238 | # Keep looping while frame is silent. |
| 239 | if rms < self.threshold: |
| 240 | # Record start of silent frames. |
| 241 | if silence_start is None: |
| 242 | silence_start = i |
| 243 | continue |
| 244 | # Keep looping while frame is not silent and silence start has not been recorded. |
| 245 | if silence_start is None: |
| 246 | continue |
| 247 | # Clear recorded silence start if interval is not enough or clip is too short |
| 248 | is_leading_silence = silence_start == 0 and i > self.max_sil_kept |
| 249 | need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length |
| 250 | if not is_leading_silence and not need_slice_middle: |
| 251 | silence_start = None |
| 252 | continue |
| 253 | # Need slicing. Record the range of silent frames to be removed. |
| 254 | if i - silence_start <= self.max_sil_kept: |
| 255 | pos = rms_list[silence_start : i + 1].argmin() + silence_start |
| 256 | if silence_start == 0: |
| 257 | sil_tags.append((0, pos)) |
| 258 | else: |
| 259 | sil_tags.append((pos, pos)) |
| 260 | clip_start = pos |
| 261 | elif i - silence_start <= self.max_sil_kept * 2: |
| 262 | pos = rms_list[i - self.max_sil_kept : silence_start + self.max_sil_kept + 1].argmin() |
| 263 | pos += i - self.max_sil_kept |
| 264 | pos_l = rms_list[silence_start : silence_start + self.max_sil_kept + 1].argmin() + silence_start |
| 265 | pos_r = rms_list[i - self.max_sil_kept : i + 1].argmin() + i - self.max_sil_kept |
| 266 | if silence_start == 0: |
| 267 | sil_tags.append((0, pos_r)) |
| 268 | clip_start = pos_r |
| 269 | else: |
| 270 | sil_tags.append((min(pos_l, pos), max(pos_r, pos))) |
| 271 | clip_start = max(pos_r, pos) |
| 272 | else: |
| 273 | pos_l = rms_list[silence_start : silence_start + self.max_sil_kept + 1].argmin() + silence_start |
| 274 | pos_r = rms_list[i - self.max_sil_kept : i + 1].argmin() + i - self.max_sil_kept |
| 275 | if silence_start == 0: |
| 276 | sil_tags.append((0, pos_r)) |
| 277 | else: |
| 278 | sil_tags.append((pos_l, pos_r)) |
| 279 | clip_start = pos_r |
| 280 | silence_start = None |
| 281 | # Deal with trailing silence. |
| 282 | total_frames = rms_list.shape[0] |
| 283 | if silence_start is not None and total_frames - silence_start >= self.min_interval: |
| 284 | silence_end = min(total_frames, silence_start + self.max_sil_kept) |
| 285 | pos = rms_list[silence_start : silence_end + 1].argmin() + silence_start |
| 286 | sil_tags.append((pos, total_frames + 1)) |
| 287 | # Apply and return slices: [chunk, start, end] |
| 288 | if len(sil_tags) == 0: |
| 289 | return [[waveform, 0, int(total_frames * self.hop_size)]] |
| 290 | else: |
| 291 | chunks = [] |
| 292 | if sil_tags[0][0] > 0: |
| 293 | chunks.append([self._apply_slice(waveform, 0, sil_tags[0][0]), 0, int(sil_tags[0][0] * self.hop_size)]) |
| 294 | for i in range(len(sil_tags) - 1): |
| 295 | chunks.append( |
| 296 | [ |
| 297 | self._apply_slice(waveform, sil_tags[i][1], sil_tags[i + 1][0]), |
| 298 | int(sil_tags[i][1] * self.hop_size), |
| 299 | int(sil_tags[i + 1][0] * self.hop_size), |
| 300 | ] |
| 301 | ) |
| 302 | if sil_tags[-1][1] < total_frames: |
| 303 | chunks.append( |
| 304 | [ |
| 305 | self._apply_slice(waveform, sil_tags[-1][1], total_frames), |
| 306 | int(sil_tags[-1][1] * self.hop_size), |
| 307 | int(total_frames * self.hop_size), |
| 308 | ] |
| 309 | ) |
| 310 | return chunks |
| 311 | |
| 312 | |
| 313 | # terminal |
| 314 | def terminate_process_tree(pid, including_parent=True): |
| 315 | try: |
| 316 | parent = psutil.Process(pid) |
| 317 | except psutil.NoSuchProcess: |
| 318 | # Process already terminated |
| 319 | return |
| 320 | |
| 321 | children = parent.children(recursive=True) |
| 322 | for child in children: |
| 323 | try: |
| 324 | os.kill(child.pid, signal.SIGTERM) # or signal.SIGKILL |
| 325 | except OSError: |
| 326 | pass |
| 327 | if including_parent: |
| 328 | try: |
| 329 | os.kill(parent.pid, signal.SIGTERM) # or signal.SIGKILL |
| 330 | except OSError: |
| 331 | pass |
| 332 | |
| 333 | |
| 334 | def terminate_process(pid): |
| 335 | if system == "Windows": |
| 336 | cmd = f"taskkill /t /f /pid {pid}" |
| 337 | os.system(cmd) |
| 338 | else: |
| 339 | terminate_process_tree(pid) |
| 340 | |
| 341 | |
| 342 | def start_training( |
| 343 | dataset_name, |
| 344 | exp_name, |
| 345 | learning_rate, |
| 346 | batch_size_per_gpu, |
| 347 | batch_size_type, |
| 348 | max_samples, |
| 349 | grad_accumulation_steps, |
| 350 | max_grad_norm, |
| 351 | epochs, |
| 352 | num_warmup_updates, |
| 353 | save_per_updates, |
| 354 | keep_last_n_checkpoints, |
| 355 | last_per_updates, |
| 356 | finetune, |
| 357 | file_checkpoint_train, |
| 358 | tokenizer_type, |
| 359 | tokenizer_file, |
| 360 | mixed_precision, |
| 361 | stream, |
| 362 | logger, |
| 363 | ch_8bit_adam, |
| 364 | ): |
| 365 | global training_process, tts_api, stop_signal |
| 366 | |
| 367 | if tts_api is not None: |
| 368 | if tts_api is not None: |
| 369 | del tts_api |
| 370 | |
| 371 | gc.collect() |
| 372 | torch.cuda.empty_cache() |
| 373 | tts_api = None |
| 374 | |
| 375 | path_project = _safe_project_path(path_data, dataset_name) |
| 376 | |
| 377 | if not os.path.isdir(path_project): |
| 378 | yield ( |
| 379 | f"There is not project with name {dataset_name}", |
| 380 | gr.update(interactive=True), |
| 381 | gr.update(interactive=False), |
| 382 | ) |
| 383 | return |
| 384 | |
| 385 | file_raw = os.path.join(path_project, "raw.arrow") |
| 386 | if not os.path.isfile(file_raw): |
| 387 | yield f"There is no file {file_raw}", gr.update(interactive=True), gr.update(interactive=False) |
| 388 | return |
| 389 | |
| 390 | # Check if a training process is already running |
| 391 | if training_process is not None: |
| 392 | return "Train run already!", gr.update(interactive=False), gr.update(interactive=True) |
| 393 | |
| 394 | yield "start train", gr.update(interactive=False), gr.update(interactive=False) |
| 395 | |
| 396 | # Command to run the training script with the specified arguments |
| 397 | |
| 398 | if tokenizer_file == "": |
| 399 | if dataset_name.endswith("_pinyin"): |
| 400 | tokenizer_type = "pinyin" |
| 401 | elif dataset_name.endswith("_char"): |
| 402 | tokenizer_type = "char" |
| 403 | else: |
| 404 | tokenizer_type = "custom" |
| 405 | |
| 406 | dataset_name = dataset_name.replace("_pinyin", "").replace("_char", "") |
| 407 | |
| 408 | if mixed_precision != "none": |
| 409 | fp16 = f"--mixed_precision={mixed_precision}" |
| 410 | else: |
| 411 | fp16 = "" |
| 412 | |
| 413 | cmd = ( |
| 414 | f'accelerate launch {fp16} "{file_train}" --exp_name {exp_name}' |
| 415 | f" --learning_rate {learning_rate}" |
| 416 | f" --batch_size_per_gpu {batch_size_per_gpu}" |
| 417 | f" --batch_size_type {batch_size_type}" |
| 418 | f" --max_samples {max_samples}" |
| 419 | f" --grad_accumulation_steps {grad_accumulation_steps}" |
| 420 | f" --max_grad_norm {max_grad_norm}" |
| 421 | f" --epochs {epochs}" |
| 422 | f" --num_warmup_updates {num_warmup_updates}" |
| 423 | f" --save_per_updates {save_per_updates}" |
| 424 | f" --keep_last_n_checkpoints {keep_last_n_checkpoints}" |
| 425 | f" --last_per_updates {last_per_updates}" |
| 426 | f" --dataset_name {dataset_name}" |
| 427 | ) |
| 428 | |
| 429 | if finetune: |
| 430 | cmd += " --finetune" |
| 431 | |
| 432 | if file_checkpoint_train != "": |
| 433 | cmd += f' --pretrain "{file_checkpoint_train}"' |
| 434 | |
| 435 | if tokenizer_file != "": |
| 436 | cmd += f" --tokenizer_path {tokenizer_file}" |
| 437 | |
| 438 | cmd += f" --tokenizer {tokenizer_type}" |
| 439 | |
| 440 | if logger != "none": |
| 441 | cmd += f" --logger {logger}" |
| 442 | |
| 443 | cmd += " --log_samples" |
| 444 | |
| 445 | if ch_8bit_adam: |
| 446 | cmd += " --bnb_optimizer" |
| 447 | |
| 448 | print("run command : \n" + cmd + "\n") |
| 449 | |
| 450 | save_settings( |
| 451 | dataset_name, |
| 452 | exp_name, |
| 453 | learning_rate, |
| 454 | batch_size_per_gpu, |
| 455 | batch_size_type, |
| 456 | max_samples, |
| 457 | grad_accumulation_steps, |
| 458 | max_grad_norm, |
| 459 | epochs, |
| 460 | num_warmup_updates, |
| 461 | save_per_updates, |
| 462 | keep_last_n_checkpoints, |
| 463 | last_per_updates, |
| 464 | finetune, |
| 465 | file_checkpoint_train, |
| 466 | tokenizer_type, |
| 467 | tokenizer_file, |
| 468 | mixed_precision, |
| 469 | logger, |
| 470 | ch_8bit_adam, |
| 471 | ) |
| 472 | |
| 473 | try: |
| 474 | if not stream: |
| 475 | # Start the training process |
| 476 | training_process = subprocess.Popen(cmd, shell=True) |
| 477 | |
| 478 | time.sleep(5) |
| 479 | yield "train start", gr.update(interactive=False), gr.update(interactive=True) |
| 480 | |
| 481 | # Wait for the training process to finish |
| 482 | training_process.wait() |
| 483 | else: |
| 484 | |
| 485 | def stream_output(pipe, output_queue): |
| 486 | try: |
| 487 | for line in iter(pipe.readline, ""): |
| 488 | output_queue.put(line) |
| 489 | except Exception as e: |
| 490 | output_queue.put(f"Error reading pipe: {str(e)}") |
| 491 | finally: |
| 492 | pipe.close() |
| 493 | |
| 494 | env = os.environ.copy() |
| 495 | env["PYTHONUNBUFFERED"] = "1" |
| 496 | |
| 497 | training_process = subprocess.Popen( |
| 498 | cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, env=env |
| 499 | ) |
| 500 | yield "Training started ...", gr.update(interactive=False), gr.update(interactive=True) |
| 501 | |
| 502 | stdout_queue = queue.Queue() |
| 503 | stderr_queue = queue.Queue() |
| 504 | |
| 505 | stdout_thread = threading.Thread(target=stream_output, args=(training_process.stdout, stdout_queue)) |
| 506 | stderr_thread = threading.Thread(target=stream_output, args=(training_process.stderr, stderr_queue)) |
| 507 | stdout_thread.daemon = True |
| 508 | stderr_thread.daemon = True |
| 509 | stdout_thread.start() |
| 510 | stderr_thread.start() |
| 511 | stop_signal = False |
| 512 | while True: |
| 513 | if stop_signal: |
| 514 | training_process.terminate() |
| 515 | time.sleep(0.5) |
| 516 | if training_process.poll() is None: |
| 517 | training_process.kill() |
| 518 | yield "Training stopped by user.", gr.update(interactive=True), gr.update(interactive=False) |
| 519 | break |
| 520 | |
| 521 | process_status = training_process.poll() |
| 522 | |
| 523 | # Handle stdout |
| 524 | try: |
| 525 | while True: |
| 526 | output = stdout_queue.get_nowait() |
| 527 | print(output, end="") |
| 528 | match = re.search( |
| 529 | r"Epoch (\d+)/(\d+):\s+(\d+)%\|.*\[(\d+:\d+)<.*?loss=(\d+\.\d+), update=(\d+)", output |
| 530 | ) |
| 531 | if match: |
| 532 | current_epoch = match.group(1) |
| 533 | total_epochs = match.group(2) |
| 534 | percent_complete = match.group(3) |
| 535 | elapsed_time = match.group(4) |
| 536 | loss = match.group(5) |
| 537 | current_update = match.group(6) |
| 538 | message = ( |
| 539 | f"Epoch: {current_epoch}/{total_epochs}, " |
| 540 | f"Progress: {percent_complete}%, " |
| 541 | f"Elapsed Time: {elapsed_time}, " |
| 542 | f"Loss: {loss}, " |
| 543 | f"Update: {current_update}" |
| 544 | ) |
| 545 | yield message, gr.update(interactive=False), gr.update(interactive=True) |
| 546 | elif output.strip(): |
| 547 | yield output, gr.update(interactive=False), gr.update(interactive=True) |
| 548 | except queue.Empty: |
| 549 | pass |
| 550 | |
| 551 | # Handle stderr |
| 552 | try: |
| 553 | while True: |
| 554 | error_output = stderr_queue.get_nowait() |
| 555 | print(error_output, end="") |
| 556 | if error_output.strip(): |
| 557 | yield f"{error_output.strip()}", gr.update(interactive=False), gr.update(interactive=True) |
| 558 | except queue.Empty: |
| 559 | pass |
| 560 | |
| 561 | if process_status is not None and stdout_queue.empty() and stderr_queue.empty(): |
| 562 | if process_status != 0: |
| 563 | yield ( |
| 564 | f"Process crashed with exit code {process_status}!", |
| 565 | gr.update(interactive=False), |
| 566 | gr.update(interactive=True), |
| 567 | ) |
| 568 | else: |
| 569 | yield ( |
| 570 | "Training complete or paused ...", |
| 571 | gr.update(interactive=False), |
| 572 | gr.update(interactive=True), |
| 573 | ) |
| 574 | break |
| 575 | |
| 576 | # Small sleep to prevent CPU thrashing |
| 577 | time.sleep(0.1) |
| 578 | |
| 579 | # Clean up |
| 580 | training_process.stdout.close() |
| 581 | training_process.stderr.close() |
| 582 | training_process.wait() |
| 583 | |
| 584 | time.sleep(1) |
| 585 | |
| 586 | if training_process is None: |
| 587 | text_info = "Train stopped !" |
| 588 | else: |
| 589 | text_info = "Train complete at end !" |
| 590 | |
| 591 | except Exception as e: # Catch all exceptions |
| 592 | # Ensure that we reset the training process variable in case of an error |
| 593 | text_info = f"An error occurred: {str(e)}" |
| 594 | |
| 595 | training_process = None |
| 596 | |
| 597 | yield text_info, gr.update(interactive=True), gr.update(interactive=False) |
| 598 | |
| 599 | |
| 600 | def stop_training(): |
| 601 | global training_process, stop_signal |
| 602 | |
| 603 | if training_process is None: |
| 604 | return "Train not running !", gr.update(interactive=True), gr.update(interactive=False) |
| 605 | terminate_process_tree(training_process.pid) |
| 606 | # training_process = None |
| 607 | stop_signal = True |
| 608 | return "Train stopped !", gr.update(interactive=True), gr.update(interactive=False) |
| 609 | |
| 610 | |
| 611 | def get_list_projects(): |
| 612 | project_list = [] |
| 613 | for folder in os.listdir(path_data): |
| 614 | path_folder = os.path.join(path_data, folder) |
| 615 | if not os.path.isdir(path_folder): |
| 616 | continue |
| 617 | folder = folder.lower() |
| 618 | if folder == "emilia_zh_en_pinyin": |
| 619 | continue |
| 620 | project_list.append(folder) |
| 621 | |
| 622 | projects_selelect = None if not project_list else project_list[-1] |
| 623 | |
| 624 | return project_list, projects_selelect |
| 625 | |
| 626 | |
| 627 | def create_data_project(name, tokenizer_type): |
| 628 | name += "_" + tokenizer_type |
| 629 | project_dir = _safe_project_path(path_data, name) |
| 630 | os.makedirs(project_dir, exist_ok=True) |
| 631 | os.makedirs(os.path.join(project_dir, "dataset"), exist_ok=True) |
| 632 | project_list, projects_selelect = get_list_projects() |
| 633 | return gr.update(choices=project_list, value=name) |
| 634 | |
| 635 | |
| 636 | def transcribe_all(name_project, audio_files, language, user=False, progress=gr.Progress()): |
| 637 | path_project = _safe_project_path(path_data, name_project) |
| 638 | path_dataset = os.path.join(path_project, "dataset") |
| 639 | path_project_wavs = os.path.join(path_project, "wavs") |
| 640 | file_metadata = os.path.join(path_project, "metadata.csv") |
| 641 | |
| 642 | if not user: |
| 643 | if audio_files is None: |
| 644 | return "You need to load an audio file." |
| 645 | |
| 646 | if os.path.isdir(path_project_wavs): |
| 647 | shutil.rmtree(path_project_wavs) |
| 648 | |
| 649 | if os.path.isfile(file_metadata): |
| 650 | os.remove(file_metadata) |
| 651 | |
| 652 | os.makedirs(path_project_wavs, exist_ok=True) |
| 653 | |
| 654 | if user: |
| 655 | file_audios = [ |
| 656 | file |
| 657 | for format in ("*.wav", "*.ogg", "*.opus", "*.mp3", "*.flac") |
| 658 | for file in glob(os.path.join(path_dataset, format)) |
| 659 | ] |
| 660 | if file_audios == []: |
| 661 | return "No audio file was found in the dataset." |
| 662 | else: |
| 663 | file_audios = audio_files |
| 664 | |
| 665 | alpha = 0.5 |
| 666 | _max = 1.0 |
| 667 | slicer = Slicer(24000) |
| 668 | |
| 669 | num = 0 |
| 670 | error_num = 0 |
| 671 | data = "" |
| 672 | for file_audio in progress.tqdm(file_audios, desc="transcribe files", total=len((file_audios))): |
| 673 | audio, _ = librosa.load(file_audio, sr=24000, mono=True) |
| 674 | |
| 675 | list_slicer = slicer.slice(audio) |
| 676 | for chunk, start, end in progress.tqdm(list_slicer, total=len(list_slicer), desc="slicer files"): |
| 677 | name_segment = os.path.join(f"segment_{num}") |
| 678 | file_segment = os.path.join(path_project_wavs, f"{name_segment}.wav") |
| 679 | |
| 680 | tmp_max = np.abs(chunk).max() |
| 681 | if tmp_max > 1: |
| 682 | chunk /= tmp_max |
| 683 | chunk = (chunk / tmp_max * (_max * alpha)) + (1 - alpha) * chunk |
| 684 | wavfile.write(file_segment, 24000, (chunk * 32767).astype(np.int16)) |
| 685 | |
| 686 | try: |
| 687 | text = transcribe(file_segment, language) |
| 688 | text = text.strip() |
| 689 | |
| 690 | data += f"{name_segment}|{text}\n" |
| 691 | |
| 692 | num += 1 |
| 693 | except: # noqa: E722 |
| 694 | error_num += 1 |
| 695 | |
| 696 | with open(file_metadata, "w", encoding="utf-8-sig") as f: |
| 697 | f.write(data) |
| 698 | |
| 699 | if error_num != []: |
| 700 | error_text = f"\nerror files : {error_num}" |
| 701 | else: |
| 702 | error_text = "" |
| 703 | |
| 704 | return f"transcribe complete samples : {num}\npath : {path_project_wavs}{error_text}" |
| 705 | |
| 706 | |
| 707 | def format_seconds_to_hms(seconds): |
| 708 | hours = int(seconds / 3600) |
| 709 | minutes = int((seconds % 3600) / 60) |
| 710 | seconds = seconds % 60 |
| 711 | return "{:02d}:{:02d}:{:02d}".format(hours, minutes, int(seconds)) |
| 712 | |
| 713 | |
| 714 | def get_correct_audio_path( |
| 715 | audio_input, |
| 716 | base_path="wavs", |
| 717 | supported_formats=("wav", "mp3", "aac", "flac", "m4a", "alac", "ogg", "aiff", "wma", "amr"), |
| 718 | ): |
| 719 | file_audio = None |
| 720 | |
| 721 | # Helper function to check if file has a supported extension |
| 722 | def has_supported_extension(file_name): |
| 723 | return any(file_name.endswith(f".{ext}") for ext in supported_formats) |
| 724 | |
| 725 | # Case 1: If it's a full path with a valid extension, use it directly |
| 726 | if os.path.isabs(audio_input) and has_supported_extension(audio_input): |
| 727 | file_audio = audio_input |
| 728 | |
| 729 | # Case 2: If it has a supported extension but is not a full path |
| 730 | elif has_supported_extension(audio_input) and not os.path.isabs(audio_input): |
| 731 | file_audio = os.path.join(base_path, audio_input) |
| 732 | |
| 733 | # Case 3: If only the name is given (no extension and not a full path) |
| 734 | elif not has_supported_extension(audio_input) and not os.path.isabs(audio_input): |
| 735 | for ext in supported_formats: |
| 736 | potential_file = os.path.join(base_path, f"{audio_input}.{ext}") |
| 737 | if os.path.exists(potential_file): |
| 738 | file_audio = potential_file |
| 739 | break |
| 740 | else: |
| 741 | file_audio = os.path.join(base_path, f"{audio_input}.{supported_formats[0]}") |
| 742 | return file_audio |
| 743 | |
| 744 | |
| 745 | def create_metadata(name_project, ch_tokenizer, progress=gr.Progress()): |
| 746 | path_project = _safe_project_path(path_data, name_project) |
| 747 | path_project_wavs = os.path.join(path_project, "wavs") |
| 748 | file_metadata = os.path.join(path_project, "metadata.csv") |
| 749 | file_raw = os.path.join(path_project, "raw.arrow") |
| 750 | file_duration = os.path.join(path_project, "duration.json") |
| 751 | file_vocab = os.path.join(path_project, "vocab.txt") |
| 752 | |
| 753 | if not os.path.isfile(file_metadata): |
| 754 | return "The file was not found in " + file_metadata, "" |
| 755 | |
| 756 | with open(file_metadata, "r", encoding="utf-8-sig") as f: |
| 757 | data = f.read() |
| 758 | |
| 759 | audio_path_list = [] |
| 760 | text_list = [] |
| 761 | duration_list = [] |
| 762 | |
| 763 | count = data.split("\n") |
| 764 | lenght = 0 |
| 765 | result = [] |
| 766 | error_files = [] |
| 767 | text_vocab_set = set() |
| 768 | for line in progress.tqdm(data.split("\n"), total=count): |
| 769 | sp_line = line.split("|") |
| 770 | if len(sp_line) != 2: |
| 771 | continue |
| 772 | name_audio, text = sp_line[:2] |
| 773 | |
| 774 | file_audio = get_correct_audio_path(name_audio, path_project_wavs) |
| 775 | |
| 776 | if not os.path.isfile(file_audio): |
| 777 | error_files.append([file_audio, "error path"]) |
| 778 | continue |
| 779 | |
| 780 | try: |
| 781 | duration = get_audio_duration(file_audio) |
| 782 | except Exception as e: |
| 783 | error_files.append([file_audio, "duration"]) |
| 784 | print(f"Error processing {file_audio}: {e}") |
| 785 | continue |
| 786 | |
| 787 | if duration < 1 or duration > 30: |
| 788 | if duration > 30: |
| 789 | error_files.append([file_audio, "duration > 30 sec"]) |
| 790 | if duration < 1: |
| 791 | error_files.append([file_audio, "duration < 1 sec "]) |
| 792 | continue |
| 793 | if len(text) < 3: |
| 794 | error_files.append([file_audio, "very short text length 3"]) |
| 795 | continue |
| 796 | |
| 797 | text = text.strip() |
| 798 | text = convert_char_to_pinyin([text], polyphone=True)[0] |
| 799 | |
| 800 | audio_path_list.append(file_audio) |
| 801 | duration_list.append(duration) |
| 802 | text_list.append(text) |
| 803 | |
| 804 | result.append({"audio_path": file_audio, "text": text, "duration": duration}) |
| 805 | if ch_tokenizer: |
| 806 | text_vocab_set.update(list(text)) |
| 807 | |
| 808 | lenght += duration |
| 809 | |
| 810 | if duration_list == []: |
| 811 | return f"Error: No audio files found in the specified path : {path_project_wavs}", "" |
| 812 | |
| 813 | min_second = round(min(duration_list), 2) |
| 814 | max_second = round(max(duration_list), 2) |
| 815 | |
| 816 | with ArrowWriter(path=file_raw) as writer: |
| 817 | for line in progress.tqdm(result, total=len(result), desc="prepare data"): |
| 818 | writer.write(line) |
| 819 | writer.finalize() |
| 820 | |
| 821 | with open(file_duration, "w") as f: |
| 822 | json.dump({"duration": duration_list}, f, ensure_ascii=False) |
| 823 | |
| 824 | new_vocal = "" |
| 825 | if not ch_tokenizer: |
| 826 | if not os.path.isfile(file_vocab): |
| 827 | file_vocab_finetune = os.path.join(path_data, "Emilia_ZH_EN_pinyin/vocab.txt") |
| 828 | if not os.path.isfile(file_vocab_finetune): |
| 829 | return "Error: Vocabulary file 'Emilia_ZH_EN_pinyin' not found!", "" |
| 830 | shutil.copy2(file_vocab_finetune, file_vocab) |
| 831 | |
| 832 | with open(file_vocab, "r", encoding="utf-8-sig") as f: |
| 833 | vocab_char_map = {} |
| 834 | for i, char in enumerate(f): |
| 835 | vocab_char_map[char[:-1]] = i |
| 836 | vocab_size = len(vocab_char_map) |
| 837 | |
| 838 | else: |
| 839 | with open(file_vocab, "w", encoding="utf-8-sig") as f: |
| 840 | for vocab in sorted(text_vocab_set): |
| 841 | f.write(vocab + "\n") |
| 842 | new_vocal += vocab + "\n" |
| 843 | vocab_size = len(text_vocab_set) |
| 844 | |
| 845 | if error_files != []: |
| 846 | error_text = "\n".join([" = ".join(item) for item in error_files]) |
| 847 | else: |
| 848 | error_text = "" |
| 849 | |
| 850 | return ( |
| 851 | f"prepare complete \nsamples : {len(text_list)}\ntime data : {format_seconds_to_hms(lenght)}\nmin sec : {min_second}\nmax sec : {max_second}\nfile_arrow : {file_raw}\nvocab : {vocab_size}\n{error_text}", |
| 852 | new_vocal, |
| 853 | ) |
| 854 | |
| 855 | |
| 856 | def check_user(value): |
| 857 | return gr.update(visible=not value), gr.update(visible=value) |
| 858 | |
| 859 | |
| 860 | def calculate_train( |
| 861 | name_project, |
| 862 | epochs, |
| 863 | learning_rate, |
| 864 | batch_size_per_gpu, |
| 865 | batch_size_type, |
| 866 | max_samples, |
| 867 | num_warmup_updates, |
| 868 | finetune, |
| 869 | ): |
| 870 | path_project = _safe_project_path(path_data, name_project) |
| 871 | file_duration = os.path.join(path_project, "duration.json") |
| 872 | |
| 873 | hop_length = 256 |
| 874 | sampling_rate = 24000 |
| 875 | |
| 876 | if not os.path.isfile(file_duration): |
| 877 | return ( |
| 878 | epochs, |
| 879 | learning_rate, |
| 880 | batch_size_per_gpu, |
| 881 | max_samples, |
| 882 | num_warmup_updates, |
| 883 | "project not found !", |
| 884 | ) |
| 885 | |
| 886 | with open(file_duration, "r") as file: |
| 887 | data = json.load(file) |
| 888 | |
| 889 | duration_list = data["duration"] |
| 890 | max_sample_length = max(duration_list) * sampling_rate / hop_length |
| 891 | total_samples = len(duration_list) |
| 892 | total_duration = sum(duration_list) |
| 893 | |
| 894 | if torch.cuda.is_available(): |
| 895 | gpu_count = torch.cuda.device_count() |
| 896 | total_memory = 0 |
| 897 | for i in range(gpu_count): |
| 898 | gpu_properties = torch.cuda.get_device_properties(i) |
| 899 | total_memory += gpu_properties.total_memory / (1024**3) # in GB |
| 900 | elif torch.xpu.is_available(): |
| 901 | gpu_count = torch.xpu.device_count() |
| 902 | total_memory = 0 |
| 903 | for i in range(gpu_count): |
| 904 | gpu_properties = torch.xpu.get_device_properties(i) |
| 905 | total_memory += gpu_properties.total_memory / (1024**3) |
| 906 | elif torch.backends.mps.is_available(): |
| 907 | gpu_count = 1 |
| 908 | total_memory = psutil.virtual_memory().available / (1024**3) |
| 909 | |
| 910 | avg_gpu_memory = total_memory / gpu_count |
| 911 | |
| 912 | # rough estimate of batch size |
| 913 | if batch_size_type == "frame": |
| 914 | batch_size_per_gpu = max(int(38400 * (avg_gpu_memory - 5) / 75), int(max_sample_length)) |
| 915 | elif batch_size_type == "sample": |
| 916 | batch_size_per_gpu = int(200 / (total_duration / total_samples)) |
| 917 | |
| 918 | if total_samples < 64: |
| 919 | max_samples = int(total_samples * 0.25) |
| 920 | |
| 921 | num_warmup_updates = max(num_warmup_updates, int(total_samples * 0.05)) |
| 922 | |
| 923 | # take 1.2M updates as the maximum |
| 924 | max_updates = 1200000 |
| 925 | |
| 926 | if batch_size_type == "frame": |
| 927 | mini_batch_duration = batch_size_per_gpu * gpu_count * hop_length / sampling_rate |
| 928 | updates_per_epoch = total_duration / mini_batch_duration |
| 929 | elif batch_size_type == "sample": |
| 930 | updates_per_epoch = total_samples / batch_size_per_gpu / gpu_count |
| 931 | |
| 932 | epochs = int(max_updates / updates_per_epoch) |
| 933 | |
| 934 | if finetune: |
| 935 | learning_rate = 1e-5 |
| 936 | else: |
| 937 | learning_rate = 7.5e-5 |
| 938 | |
| 939 | return ( |
| 940 | epochs, |
| 941 | learning_rate, |
| 942 | batch_size_per_gpu, |
| 943 | max_samples, |
| 944 | num_warmup_updates, |
| 945 | total_samples, |
| 946 | ) |
| 947 | |
| 948 | |
| 949 | def prune_checkpoint(checkpoint_path: str, new_checkpoint_path: str, save_ema: bool, safetensors: bool) -> str: |
| 950 | try: |
| 951 | checkpoint = torch.load(checkpoint_path, weights_only=True) |
| 952 | print("Original Checkpoint Keys:", checkpoint.keys()) |
| 953 | |
| 954 | to_retain = "ema_model_state_dict" if save_ema else "model_state_dict" |
| 955 | try: |
| 956 | model_state_dict_to_retain = checkpoint[to_retain] |
| 957 | except KeyError: |
| 958 | return f"{to_retain} not found in the checkpoint." |
| 959 | |
| 960 | if safetensors: |
| 961 | new_checkpoint_path = new_checkpoint_path.replace(".pt", ".safetensors") |
| 962 | save_file(model_state_dict_to_retain, new_checkpoint_path) |
| 963 | else: |
| 964 | new_checkpoint_path = new_checkpoint_path.replace(".safetensors", ".pt") |
| 965 | new_checkpoint = {"ema_model_state_dict": model_state_dict_to_retain} |
| 966 | torch.save(new_checkpoint, new_checkpoint_path) |
| 967 | |
| 968 | return f"New checkpoint saved at: {new_checkpoint_path}" |
| 969 | |
| 970 | except Exception as e: |
| 971 | return f"An error occurred: {e}" |
| 972 | |
| 973 | |
| 974 | def expand_model_embeddings(ckpt_path, new_ckpt_path, num_new_tokens=42): |
| 975 | seed = 666 |
| 976 | random.seed(seed) |
| 977 | os.environ["PYTHONHASHSEED"] = str(seed) |
| 978 | torch.manual_seed(seed) |
| 979 | torch.cuda.manual_seed(seed) |
| 980 | torch.cuda.manual_seed_all(seed) |
| 981 | torch.backends.cudnn.deterministic = True |
| 982 | torch.backends.cudnn.benchmark = False |
| 983 | |
| 984 | if ckpt_path.endswith(".safetensors"): |
| 985 | ckpt = load_file(ckpt_path, device="cpu") |
| 986 | ckpt = {"ema_model_state_dict": ckpt} |
| 987 | elif ckpt_path.endswith(".pt"): |
| 988 | ckpt = torch.load(ckpt_path, map_location="cpu") |
| 989 | |
| 990 | ema_sd = ckpt.get("ema_model_state_dict", {}) |
| 991 | embed_key_ema = "ema_model.transformer.text_embed.text_embed.weight" |
| 992 | old_embed_ema = ema_sd[embed_key_ema] |
| 993 | |
| 994 | vocab_old = old_embed_ema.size(0) |
| 995 | embed_dim = old_embed_ema.size(1) |
| 996 | vocab_new = vocab_old + num_new_tokens |
| 997 | |
| 998 | def expand_embeddings(old_embeddings): |
| 999 | new_embeddings = torch.zeros((vocab_new, embed_dim)) |
| 1000 | new_embeddings[:vocab_old] = old_embeddings |
| 1001 | new_embeddings[vocab_old:] = torch.randn((num_new_tokens, embed_dim)) |
| 1002 | return new_embeddings |
| 1003 | |
| 1004 | ema_sd[embed_key_ema] = expand_embeddings(ema_sd[embed_key_ema]) |
| 1005 | |
| 1006 | if new_ckpt_path.endswith(".safetensors"): |
| 1007 | save_file(ema_sd, new_ckpt_path) |
| 1008 | elif new_ckpt_path.endswith(".pt"): |
| 1009 | torch.save(ckpt, new_ckpt_path) |
| 1010 | |
| 1011 | return vocab_new |
| 1012 | |
| 1013 | |
| 1014 | def vocab_count(text): |
| 1015 | return str(len(text.split(","))) |
| 1016 | |
| 1017 | |
| 1018 | def vocab_extend(project_name, symbols, model_type): |
| 1019 | if symbols == "": |
| 1020 | return "Symbols empty!" |
| 1021 | |
| 1022 | name_project = project_name |
| 1023 | path_project = _safe_project_path(path_data, name_project) |
| 1024 | file_vocab_project = os.path.join(path_project, "vocab.txt") |
| 1025 | |
| 1026 | file_vocab = os.path.join(path_data, "Emilia_ZH_EN_pinyin/vocab.txt") |
| 1027 | if not os.path.isfile(file_vocab): |
| 1028 | return f"the file {file_vocab} not found !" |
| 1029 | |
| 1030 | symbols = symbols.split(",") |
| 1031 | if symbols == []: |
| 1032 | return "Symbols to extend not found." |
| 1033 | |
| 1034 | with open(file_vocab, "r", encoding="utf-8-sig") as f: |
| 1035 | data = f.read() |
| 1036 | vocab = data.split("\n") |
| 1037 | vocab_check = set(vocab) |
| 1038 | |
| 1039 | miss_symbols = [] |
| 1040 | for item in symbols: |
| 1041 | item = item.replace(" ", "") |
| 1042 | if item in vocab_check: |
| 1043 | continue |
| 1044 | miss_symbols.append(item) |
| 1045 | |
| 1046 | if miss_symbols == []: |
| 1047 | return "Symbols are okay no need to extend." |
| 1048 | |
| 1049 | size_vocab = len(vocab) |
| 1050 | vocab.pop() |
| 1051 | for item in miss_symbols: |
| 1052 | vocab.append(item) |
| 1053 | |
| 1054 | vocab.append("") |
| 1055 | |
| 1056 | with open(file_vocab_project, "w", encoding="utf-8") as f: |
| 1057 | f.write("\n".join(vocab)) |
| 1058 | |
| 1059 | if model_type == "F5TTS_v1_Base": |
| 1060 | ckpt_path = str(cached_path("hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors")) |
| 1061 | elif model_type == "F5TTS_Base": |
| 1062 | ckpt_path = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt")) |
| 1063 | elif model_type == "E2TTS_Base": |
| 1064 | ckpt_path = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.pt")) |
| 1065 | |
| 1066 | vocab_size_new = len(miss_symbols) |
| 1067 | |
| 1068 | dataset_name = name_project.replace("_pinyin", "").replace("_char", "") |
| 1069 | new_ckpt_path = _safe_project_path(path_project_ckpts, dataset_name) |
| 1070 | os.makedirs(new_ckpt_path, exist_ok=True) |
| 1071 | |
| 1072 | # Add pretrained_ prefix to model when copying for consistency with finetune_cli.py |
| 1073 | new_ckpt_file = os.path.join(new_ckpt_path, "pretrained_" + os.path.basename(ckpt_path)) |
| 1074 | |
| 1075 | size = expand_model_embeddings(ckpt_path, new_ckpt_file, num_new_tokens=vocab_size_new) |
| 1076 | |
| 1077 | vocab_new = "\n".join(miss_symbols) |
| 1078 | return f"vocab old size : {size_vocab}\nvocab new size : {size}\nvocab add : {vocab_size_new}\nnew symbols :\n{vocab_new}" |
| 1079 | |
| 1080 | |
| 1081 | def vocab_check(project_name, tokenizer_type): |
| 1082 | name_project = project_name |
| 1083 | path_project = _safe_project_path(path_data, name_project) |
| 1084 | |
| 1085 | file_metadata = os.path.join(path_project, "metadata.csv") |
| 1086 | |
| 1087 | file_vocab = os.path.join(path_data, "Emilia_ZH_EN_pinyin/vocab.txt") |
| 1088 | if not os.path.isfile(file_vocab): |
| 1089 | return f"the file {file_vocab} not found !", "" |
| 1090 | |
| 1091 | with open(file_vocab, "r", encoding="utf-8-sig") as f: |
| 1092 | data = f.read() |
| 1093 | vocab = data.split("\n") |
| 1094 | vocab = set(vocab) |
| 1095 | |
| 1096 | if not os.path.isfile(file_metadata): |
| 1097 | return f"the file {file_metadata} not found !", "" |
| 1098 | |
| 1099 | with open(file_metadata, "r", encoding="utf-8-sig") as f: |
| 1100 | data = f.read() |
| 1101 | |
| 1102 | miss_symbols = [] |
| 1103 | miss_symbols_keep = {} |
| 1104 | for item in data.split("\n"): |
| 1105 | sp = item.split("|") |
| 1106 | if len(sp) != 2: |
| 1107 | continue |
| 1108 | |
| 1109 | text = sp[1].strip() |
| 1110 | if tokenizer_type == "pinyin": |
| 1111 | text = convert_char_to_pinyin([text], polyphone=True)[0] |
| 1112 | |
| 1113 | for t in text: |
| 1114 | if t not in vocab and t not in miss_symbols_keep: |
| 1115 | miss_symbols.append(t) |
| 1116 | miss_symbols_keep[t] = t |
| 1117 | |
| 1118 | if miss_symbols == []: |
| 1119 | vocab_miss = "" |
| 1120 | info = "You can train using your language !" |
| 1121 | else: |
| 1122 | vocab_miss = ",".join(miss_symbols) |
| 1123 | info = f"The following {len(miss_symbols)} symbols are missing in your language\n\n" |
| 1124 | |
| 1125 | return info, vocab_miss |
| 1126 | |
| 1127 | |
| 1128 | def get_random_sample_prepare(project_name): |
| 1129 | name_project = project_name |
| 1130 | path_project = _safe_project_path(path_data, name_project) |
| 1131 | file_arrow = os.path.join(path_project, "raw.arrow") |
| 1132 | if not os.path.isfile(file_arrow): |
| 1133 | return "", None |
| 1134 | dataset = Dataset_.from_file(file_arrow) |
| 1135 | random_sample = dataset.shuffle(seed=random.randint(0, 1000)).select([0]) |
| 1136 | text = "[" + " , ".join(["' " + t + " '" for t in random_sample["text"][0]]) + "]" |
| 1137 | audio_path = random_sample["audio_path"][0] |
| 1138 | return text, audio_path |
| 1139 | |
| 1140 | |
| 1141 | def get_random_sample_transcribe(project_name): |
| 1142 | name_project = project_name |
| 1143 | path_project = _safe_project_path(path_data, name_project) |
| 1144 | file_metadata = os.path.join(path_project, "metadata.csv") |
| 1145 | if not os.path.isfile(file_metadata): |
| 1146 | return "", None |
| 1147 | |
| 1148 | data = "" |
| 1149 | with open(file_metadata, "r", encoding="utf-8-sig") as f: |
| 1150 | data = f.read() |
| 1151 | |
| 1152 | list_data = [] |
| 1153 | for item in data.split("\n"): |
| 1154 | sp = item.split("|") |
| 1155 | if len(sp) != 2: |
| 1156 | continue |
| 1157 | |
| 1158 | # fixed audio when it is absolute |
| 1159 | file_audio = get_correct_audio_path(sp[0], os.path.join(path_project, "wavs")) |
| 1160 | list_data.append([file_audio, sp[1]]) |
| 1161 | |
| 1162 | if list_data == []: |
| 1163 | return "", None |
| 1164 | |
| 1165 | random_item = random.choice(list_data) |
| 1166 | |
| 1167 | return random_item[1], random_item[0] |
| 1168 | |
| 1169 | |
| 1170 | def get_random_sample_infer(project_name): |
| 1171 | text, audio = get_random_sample_transcribe(project_name) |
| 1172 | return ( |
| 1173 | text, |
| 1174 | text, |
| 1175 | audio, |
| 1176 | ) |
| 1177 | |
| 1178 | |
| 1179 | def infer( |
| 1180 | project, file_checkpoint, exp_name, ref_text, ref_audio, gen_text, nfe_step, use_ema, speed, seed, remove_silence |
| 1181 | ): |
| 1182 | global last_checkpoint, last_device, tts_api, last_ema |
| 1183 | |
| 1184 | if not os.path.isfile(file_checkpoint): |
| 1185 | return None, "checkpoint not found!" |
| 1186 | |
| 1187 | if training_process is not None: |
| 1188 | device_test = "cpu" |
| 1189 | else: |
| 1190 | device_test = None |
| 1191 | |
| 1192 | if last_checkpoint != file_checkpoint or last_device != device_test or last_ema != use_ema or tts_api is None: |
| 1193 | if last_checkpoint != file_checkpoint: |
| 1194 | last_checkpoint = file_checkpoint |
| 1195 | |
| 1196 | if last_device != device_test: |
| 1197 | last_device = device_test |
| 1198 | |
| 1199 | if last_ema != use_ema: |
| 1200 | last_ema = use_ema |
| 1201 | |
| 1202 | vocab_file = os.path.join(path_data, project, "vocab.txt") |
| 1203 | |
| 1204 | tts_api = F5TTS( |
| 1205 | model=exp_name, ckpt_file=file_checkpoint, vocab_file=vocab_file, device=device_test, use_ema=use_ema |
| 1206 | ) |
| 1207 | |
| 1208 | print("update >> ", device_test, file_checkpoint, use_ema) |
| 1209 | |
| 1210 | if seed == -1: # -1 used for random |
| 1211 | seed = None |
| 1212 | |
| 1213 | with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f: |
| 1214 | tts_api.infer( |
| 1215 | ref_file=ref_audio, |
| 1216 | ref_text=ref_text.strip(), |
| 1217 | gen_text=gen_text.strip(), |
| 1218 | nfe_step=nfe_step, |
| 1219 | speed=speed, |
| 1220 | remove_silence=remove_silence, |
| 1221 | file_wave=f.name, |
| 1222 | seed=seed, |
| 1223 | ) |
| 1224 | return f.name, tts_api.device, str(tts_api.seed) |
| 1225 | |
| 1226 | |
| 1227 | def check_finetune(finetune): |
| 1228 | return gr.update(interactive=finetune), gr.update(interactive=finetune), gr.update(interactive=finetune) |
| 1229 | |
| 1230 | |
| 1231 | def get_checkpoints_project(project_name, is_gradio=True): |
| 1232 | if project_name is None: |
| 1233 | return [], "" |
| 1234 | project_name = project_name.replace("_pinyin", "").replace("_char", "") |
| 1235 | |
| 1236 | if os.path.isdir(path_project_ckpts): |
| 1237 | files_checkpoints = glob(os.path.join(path_project_ckpts, project_name, "*.pt")) |
| 1238 | # Separate pretrained and regular checkpoints |
| 1239 | pretrained_checkpoints = [f for f in files_checkpoints if "pretrained_" in os.path.basename(f)] |
| 1240 | regular_checkpoints = [ |
| 1241 | f |
| 1242 | for f in files_checkpoints |
| 1243 | if "pretrained_" not in os.path.basename(f) and "model_last.pt" not in os.path.basename(f) |
| 1244 | ] |
| 1245 | last_checkpoint = [f for f in files_checkpoints if "model_last.pt" in os.path.basename(f)] |
| 1246 | |
| 1247 | # Sort regular checkpoints by number |
| 1248 | regular_checkpoints = sorted( |
| 1249 | regular_checkpoints, key=lambda x: int(os.path.basename(x).split("_")[1].split(".")[0]) |
| 1250 | ) |
| 1251 | |
| 1252 | # Combine in order: pretrained, regular, last |
| 1253 | files_checkpoints = pretrained_checkpoints + regular_checkpoints + last_checkpoint |
| 1254 | else: |
| 1255 | files_checkpoints = [] |
| 1256 | |
| 1257 | selelect_checkpoint = None if not files_checkpoints else files_checkpoints[0] |
| 1258 | |
| 1259 | if is_gradio: |
| 1260 | return gr.update(choices=files_checkpoints, value=selelect_checkpoint) |
| 1261 | |
| 1262 | return files_checkpoints, selelect_checkpoint |
| 1263 | |
| 1264 | |
| 1265 | def get_audio_project(project_name, is_gradio=True): |
| 1266 | if project_name is None: |
| 1267 | return [], "" |
| 1268 | project_name = project_name.replace("_pinyin", "").replace("_char", "") |
| 1269 | |
| 1270 | if os.path.isdir(path_project_ckpts): |
| 1271 | files_audios = glob(os.path.join(path_project_ckpts, project_name, "samples", "*.wav")) |
| 1272 | files_audios = sorted(files_audios, key=lambda x: int(os.path.basename(x).split("_")[1].split(".")[0])) |
| 1273 | |
| 1274 | files_audios = [item.replace("_gen.wav", "") for item in files_audios if item.endswith("_gen.wav")] |
| 1275 | else: |
| 1276 | files_audios = [] |
| 1277 | |
| 1278 | selelect_checkpoint = None if not files_audios else files_audios[0] |
| 1279 | |
| 1280 | if is_gradio: |
| 1281 | return gr.update(choices=files_audios, value=selelect_checkpoint) |
| 1282 | |
| 1283 | return files_audios, selelect_checkpoint |
| 1284 | |
| 1285 | |
| 1286 | def get_gpu_stats(): |
| 1287 | gpu_stats = "" |
| 1288 | |
| 1289 | if torch.cuda.is_available(): |
| 1290 | gpu_count = torch.cuda.device_count() |
| 1291 | for i in range(gpu_count): |
| 1292 | gpu_name = torch.cuda.get_device_name(i) |
| 1293 | gpu_properties = torch.cuda.get_device_properties(i) |
| 1294 | total_memory = gpu_properties.total_memory / (1024**3) # in GB |
| 1295 | allocated_memory = torch.cuda.memory_allocated(i) / (1024**2) # in MB |
| 1296 | reserved_memory = torch.cuda.memory_reserved(i) / (1024**2) # in MB |
| 1297 | |
| 1298 | gpu_stats += ( |
| 1299 | f"GPU {i} Name: {gpu_name}\n" |
| 1300 | f"Total GPU memory (GPU {i}): {total_memory:.2f} GB\n" |
| 1301 | f"Allocated GPU memory (GPU {i}): {allocated_memory:.2f} MB\n" |
| 1302 | f"Reserved GPU memory (GPU {i}): {reserved_memory:.2f} MB\n\n" |
| 1303 | ) |
| 1304 | elif torch.xpu.is_available(): |
| 1305 | gpu_count = torch.xpu.device_count() |
| 1306 | for i in range(gpu_count): |
| 1307 | gpu_name = torch.xpu.get_device_name(i) |
| 1308 | gpu_properties = torch.xpu.get_device_properties(i) |
| 1309 | total_memory = gpu_properties.total_memory / (1024**3) # in GB |
| 1310 | allocated_memory = torch.xpu.memory_allocated(i) / (1024**2) # in MB |
| 1311 | reserved_memory = torch.xpu.memory_reserved(i) / (1024**2) # in MB |
| 1312 | |
| 1313 | gpu_stats += ( |
| 1314 | f"GPU {i} Name: {gpu_name}\n" |
| 1315 | f"Total GPU memory (GPU {i}): {total_memory:.2f} GB\n" |
| 1316 | f"Allocated GPU memory (GPU {i}): {allocated_memory:.2f} MB\n" |
| 1317 | f"Reserved GPU memory (GPU {i}): {reserved_memory:.2f} MB\n\n" |
| 1318 | ) |
| 1319 | elif torch.backends.mps.is_available(): |
| 1320 | gpu_count = 1 |
| 1321 | gpu_stats += "MPS GPU\n" |
| 1322 | total_memory = psutil.virtual_memory().total / ( |
| 1323 | 1024**3 |
| 1324 | ) # Total system memory (MPS doesn't have its own memory) |
| 1325 | allocated_memory = 0 |
| 1326 | reserved_memory = 0 |
| 1327 | |
| 1328 | gpu_stats += ( |
| 1329 | f"Total system memory: {total_memory:.2f} GB\n" |
| 1330 | f"Allocated GPU memory (MPS): {allocated_memory:.2f} MB\n" |
| 1331 | f"Reserved GPU memory (MPS): {reserved_memory:.2f} MB\n" |
| 1332 | ) |
| 1333 | |
| 1334 | else: |
| 1335 | gpu_stats = "No GPU available" |
| 1336 | |
| 1337 | return gpu_stats |
| 1338 | |
| 1339 | |
| 1340 | def get_cpu_stats(): |
| 1341 | cpu_usage = psutil.cpu_percent(interval=1) |
| 1342 | memory_info = psutil.virtual_memory() |
| 1343 | memory_used = memory_info.used / (1024**2) |
| 1344 | memory_total = memory_info.total / (1024**2) |
| 1345 | memory_percent = memory_info.percent |
| 1346 | |
| 1347 | pid = os.getpid() |
| 1348 | process = psutil.Process(pid) |
| 1349 | nice_value = process.nice() |
| 1350 | |
| 1351 | cpu_stats = ( |
| 1352 | f"CPU Usage: {cpu_usage:.2f}%\n" |
| 1353 | f"System Memory: {memory_used:.2f} MB used / {memory_total:.2f} MB total ({memory_percent}% used)\n" |
| 1354 | f"Process Priority (Nice value): {nice_value}" |
| 1355 | ) |
| 1356 | |
| 1357 | return cpu_stats |
| 1358 | |
| 1359 | |
| 1360 | def get_combined_stats(): |
| 1361 | gpu_stats = get_gpu_stats() |
| 1362 | cpu_stats = get_cpu_stats() |
| 1363 | combined_stats = f"### GPU Stats\n{gpu_stats}\n\n### CPU Stats\n{cpu_stats}" |
| 1364 | return combined_stats |
| 1365 | |
| 1366 | |
| 1367 | def get_audio_select(file_sample): |
| 1368 | select_audio_ref = file_sample |
| 1369 | select_audio_gen = file_sample |
| 1370 | |
| 1371 | if file_sample is not None: |
| 1372 | select_audio_ref += "_ref.wav" |
| 1373 | select_audio_gen += "_gen.wav" |
| 1374 | |
| 1375 | return select_audio_ref, select_audio_gen |
| 1376 | |
| 1377 | |
| 1378 | with gr.Blocks() as app: |
| 1379 | gr.Markdown( |
| 1380 | """ |
| 1381 | # F5 TTS Automatic Finetune |
| 1382 | |
| 1383 | This is a local web UI for F5 TTS finetuning support. This app supports the following TTS models: |
| 1384 | |
| 1385 | * [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching) |
| 1386 | * [E2 TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS) |
| 1387 | |
| 1388 | The pretrained checkpoints support English and Chinese. |
| 1389 | |
| 1390 | For tutorial and updates check here (https://github.com/SWivid/F5-TTS/discussions/143) |
| 1391 | """ |
| 1392 | ) |
| 1393 | |
| 1394 | with gr.Row(): |
| 1395 | projects, projects_selelect = get_list_projects() |
| 1396 | tokenizer_type = gr.Radio(label="Tokenizer Type", choices=["pinyin", "char", "custom"], value="pinyin") |
| 1397 | project_name = gr.Textbox(label="Project Name", value="my_speak") |
| 1398 | bt_create = gr.Button("Create a New Project") |
| 1399 | |
| 1400 | with gr.Row(): |
| 1401 | cm_project = gr.Dropdown( |
| 1402 | choices=projects, value=projects_selelect, label="Project", allow_custom_value=True, scale=6 |
| 1403 | ) |
| 1404 | ch_refresh_project = gr.Button("Refresh", scale=1) |
| 1405 | |
| 1406 | bt_create.click(fn=create_data_project, inputs=[project_name, tokenizer_type], outputs=[cm_project]) |
| 1407 | |
| 1408 | with gr.Tabs(): |
| 1409 | with gr.TabItem("Transcribe Data"): |
| 1410 | gr.Markdown("""```plaintext |
| 1411 | Skip this step if you have your dataset, metadata.csv, and a folder wavs with all the audio files. |
| 1412 | ```""") |
| 1413 | |
| 1414 | ch_manual = gr.Checkbox(label="Audio from Path", value=False) |
| 1415 | |
| 1416 | mark_info_transcribe = gr.Markdown( |
| 1417 | """```plaintext |
| 1418 | Place your 'wavs' folder and 'metadata.csv' file in the '{your_project_name}' directory. |
| 1419 | |
| 1420 | my_speak/ |
| 1421 | │ |
| 1422 | └── dataset/ |
| 1423 | ├── audio1.wav |
| 1424 | └── audio2.wav |
| 1425 | ... |
| 1426 | ```""", |
| 1427 | visible=False, |
| 1428 | ) |
| 1429 | |
| 1430 | audio_speaker = gr.File(label="Voice", type="filepath", file_count="multiple") |
| 1431 | txt_lang = gr.Textbox(label="Language", value="English") |
| 1432 | bt_transcribe = bt_create = gr.Button("Transcribe") |
| 1433 | txt_info_transcribe = gr.Textbox(label="Info", value="") |
| 1434 | bt_transcribe.click( |
| 1435 | fn=transcribe_all, |
| 1436 | inputs=[cm_project, audio_speaker, txt_lang, ch_manual], |
| 1437 | outputs=[txt_info_transcribe], |
| 1438 | ) |
| 1439 | ch_manual.change(fn=check_user, inputs=[ch_manual], outputs=[audio_speaker, mark_info_transcribe]) |
| 1440 | |
| 1441 | random_sample_transcribe = gr.Button("Random Sample") |
| 1442 | |
| 1443 | with gr.Row(): |
| 1444 | random_text_transcribe = gr.Textbox(label="Text") |
| 1445 | random_audio_transcribe = gr.Audio(label="Audio", type="filepath") |
| 1446 | |
| 1447 | random_sample_transcribe.click( |
| 1448 | fn=get_random_sample_transcribe, |
| 1449 | inputs=[cm_project], |
| 1450 | outputs=[random_text_transcribe, random_audio_transcribe], |
| 1451 | ) |
| 1452 | |
| 1453 | with gr.TabItem("Vocab Check"): |
| 1454 | gr.Markdown("""```plaintext |
| 1455 | Check the vocabulary for fine-tuning Emilia_ZH_EN to ensure all symbols are included. For fine-tuning a new language. |
| 1456 | ```""") |
| 1457 | |
| 1458 | check_button = gr.Button("Check Vocab") |
| 1459 | txt_info_check = gr.Textbox(label="Info", value="") |
| 1460 | |
| 1461 | gr.Markdown("""```plaintext |
| 1462 | Using the extended model, you can finetune to a new language that is missing symbols in the vocab. This creates a new model with a new vocabulary size and saves it in your ckpts/project folder. |
| 1463 | ```""") |
| 1464 | |
| 1465 | exp_name_extend = gr.Radio( |
| 1466 | label="Model", choices=["F5TTS_v1_Base", "F5TTS_Base", "E2TTS_Base"], value="F5TTS_v1_Base" |
| 1467 | ) |
| 1468 | |
| 1469 | with gr.Row(): |
| 1470 | txt_extend = gr.Textbox( |
| 1471 | label="Symbols", |
| 1472 | value="", |
| 1473 | placeholder="To add new symbols, make sure to use ',' for each symbol", |
| 1474 | scale=6, |
| 1475 | ) |
| 1476 | txt_count_symbol = gr.Textbox(label="New Vocab Size", value="", scale=1) |
| 1477 | |
| 1478 | extend_button = gr.Button("Extend") |
| 1479 | txt_info_extend = gr.Textbox(label="Info", value="") |
| 1480 | |
| 1481 | txt_extend.change(vocab_count, inputs=[txt_extend], outputs=[txt_count_symbol]) |
| 1482 | check_button.click( |
| 1483 | fn=vocab_check, inputs=[cm_project, tokenizer_type], outputs=[txt_info_check, txt_extend] |
| 1484 | ) |
| 1485 | extend_button.click( |
| 1486 | fn=vocab_extend, inputs=[cm_project, txt_extend, exp_name_extend], outputs=[txt_info_extend] |
| 1487 | ) |
| 1488 | |
| 1489 | with gr.TabItem("Prepare Data"): |
| 1490 | gr.Markdown("""```plaintext |
| 1491 | Skip this step if you have your dataset, raw.arrow, duration.json, and vocab.txt |
| 1492 | ```""") |
| 1493 | |
| 1494 | gr.Markdown( |
| 1495 | """```plaintext |
| 1496 | Place all your "wavs" folder and your "metadata.csv" file in your project name directory. |
| 1497 | |
| 1498 | Supported audio formats: "wav", "mp3", "aac", "flac", "m4a", "alac", "ogg", "aiff", "wma", "amr" |
| 1499 | |
| 1500 | Example wav format: |
| 1501 | my_speak/ |
| 1502 | │ |
| 1503 | ├── wavs/ |
| 1504 | │ ├── audio1.wav |
| 1505 | │ └── audio2.wav |
| 1506 | | ... |
| 1507 | │ |
| 1508 | └── metadata.csv |
| 1509 | |
| 1510 | File format metadata.csv: |
| 1511 | |
| 1512 | audio1|text1 or audio1.wav|text1 or your_path/audio1.wav|text1 |
| 1513 | audio2|text1 or audio2.wav|text1 or your_path/audio2.wav|text1 |
| 1514 | ... |
| 1515 | |
| 1516 | ```""" |
| 1517 | ) |
| 1518 | ch_tokenizern = gr.Checkbox(label="Create Vocabulary", value=False, visible=False) |
| 1519 | |
| 1520 | bt_prepare = bt_create = gr.Button("Prepare") |
| 1521 | txt_info_prepare = gr.Textbox(label="Info", value="") |
| 1522 | txt_vocab_prepare = gr.Textbox(label="Vocab", value="") |
| 1523 | |
| 1524 | bt_prepare.click( |
| 1525 | fn=create_metadata, inputs=[cm_project, ch_tokenizern], outputs=[txt_info_prepare, txt_vocab_prepare] |
| 1526 | ) |
| 1527 | |
| 1528 | random_sample_prepare = gr.Button("Random Sample") |
| 1529 | |
| 1530 | with gr.Row(): |
| 1531 | random_text_prepare = gr.Textbox(label="Tokenizer") |
| 1532 | random_audio_prepare = gr.Audio(label="Audio", type="filepath") |
| 1533 | |
| 1534 | random_sample_prepare.click( |
| 1535 | fn=get_random_sample_prepare, inputs=[cm_project], outputs=[random_text_prepare, random_audio_prepare] |
| 1536 | ) |
| 1537 | |
| 1538 | with gr.TabItem("Train Model"): |
| 1539 | gr.Markdown("""```plaintext |
| 1540 | The auto-setting is still experimental. Set a large value of epoch if not sure; and keep last N checkpoints if limited disk space. |
| 1541 | If you encounter a memory error, try reducing the batch size per GPU to a smaller number. |
| 1542 | ```""") |
| 1543 | with gr.Row(): |
| 1544 | exp_name = gr.Radio(label="Model", choices=["F5TTS_v1_Base", "F5TTS_Base", "E2TTS_Base"]) |
| 1545 | tokenizer_file = gr.Textbox(label="Tokenizer File") |
| 1546 | file_checkpoint_train = gr.Textbox(label="Path to the Pretrained Checkpoint") |
| 1547 | |
| 1548 | with gr.Row(): |
| 1549 | ch_finetune = bt_create = gr.Checkbox(label="Finetune") |
| 1550 | lb_samples = gr.Label(label="Samples") |
| 1551 | bt_calculate = bt_create = gr.Button("Auto Settings") |
| 1552 | |
| 1553 | with gr.Row(): |
| 1554 | epochs = gr.Number(label="Epochs") |
| 1555 | learning_rate = gr.Number(label="Learning Rate", step=0.5e-5) |
| 1556 | max_grad_norm = gr.Number(label="Max Gradient Norm") |
| 1557 | num_warmup_updates = gr.Number(label="Warmup Updates") |
| 1558 | |
| 1559 | with gr.Row(): |
| 1560 | batch_size_type = gr.Radio( |
| 1561 | label="Batch Size Type", |
| 1562 | choices=["frame", "sample"], |
| 1563 | info="frame is calculated as seconds * sampling_rate / hop_length", |
| 1564 | ) |
| 1565 | batch_size_per_gpu = gr.Number(label="Batch Size per GPU", info="N frames or N samples") |
| 1566 | grad_accumulation_steps = gr.Number( |
| 1567 | label="Gradient Accumulation Steps", info="Effective batch size is multiplied by this value" |
| 1568 | ) |
| 1569 | max_samples = gr.Number(label="Max Samples", info="Maximum number of samples per single GPU batch") |
| 1570 | |
| 1571 | with gr.Row(): |
| 1572 | save_per_updates = gr.Number( |
| 1573 | label="Save per Updates", |
| 1574 | info="Save intermediate checkpoints every N updates", |
| 1575 | minimum=10, |
| 1576 | ) |
| 1577 | keep_last_n_checkpoints = gr.Number( |
| 1578 | label="Keep Last N Checkpoints", |
| 1579 | step=1, |
| 1580 | precision=0, |
| 1581 | info="-1 to keep all, 0 to not save intermediate, > 0 to keep last N", |
| 1582 | minimum=-1, |
| 1583 | ) |
| 1584 | last_per_updates = gr.Number( |
| 1585 | label="Last per Updates", |
| 1586 | info="Save latest checkpoint with suffix _last.pt every N updates", |
| 1587 | minimum=10, |
| 1588 | ) |
| 1589 | gr.Radio(label="") # placeholder |
| 1590 | |
| 1591 | with gr.Row(): |
| 1592 | ch_8bit_adam = gr.Checkbox(label="Use 8-bit Adam optimizer") |
| 1593 | mixed_precision = gr.Radio(label="Mixed Precision", choices=["none", "fp16", "bf16"]) |
| 1594 | cd_logger = gr.Radio(label="Logger", choices=["none", "wandb", "tensorboard"]) |
| 1595 | with gr.Column(): |
| 1596 | start_button = gr.Button("Start Training") |
| 1597 | stop_button = gr.Button("Stop Training", interactive=False) |
| 1598 | |
| 1599 | if projects_selelect is not None: |
| 1600 | ( |
| 1601 | exp_name_value, |
| 1602 | learning_rate_value, |
| 1603 | batch_size_per_gpu_value, |
| 1604 | batch_size_type_value, |
| 1605 | max_samples_value, |
| 1606 | grad_accumulation_steps_value, |
| 1607 | max_grad_norm_value, |
| 1608 | epochs_value, |
| 1609 | num_warmup_updates_value, |
| 1610 | save_per_updates_value, |
| 1611 | keep_last_n_checkpoints_value, |
| 1612 | last_per_updates_value, |
| 1613 | finetune_value, |
| 1614 | file_checkpoint_train_value, |
| 1615 | tokenizer_type_value, |
| 1616 | tokenizer_file_value, |
| 1617 | mixed_precision_value, |
| 1618 | logger_value, |
| 1619 | bnb_optimizer_value, |
| 1620 | ) = load_settings(projects_selelect) |
| 1621 | |
| 1622 | # Assigning values to the respective components |
| 1623 | exp_name.value = exp_name_value |
| 1624 | learning_rate.value = learning_rate_value |
| 1625 | batch_size_per_gpu.value = batch_size_per_gpu_value |
| 1626 | batch_size_type.value = batch_size_type_value |
| 1627 | max_samples.value = max_samples_value |
| 1628 | grad_accumulation_steps.value = grad_accumulation_steps_value |
| 1629 | max_grad_norm.value = max_grad_norm_value |
| 1630 | epochs.value = epochs_value |
| 1631 | num_warmup_updates.value = num_warmup_updates_value |
| 1632 | save_per_updates.value = save_per_updates_value |
| 1633 | keep_last_n_checkpoints.value = keep_last_n_checkpoints_value |
| 1634 | last_per_updates.value = last_per_updates_value |
| 1635 | ch_finetune.value = finetune_value |
| 1636 | file_checkpoint_train.value = file_checkpoint_train_value |
| 1637 | tokenizer_type.value = tokenizer_type_value |
| 1638 | tokenizer_file.value = tokenizer_file_value |
| 1639 | mixed_precision.value = mixed_precision_value |
| 1640 | cd_logger.value = logger_value |
| 1641 | ch_8bit_adam.value = bnb_optimizer_value |
| 1642 | |
| 1643 | ch_stream = gr.Checkbox(label="Stream Output Experiment", value=True) |
| 1644 | txt_info_train = gr.Textbox(label="Info", value="") |
| 1645 | |
| 1646 | list_audios, select_audio = get_audio_project(projects_selelect, False) |
| 1647 | |
| 1648 | select_audio_ref = select_audio |
| 1649 | select_audio_gen = select_audio |
| 1650 | |
| 1651 | if select_audio is not None: |
| 1652 | select_audio_ref += "_ref.wav" |
| 1653 | select_audio_gen += "_gen.wav" |
| 1654 | |
| 1655 | with gr.Row(): |
| 1656 | ch_list_audio = gr.Dropdown( |
| 1657 | choices=list_audios, |
| 1658 | value=select_audio, |
| 1659 | label="Audios", |
| 1660 | allow_custom_value=True, |
| 1661 | scale=6, |
| 1662 | interactive=True, |
| 1663 | ) |
| 1664 | bt_stream_audio = gr.Button("Refresh", scale=1) |
| 1665 | bt_stream_audio.click(fn=get_audio_project, inputs=[cm_project], outputs=[ch_list_audio]) |
| 1666 | cm_project.change(fn=get_audio_project, inputs=[cm_project], outputs=[ch_list_audio]) |
| 1667 | |
| 1668 | with gr.Row(): |
| 1669 | audio_ref_stream = gr.Audio(label="Original", type="filepath", value=select_audio_ref) |
| 1670 | audio_gen_stream = gr.Audio(label="Generate", type="filepath", value=select_audio_gen) |
| 1671 | |
| 1672 | ch_list_audio.change( |
| 1673 | fn=get_audio_select, |
| 1674 | inputs=[ch_list_audio], |
| 1675 | outputs=[audio_ref_stream, audio_gen_stream], |
| 1676 | ) |
| 1677 | |
| 1678 | start_button.click( |
| 1679 | fn=start_training, |
| 1680 | inputs=[ |
| 1681 | cm_project, |
| 1682 | exp_name, |
| 1683 | learning_rate, |
| 1684 | batch_size_per_gpu, |
| 1685 | batch_size_type, |
| 1686 | max_samples, |
| 1687 | grad_accumulation_steps, |
| 1688 | max_grad_norm, |
| 1689 | epochs, |
| 1690 | num_warmup_updates, |
| 1691 | save_per_updates, |
| 1692 | keep_last_n_checkpoints, |
| 1693 | last_per_updates, |
| 1694 | ch_finetune, |
| 1695 | file_checkpoint_train, |
| 1696 | tokenizer_type, |
| 1697 | tokenizer_file, |
| 1698 | mixed_precision, |
| 1699 | ch_stream, |
| 1700 | cd_logger, |
| 1701 | ch_8bit_adam, |
| 1702 | ], |
| 1703 | outputs=[txt_info_train, start_button, stop_button], |
| 1704 | ) |
| 1705 | stop_button.click(fn=stop_training, outputs=[txt_info_train, start_button, stop_button]) |
| 1706 | |
| 1707 | bt_calculate.click( |
| 1708 | fn=calculate_train, |
| 1709 | inputs=[ |
| 1710 | cm_project, |
| 1711 | epochs, |
| 1712 | learning_rate, |
| 1713 | batch_size_per_gpu, |
| 1714 | batch_size_type, |
| 1715 | max_samples, |
| 1716 | num_warmup_updates, |
| 1717 | ch_finetune, |
| 1718 | ], |
| 1719 | outputs=[ |
| 1720 | epochs, |
| 1721 | learning_rate, |
| 1722 | batch_size_per_gpu, |
| 1723 | max_samples, |
| 1724 | num_warmup_updates, |
| 1725 | lb_samples, |
| 1726 | ], |
| 1727 | ) |
| 1728 | |
| 1729 | ch_finetune.change( |
| 1730 | check_finetune, inputs=[ch_finetune], outputs=[file_checkpoint_train, tokenizer_file, tokenizer_type] |
| 1731 | ) |
| 1732 | |
| 1733 | def setup_load_settings(): |
| 1734 | output_components = [ |
| 1735 | exp_name, |
| 1736 | learning_rate, |
| 1737 | batch_size_per_gpu, |
| 1738 | batch_size_type, |
| 1739 | max_samples, |
| 1740 | grad_accumulation_steps, |
| 1741 | max_grad_norm, |
| 1742 | epochs, |
| 1743 | num_warmup_updates, |
| 1744 | save_per_updates, |
| 1745 | keep_last_n_checkpoints, |
| 1746 | last_per_updates, |
| 1747 | ch_finetune, |
| 1748 | file_checkpoint_train, |
| 1749 | tokenizer_type, |
| 1750 | tokenizer_file, |
| 1751 | mixed_precision, |
| 1752 | cd_logger, |
| 1753 | ch_8bit_adam, |
| 1754 | ] |
| 1755 | return output_components |
| 1756 | |
| 1757 | outputs = setup_load_settings() |
| 1758 | |
| 1759 | cm_project.change( |
| 1760 | fn=load_settings, |
| 1761 | inputs=[cm_project], |
| 1762 | outputs=outputs, |
| 1763 | ) |
| 1764 | |
| 1765 | ch_refresh_project.click( |
| 1766 | fn=load_settings, |
| 1767 | inputs=[cm_project], |
| 1768 | outputs=outputs, |
| 1769 | ) |
| 1770 | |
| 1771 | with gr.TabItem("Test Model"): |
| 1772 | gr.Markdown("""```plaintext |
| 1773 | Check the use_ema setting (True or False) for your model to see what works best for you. Set seed to -1 for random. |
| 1774 | ```""") |
| 1775 | exp_name = gr.Radio( |
| 1776 | label="Model", choices=["F5TTS_v1_Base", "F5TTS_Base", "E2TTS_Base"], value="F5TTS_v1_Base" |
| 1777 | ) |
| 1778 | list_checkpoints, checkpoint_select = get_checkpoints_project(projects_selelect, False) |
| 1779 | |
| 1780 | with gr.Row(): |
| 1781 | nfe_step = gr.Number(label="NFE Step", value=32) |
| 1782 | speed = gr.Slider(label="Speed", value=1.0, minimum=0.3, maximum=2.0, step=0.1) |
| 1783 | seed = gr.Number(label="Random Seed", value=-1, minimum=-1) |
| 1784 | remove_silence = gr.Checkbox(label="Remove Silence") |
| 1785 | |
| 1786 | with gr.Row(): |
| 1787 | ch_use_ema = gr.Checkbox( |
| 1788 | label="Use EMA", value=True, info="Turn off at early stage might offer better results" |
| 1789 | ) |
| 1790 | cm_checkpoint = gr.Dropdown( |
| 1791 | choices=list_checkpoints, value=checkpoint_select, label="Checkpoints", allow_custom_value=True |
| 1792 | ) |
| 1793 | bt_checkpoint_refresh = gr.Button("Refresh") |
| 1794 | |
| 1795 | random_sample_infer = gr.Button("Random Sample") |
| 1796 | |
| 1797 | ref_text = gr.Textbox(label="Reference Text") |
| 1798 | ref_audio = gr.Audio(label="Reference Audio", type="filepath") |
| 1799 | gen_text = gr.Textbox(label="Text to Generate") |
| 1800 | |
| 1801 | random_sample_infer.click( |
| 1802 | fn=get_random_sample_infer, inputs=[cm_project], outputs=[ref_text, gen_text, ref_audio] |
| 1803 | ) |
| 1804 | |
| 1805 | with gr.Row(): |
| 1806 | txt_info_gpu = gr.Textbox("", label="Inference on Device :") |
| 1807 | seed_info = gr.Textbox(label="Used Random Seed :") |
| 1808 | check_button_infer = gr.Button("Inference") |
| 1809 | |
| 1810 | gen_audio = gr.Audio(label="Generated Audio", type="filepath") |
| 1811 | |
| 1812 | check_button_infer.click( |
| 1813 | fn=infer, |
| 1814 | inputs=[ |
| 1815 | cm_project, |
| 1816 | cm_checkpoint, |
| 1817 | exp_name, |
| 1818 | ref_text, |
| 1819 | ref_audio, |
| 1820 | gen_text, |
| 1821 | nfe_step, |
| 1822 | ch_use_ema, |
| 1823 | speed, |
| 1824 | seed, |
| 1825 | remove_silence, |
| 1826 | ], |
| 1827 | outputs=[gen_audio, txt_info_gpu, seed_info], |
| 1828 | ) |
| 1829 | |
| 1830 | bt_checkpoint_refresh.click(fn=get_checkpoints_project, inputs=[cm_project], outputs=[cm_checkpoint]) |
| 1831 | cm_project.change(fn=get_checkpoints_project, inputs=[cm_project], outputs=[cm_checkpoint]) |
| 1832 | |
| 1833 | with gr.TabItem("Prune Checkpoint"): |
| 1834 | gr.Markdown("""```plaintext |
| 1835 | Reduce the Base model size from 5GB to 1.3GB. The new checkpoint file prunes out optimizer and etc., can be used for inference or finetuning afterward, but not able to resume pretraining. |
| 1836 | ```""") |
| 1837 | txt_path_checkpoint = gr.Textbox(label="Path to Checkpoint:") |
| 1838 | txt_path_checkpoint_small = gr.Textbox(label="Path to Output:") |
| 1839 | with gr.Row(): |
| 1840 | ch_save_ema = gr.Checkbox(label="Save EMA checkpoint", value=True) |
| 1841 | ch_safetensors = gr.Checkbox(label="Save with safetensors format", value=True) |
| 1842 | txt_info_reduse = gr.Textbox(label="Info", value="") |
| 1843 | reduse_button = gr.Button("Prune") |
| 1844 | reduse_button.click( |
| 1845 | fn=prune_checkpoint, |
| 1846 | inputs=[txt_path_checkpoint, txt_path_checkpoint_small, ch_save_ema, ch_safetensors], |
| 1847 | outputs=[txt_info_reduse], |
| 1848 | ) |
| 1849 | |
| 1850 | with gr.TabItem("System Info"): |
| 1851 | output_box = gr.Textbox(label="GPU and CPU Information", lines=20) |
| 1852 | |
| 1853 | def update_stats(): |
| 1854 | return get_combined_stats() |
| 1855 | |
| 1856 | update_button = gr.Button("Update Stats") |
| 1857 | update_button.click(fn=update_stats, outputs=output_box) |
| 1858 | |
| 1859 | def auto_update(): |
| 1860 | yield gr.update(value=update_stats()) |
| 1861 | |
| 1862 | gr.update(fn=auto_update, inputs=[], outputs=output_box) |
| 1863 | |
| 1864 | |
| 1865 | @click.command() |
| 1866 | @click.option("--port", "-p", default=None, type=int, help="Port to run the app on") |
| 1867 | @click.option("--host", "-H", default=None, help="Host to run the app on") |
| 1868 | @click.option( |
| 1869 | "--share", |
| 1870 | "-s", |
| 1871 | default=False, |
| 1872 | is_flag=True, |
| 1873 | help="Share the app via Gradio share link", |
| 1874 | ) |
| 1875 | @click.option("--api", "-a", default=True, is_flag=True, help="Allow API access") |
| 1876 | def main(port, host, share, api): |
| 1877 | global app |
| 1878 | print("Starting app...") |
| 1879 | app.queue(api_open=api).launch(server_name=host, server_port=port, share=share, show_api=api) |
| 1880 | |
| 1881 | |
| 1882 | if __name__ == "__main__": |
| 1883 | main() |
| 1884 |