Upload 12 files

README.md

@@ -1,10 +1,3 @@
----
-license: mit
-title: Other guys fine tune xtts web ui idk
-emoji: 🐢
-short_description: Other guys fine tune xtts web ui idk
-sdk: gradio
----
 # xtts-finetune-webui
 
 This webui is a slightly modified copy of the [official webui](https://github.com/coqui-ai/TTS/pull/3296) for finetune xtts.
@@ -58,6 +51,9 @@ If you are looking for an option for normal XTTS use look here [https://github.c
 
 ![image](https://github.com/daswer123/xtts-finetune-webui/assets/22278673/aa05bcd4-8642-4de4-8f2f-bc0f5571af63)
 
+## Google colab
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/DrewThomasson/xtts-finetune-webui/blob/main/notebook/xtts_finetune_webui.ipynb)
+
 ## Install
 
 1. Make sure you have `Cuda` installed
@@ -77,4 +73,9 @@ If you are looking for an option for normal XTTS use look here [https://github.c
 1. Run `bash install.sh`
 2. To start the server start `start.sh`
 3. Go to the local address `127.0.0.1:5003`
-~
+
+### On Apple Silicon Mac (python 3.10 env)
+1. Run `pip install --no-deps -r apple_silicon_requirements.txt`
+2. To start the server `python xtts_demo.py`
+3. Go to the local address `127.0.0.1:5003`
+~                                            

apple_silicon_requirements.txt

@@ -0,0 +1,189 @@
+absl-py==2.1.0
+aiofiles==23.2.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+altair==5.3.0
+annotated-types==0.7.0
+anyascii==0.3.2
+anyio==3.7.1
+async-timeout==4.0.3
+attrs==23.2.0
+audioread==3.0.1
+av==12.2.0
+Babel==2.15.0
+bangla==0.0.2
+blinker==1.8.2
+blis==0.7.11
+bnnumerizer==0.0.2
+bnunicodenormalizer==0.1.7
+catalogue==2.0.10
+certifi==2024.7.4
+cffi==1.16.0
+charset-normalizer==3.3.2
+click==8.1.7
+cloudpathlib==0.16.0
+colorama==0.4.6
+coloredlogs==15.0.1
+confection==0.1.5
+contourpy==1.2.1
+coqpit==0.0.17
+coqui-tts==0.24.2
+coqui-tts-trainer==0.1.4
+ctranslate2==4.3.1
+cutlet==0.4.0
+cycler==0.12.1
+cymem==2.0.8
+Cython==3.0.10
+dateparser==1.1.8
+decorator==5.1.1
+dnspython==2.6.1
+docopt==0.6.2
+einops==0.8.0
+email_validator==2.2.0
+encodec==0.1.1
+exceptiongroup==1.2.2
+fastapi==0.103.1
+fastapi-cli==0.0.4
+faster-whisper==1.0.2
+ffmpy==0.3.2
+filelock==3.15.4
+Flask==3.0.3
+flatbuffers==24.3.25
+fonttools==4.53.1
+frozenlist==1.4.1
+fsspec==2024.6.1
+fugashi==1.3.2
+g2pkk==0.1.2
+gradio==4.44.1
+gradio_client==1.3.0
+grpcio==1.64.1
+gruut==2.4.0
+gruut-ipa==0.13.0
+gruut_lang_de==2.0.1
+gruut_lang_en==2.0.1
+gruut_lang_es==2.0.1
+gruut_lang_fr==2.0.2
+h11==0.14.0
+hangul-romanize==0.1.0
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.5
+humanfriendly==10.0
+idna==3.7
+importlib_resources==6.4.0
+inflect==7.3.1
+itsdangerous==2.2.0
+jaconv==0.4.0
+jamo==0.4.1
+jieba==0.42.1
+Jinja2==3.1.4
+joblib==1.4.2
+jsonlines==1.2.0
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+langcodes==3.4.0
+language_data==1.2.0
+lazy_loader==0.4
+librosa==0.10.2.post1
+llvmlite==0.43.0
+marisa-trie==1.2.0
+Markdown==3.6
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.8.4
+mdurl==0.1.2
+mecab-python3==1.0.9
+mojimoji==0.0.13
+more-itertools==10.3.0
+mpmath==1.3.0
+msgpack==1.0.8
+multidict==6.0.5
+murmurhash==1.0.10
+networkx==2.8.8
+nltk==3.8.1
+num2words==0.5.13
+numba==0.60.0
+numpy==1.26.4
+onnxruntime==1.18.1
+orjson==3.10.6
+packaging==24.1
+pandas==1.5.3
+pillow==10.4.0
+platformdirs==4.2.2
+pooch==1.8.2
+preshed==3.0.9
+protobuf==4.25.3
+psutil==6.0.0
+pycparser==2.22
+pydantic==2.3.0
+pydantic_core==2.6.3
+pydub==0.25.1
+pygame==2.6.0
+Pygments==2.18.0
+pynndescent==0.5.13
+pyparsing==3.1.2
+pypinyin==0.51.0
+pysbd==0.3.4
+python-crfsuite==0.9.10
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.9
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+rich==13.7.1
+rpds-py==0.19.0
+ruff==0.5.2
+safetensors==0.4.3
+scikit-learn==1.5.1
+scipy==1.11.4
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+smart-open==6.4.0
+sniffio==1.3.1
+soundfile==0.12.1
+soxr==0.3.7
+spacy==3.7.4
+spacy-legacy==3.0.12
+spacy-loggers==1.0.5
+srsly==2.4.8
+starlette==0.27.0
+SudachiDict-core==20240409
+SudachiPy==0.6.8
+sympy==1.13.0
+tensorboard==2.17.0
+tensorboard-data-server==0.7.2
+thinc==8.2.5
+threadpoolctl==3.5.0
+tokenizers==0.19.1
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.3.1
+torchaudio==2.3.1
+tqdm==4.66.4
+trainer==0.0.36
+transformers==4.42.4
+TTS==0.21.3
+typeguard==4.3.0
+typer==0.12.5
+typing_extensions==4.12.2
+tzdata==2024.1
+tzlocal==5.2
+umap-learn==0.5.6
+Unidecode==1.3.8
+unidic-lite==1.0.8
+urllib3==2.2.2
+uvicorn==0.30.1
+uvloop==0.19.0
+wasabi==1.1.3
+watchfiles==0.22.0
+weasel==0.3.4
+websockets==11.0.3
+Werkzeug==3.0.3
+wrapt==1.16.0
+yarl==1.9.4

notebook/xtts_finetune_webui.ipynb

@@ -0,0 +1,154 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4",
+      "authorship_tag": "ABX9TyP4Z6m49+bXNW/J1fP7ZIEB",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/DrewThomasson/xtts-finetune-webui/blob/main/notebook/xtts_finetune_webui.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Welcome to the *xtts*-finetune-webui gradio gui!\n",
+        "\n",
+        "This webui is a slightly modified copy of the official webui for finetune xtts.\n",
+        "\n",
+        "If you are looking for an option for normal XTTS use look here https://github.com/daswer123/xtts-webui"
+      ],
+      "metadata": {
+        "id": "OVjEG_yGoC2W"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "44HpAIVRfJve"
+      },
+      "outputs": [],
+      "source": [
+        "# @title 🛠️ Install requirments\n",
+        "#!DEBIAN_FRONTEND=noninteractive\n",
+        "!sudo apt-get update # && sudo apt-get -y upgrade\n",
+        "!sudo apt-get -y install libegl1\n",
+        "!sudo apt-get -y install libopengl0\n",
+        "!sudo apt-get -y install libxcb-cursor0\n",
+        "!pip install -r https://raw.githubusercontent.com/daswer123/xtts-finetune-webui/main/requirements.txt\n",
+        "!pip install gradio==4.44.1\n",
+        "!pip install fastapi==0.103.1\n",
+        "!pip install pydantic==2.3.0"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# @title 🚀 Run interface\n",
+        "%cd /content/\n",
+        "!git clone https://github.com/DrewThomasson/xtts-finetune-webui.git\n",
+        "%cd /content/xtts-finetune-webui\n",
+        "!python xtts_demo.py --share"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "62Da1Q5AgN3W"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import shutil\n",
+        "import requests\n",
+        "import os\n",
+        "from tqdm import tqdm  # Progress bar library\n",
+        "\n",
+        "# Define the paths\n",
+        "finetune_dir = '/content/xtts-finetune-webui/finetune_models/ready'  # @param {type:\"string\"}\n",
+        "dataset_dir = '/content/xtts-finetune-webui/finetune_models/dataset'  # @param {type:\"string\"}\n",
+        "\n",
+        "# Create a temporary directory to collect both folders before zipping\n",
+        "temp_dir = \"/content/temp_finetune_dataset\"\n",
+        "os.makedirs(temp_dir, exist_ok=True)\n",
+        "\n",
+        "# Copy both directories into the temporary directory with a progress bar\n",
+        "def copy_with_progress(src, dst):\n",
+        "    total_files = sum(len(files) for _, _, files in os.walk(src))\n",
+        "    with tqdm(total=total_files, desc=f\"Copying {os.path.basename(src)}\") as pbar:\n",
+        "        for root, _, files in os.walk(src):\n",
+        "            rel_path = os.path.relpath(root, src)\n",
+        "            target_path = os.path.join(dst, rel_path)\n",
+        "            os.makedirs(target_path, exist_ok=True)\n",
+        "            for file in files:\n",
+        "                shutil.copy(os.path.join(root, file), target_path)\n",
+        "                pbar.update(1)\n",
+        "\n",
+        "copy_with_progress(finetune_dir, os.path.join(temp_dir, \"ready\"))\n",
+        "copy_with_progress(dataset_dir, os.path.join(temp_dir, \"dataset\"))\n",
+        "\n",
+        "# Create a zip file of the combined directories with progress\n",
+        "zip_filename = \"finetune_and_dataset.zip\"\n",
+        "with tqdm(total=100, desc=\"Zipping files\") as pbar:\n",
+        "    shutil.make_archive(\"finetune_and_dataset\", 'zip', root_dir=temp_dir)\n",
+        "    pbar.update(100)\n",
+        "\n",
+        "# Define a function to stream the upload with a progress bar\n",
+        "def upload_with_progress(file_path, url):\n",
+        "    file_size = os.path.getsize(file_path)\n",
+        "    with open(file_path, 'rb') as f, tqdm(\n",
+        "        total=file_size, unit='B', unit_scale=True, desc=\"Uploading\"\n",
+        "    ) as progress:\n",
+        "        response = requests.post(\n",
+        "            url,\n",
+        "            files={\"file\": (file_path, f)},\n",
+        "            stream=True,\n",
+        "            headers={\"Connection\": \"keep-alive\"},\n",
+        "        )\n",
+        "        # Update the progress bar as chunks are sent\n",
+        "        for chunk in response.iter_content(chunk_size=4096):\n",
+        "            if chunk:\n",
+        "                progress.update(len(chunk))\n",
+        "    return response\n",
+        "\n",
+        "# Upload the zip file to file.io with a progress bar\n",
+        "response = upload_with_progress(zip_filename, \"https://file.io/?expires=1d\")\n",
+        "\n",
+        "# Parse the response and display the download link\n",
+        "if response.status_code == 200:\n",
+        "    download_link = response.json().get('link', 'Error: No link found.')\n",
+        "    print(f\"Your file is ready: {download_link}\")\n",
+        "else:\n",
+        "    print(f\"Failed to upload: {response.status_code} - {response.text}\")\n"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "MYBWgKevr6S3"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

requirements.txt

@@ -1,11 +1,7 @@
-# Python packages for your Hugging Face Space
 faster_whisper==1.0.2
 gradio==4.13.0
 spacy==3.7.4
-coqui-tts[languages]==0.24.1
+coqui-tts[languages] == 0.24.2
+
 cutlet
 fugashi[unidic-lite]
-
-# CUDA-enabled PyTorch and Torchaudio
-torch==2.1.1+cu118
-torchaudio==2.1.1+cu118 --index-url https://download.pytorch.org/whl/cu118

xtts_demo.py

@@ -0,0 +1,697 @@
+import argparse
+import os
+import sys
+import tempfile
+from pathlib import Path
+
+import os
+import shutil
+import glob
+
+import gradio as gr
+import librosa.display
+import numpy as np
+
+import torch
+import torchaudio
+import traceback
+from utils.formatter import format_audio_list,find_latest_best_model, list_audios
+from utils.gpt_train import train_gpt
+
+from faster_whisper import WhisperModel
+
+from TTS.tts.configs.xtts_config import XttsConfig
+from TTS.tts.models.xtts import Xtts
+
+from TTS.tts.configs.xtts_config import XttsConfig
+from TTS.tts.models.xtts import Xtts
+
+# Clear logs
+def remove_log_file(file_path):
+     log_file = Path(file_path)
+
+     if log_file.exists() and log_file.is_file():
+         log_file.unlink()
+
+# remove_log_file(str(Path.cwd() / "log.out"))
+
+def clear_gpu_cache():
+    # clear the GPU cache
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+XTTS_MODEL = None
+def load_model(xtts_checkpoint, xtts_config, xtts_vocab,xtts_speaker):
+    global XTTS_MODEL
+    clear_gpu_cache()
+    if not xtts_checkpoint or not xtts_config or not xtts_vocab:
+        return "You need to run the previous steps or manually set the `XTTS checkpoint path`, `XTTS config path`, and `XTTS vocab path` fields !!"
+    config = XttsConfig()
+    config.load_json(xtts_config)
+    XTTS_MODEL = Xtts.init_from_config(config)
+    print("Loading XTTS model! ")
+    XTTS_MODEL.load_checkpoint(config, checkpoint_path=xtts_checkpoint, vocab_path=xtts_vocab,speaker_file_path=xtts_speaker, use_deepspeed=False)
+    if torch.cuda.is_available():
+        XTTS_MODEL.cuda()
+
+    print("Model Loaded!")
+    return "Model Loaded!"
+
+def run_tts(lang, tts_text, speaker_audio_file, temperature, length_penalty,repetition_penalty,top_k,top_p,sentence_split,use_config):
+    if XTTS_MODEL is None or not speaker_audio_file:
+        return "You need to run the previous step to load the model !!", None, None
+
+    gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(audio_path=speaker_audio_file, gpt_cond_len=XTTS_MODEL.config.gpt_cond_len, max_ref_length=XTTS_MODEL.config.max_ref_len, sound_norm_refs=XTTS_MODEL.config.sound_norm_refs)
+    
+    if use_config:
+        out = XTTS_MODEL.inference(
+            text=tts_text,
+            language=lang,
+            gpt_cond_latent=gpt_cond_latent,
+            speaker_embedding=speaker_embedding,
+            temperature=XTTS_MODEL.config.temperature, # Add custom parameters here
+            length_penalty=XTTS_MODEL.config.length_penalty,
+            repetition_penalty=XTTS_MODEL.config.repetition_penalty,
+            top_k=XTTS_MODEL.config.top_k,
+            top_p=XTTS_MODEL.config.top_p,
+            enable_text_splitting = True
+        )
+    else:
+        out = XTTS_MODEL.inference(
+            text=tts_text,
+            language=lang,
+            gpt_cond_latent=gpt_cond_latent,
+            speaker_embedding=speaker_embedding,
+            temperature=temperature, # Add custom parameters here
+            length_penalty=length_penalty,
+            repetition_penalty=float(repetition_penalty),
+            top_k=top_k,
+            top_p=top_p,
+            enable_text_splitting = sentence_split
+        )
+
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
+        out["wav"] = torch.tensor(out["wav"]).unsqueeze(0)
+        out_path = fp.name
+        torchaudio.save(out_path, out["wav"], 24000)
+
+    return "Speech generated !", out_path, speaker_audio_file
+
+
+def load_params_tts(out_path,version):
+    
+    out_path = Path(out_path)
+
+    # base_model_path = Path.cwd() / "models" / version 
+
+    # if not base_model_path.exists():
+    #     return "Base model not found !","","",""
+
+    ready_model_path = out_path / "ready" 
+
+    vocab_path =  ready_model_path / "vocab.json"
+    config_path = ready_model_path / "config.json"
+    speaker_path =  ready_model_path / "speakers_xtts.pth"
+    reference_path  = ready_model_path / "reference.wav"
+
+    model_path = ready_model_path / "model.pth"
+
+    if not model_path.exists():
+        model_path = ready_model_path / "unoptimize_model.pth"
+        if not model_path.exists():
+          return "Params for TTS not found", "", "", ""         
+
+    return "Params for TTS loaded", model_path, config_path, vocab_path,speaker_path, reference_path
+     
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(
+        description="""XTTS fine-tuning demo\n\n"""
+        """
+        Example runs:
+        python3 TTS/demos/xtts_ft_demo/xtts_demo.py --port 
+        """,
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+    parser.add_argument(
+        "--share",
+        action="store_true",
+        default=False,
+        help="Enable sharing of the Gradio interface via public link.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        help="Port to run the gradio demo. Default: 5003",
+        default=5003,
+    )
+    parser.add_argument(
+        "--out_path",
+        type=str,
+        help="Output path (where data and checkpoints will be saved) Default: output/",
+        default=str(Path.cwd() / "finetune_models"),
+    )
+
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        help="Number of epochs to train. Default: 6",
+        default=6,
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        help="Batch size. Default: 2",
+        default=2,
+    )
+    parser.add_argument(
+        "--grad_acumm",
+        type=int,
+        help="Grad accumulation steps. Default: 1",
+        default=1,
+    )
+    parser.add_argument(
+        "--max_audio_length",
+        type=int,
+        help="Max permitted audio size in seconds. Default: 11",
+        default=11,
+    )
+
+    args = parser.parse_args()
+
+    with gr.Blocks() as demo:
+        with gr.Tab("1 - Data processing"):
+            out_path = gr.Textbox(
+                label="Output path (where data and checkpoints will be saved):",
+                value=args.out_path,
+            )
+            # upload_file = gr.Audio(
+            #     sources="upload",
+            #     label="Select here the audio files that you want to use for XTTS trainining !",
+            #     type="filepath",
+            # )
+            upload_file = gr.File(
+                file_count="multiple",
+                label="Select here the audio files that you want to use for XTTS trainining (Supported formats: wav, mp3, and flac)",
+            )
+            
+            audio_folder_path = gr.Textbox(
+                label="Path to the folder with audio files (optional):",
+                value="",
+            )
+
+            whisper_model = gr.Dropdown(
+                label="Whisper Model",
+                value="large-v3",
+                choices=[
+                    "large-v3",
+                    "large-v2",
+                    "large",
+                    "medium",
+                    "small"
+                ],
+            )
+
+            lang = gr.Dropdown(
+                label="Dataset Language",
+                value="en",
+                choices=[
+                    "en",
+                    "es",
+                    "fr",
+                    "de",
+                    "it",
+                    "pt",
+                    "pl",
+                    "tr",
+                    "ru",
+                    "nl",
+                    "cs",
+                    "ar",
+                    "zh",
+                    "hu",
+                    "ko",
+                    "ja"
+                ],
+            )
+            progress_data = gr.Label(
+                label="Progress:"
+            )
+            # demo.load(read_logs, None, logs, every=1)
+
+            prompt_compute_btn = gr.Button(value="Step 1 - Create dataset")
+        
+            def preprocess_dataset(audio_path, audio_folder_path, language, whisper_model, out_path, train_csv, eval_csv, progress=gr.Progress(track_tqdm=True)):
+                clear_gpu_cache()
+            
+                train_csv = ""
+                eval_csv = ""
+            
+                out_path = os.path.join(out_path, "dataset")
+                os.makedirs(out_path, exist_ok=True)
+            
+                if audio_folder_path:
+                    audio_files = list(list_audios(audio_folder_path))
+                else:
+                    audio_files = audio_path
+            
+                if not audio_files:
+                    return "No audio files found! Please provide files via Gradio or specify a folder path.", "", ""
+                else:
+                    try:
+                        # Loading Whisper
+                        device = "cuda" if torch.cuda.is_available() else "cpu" 
+                        
+                        # Detect compute type 
+                        if torch.cuda.is_available():
+                            compute_type = "float16"
+                        else:
+                            compute_type = "float32"
+                        
+                        asr_model = WhisperModel(whisper_model, device=device, compute_type=compute_type)
+                        train_meta, eval_meta, audio_total_size = format_audio_list(audio_files, asr_model=asr_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 an error !! Please check the console to verify the full error message! \n Error summary: {error}", "", ""
+            
+                # clear_gpu_cache()
+            
+                # if audio total len is less than 2 minutes raise an error
+                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
+
+
+        with gr.Tab("2 - Fine-tuning XTTS Encoder"):
+            load_params_btn = gr.Button(value="Load Params from output folder")
+            version = gr.Dropdown(
+                label="XTTS base version",
+                value="v2.0.2",
+                choices=[
+                    "v2.0.3",
+                    "v2.0.2",
+                    "v2.0.1",
+                    "v2.0.0",
+                    "main"
+                ],
+            )
+            train_csv = gr.Textbox(
+                label="Train CSV:",
+            )
+            eval_csv = gr.Textbox(
+                label="Eval CSV:",
+            )
+            custom_model = gr.Textbox(
+                label="(Optional) Custom model.pth file , leave blank if you want to use the base file.",
+                value="",
+            )
+            num_epochs =  gr.Slider(
+                label="Number of epochs:",
+                minimum=1,
+                maximum=100,
+                step=1,
+                value=args.num_epochs,
+            )
+            batch_size = gr.Slider(
+                label="Batch size:",
+                minimum=2,
+                maximum=512,
+                step=1,
+                value=args.batch_size,
+            )
+            grad_acumm = gr.Slider(
+                label="Grad accumulation steps:",
+                minimum=2,
+                maximum=128,
+                step=1,
+                value=args.grad_acumm,
+            )
+            max_audio_length = gr.Slider(
+                label="Max permitted audio size in seconds:",
+                minimum=2,
+                maximum=20,
+                step=1,
+                value=args.max_audio_length,
+            )
+            clear_train_data = gr.Dropdown(
+                label="Clear train data, you will delete selected folder, after optimizing",
+                value="none",
+                choices=[
+                    "none",
+                    "run",
+                    "dataset",
+                    "all"
+                ])
+            
+            progress_train = gr.Label(
+                label="Progress:"
+            )
+
+            # demo.load(read_logs, None, logs_tts_train, every=1)
+            train_btn = gr.Button(value="Step 2 - Run the training")
+            optimize_model_btn = gr.Button(value="Step 2.5 - Optimize the model")
+            
+            import os
+            import shutil
+            from pathlib import Path
+            import traceback
+            
+            def train_model(custom_model, version, language, train_csv, eval_csv, num_epochs, batch_size, grad_acumm, output_path, max_audio_length):
+                clear_gpu_cache()
+            
+                run_dir = Path(output_path) / "run"
+            
+                # Remove train dir
+                if run_dir.exists():
+                    shutil.rmtree(run_dir)
+                
+                # Check if the dataset language matches the language you specified 
+                lang_file_path = Path(output_path) / "dataset" / "lang.txt"
+            
+                # Check if lang.txt already exists and contains a different language
+                current_language = None
+                if lang_file_path.exists():
+                    with open(lang_file_path, 'r', encoding='utf-8') as existing_lang_file:
+                        current_language = existing_lang_file.read().strip()
+                        if current_language != language:
+                            print("The language that was prepared for the dataset does not match the specified language. Change the language to the one specified in the dataset")
+                            language = current_language
+                        
+                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:
+                    # convert seconds to waveform frames
+                    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_acumm, 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")
+            
+                # Move reference audio to output folder and rename it
+                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):
+                # print(out_path)
+                out_path = Path(out_path)  # Ensure that out_path is a Path object.
+            
+                ready_dir = out_path / "ready"
+                run_dir = out_path / "run"
+                dataset_dir = out_path / "dataset"
+            
+                # Clear specified training data directories.
+                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}")
+            
+                # Get full path to model
+                model_path = ready_dir / "unoptimize_model.pth"
+
+                if not model_path.is_file():
+                    return "Unoptimized model not found in ready folder", ""
+            
+                # Load the checkpoint and remove unnecessary parts.
+                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]
+
+                # Make sure out_path is a Path object or convert it to Path
+                os.remove(model_path)
+
+                  # Save the optimized model.
+                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)
+
+                clear_gpu_cache()
+        
+                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"
+
+                # Write the target language to lang.txt in the output directory
+                lang_file_path =  dataset_path / "lang.txt"
+
+                # Check if lang.txt already exists and contains a different language
+                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()
+
+                clear_gpu_cache()
+
+                print(current_language)
+                return "The data has been updated", eval_train, eval_csv, current_language
+
+        with gr.Tab("3 - Inference"):
+            with gr.Row():
+                with gr.Column() as col1:
+                    load_params_tts_btn = gr.Button(value="Load params for TTS from output folder")
+                    xtts_checkpoint = gr.Textbox(
+                        label="XTTS checkpoint path:",
+                        value="",
+                    )
+                    xtts_config = gr.Textbox(
+                        label="XTTS config path:",
+                        value="",
+                    )
+
+                    xtts_vocab = gr.Textbox(
+                        label="XTTS vocab path:",
+                        value="",
+                    )
+                    xtts_speaker = gr.Textbox(
+                        label="XTTS speaker path:",
+                        value="",
+                    )
+                    progress_load = gr.Label(
+                        label="Progress:"
+                    )
+                    load_btn = gr.Button(value="Step 3 - Load Fine-tuned XTTS model")
+
+                with gr.Column() as col2:
+                    speaker_reference_audio = gr.Textbox(
+                        label="Speaker reference audio:",
+                        value="",
+                    )
+                    tts_language = gr.Dropdown(
+                        label="Language",
+                        value="en",
+                        choices=[
+                            "en",
+                            "es",
+                            "fr",
+                            "de",
+                            "it",
+                            "pt",
+                            "pl",
+                            "tr",
+                            "ru",
+                            "nl",
+                            "cs",
+                            "ar",
+                            "zh",
+                            "hu",
+                            "ko",
+                            "ja",
+                        ]
+                    )
+                    tts_text = gr.Textbox(
+                        label="Input Text.",
+                        value="This model sounds really good and above all, it's reasonably fast.",
+                    )
+                    with gr.Accordion("Advanced settings", open=False) as acr:
+                        temperature = gr.Slider(
+                            label="temperature",
+                            minimum=0,
+                            maximum=1,
+                            step=0.05,
+                            value=0.75,
+                        )
+                        length_penalty  = gr.Slider(
+                            label="length_penalty",
+                            minimum=-10.0,
+                            maximum=10.0,
+                            step=0.5,
+                            value=1,
+                        )
+                        repetition_penalty = gr.Slider(
+                            label="repetition penalty",
+                            minimum=1,
+                            maximum=10,
+                            step=0.5,
+                            value=5,
+                        )
+                        top_k = gr.Slider(
+                            label="top_k",
+                            minimum=1,
+                            maximum=100,
+                            step=1,
+                            value=50,
+                        )
+                        top_p = gr.Slider(
+                            label="top_p",
+                            minimum=0,
+                            maximum=1,
+                            step=0.05,
+                            value=0.85,
+                        )
+                        sentence_split = gr.Checkbox(
+                            label="Enable text splitting",
+                            value=True,
+                        )
+                        use_config = gr.Checkbox(
+                            label="Use Inference settings from config, if disabled use the settings above",
+                            value=False,
+                        )
+                    tts_btn = gr.Button(value="Step 4 - Inference")
+
+                with gr.Column() as col3:
+                    progress_gen = gr.Label(
+                        label="Progress:"
+                    )
+                    tts_output_audio = gr.Audio(label="Generated Audio.")
+                    reference_audio = gr.Audio(label="Reference audio used.")
+
+            prompt_compute_btn.click(
+                fn=preprocess_dataset,
+                inputs=[
+                    upload_file,
+                    audio_folder_path,
+                    lang,
+                    whisper_model,
+                    out_path,
+                    train_csv,
+                    eval_csv
+                ],
+                outputs=[
+                    progress_data,
+                    train_csv,
+                    eval_csv,
+                ],
+            )
+
+
+            load_params_btn.click(
+                fn=load_params,
+                inputs=[out_path],
+                outputs=[
+                    progress_train,
+                    train_csv,
+                    eval_csv,
+                    lang
+                ]
+            )
+
+
+            train_btn.click(
+                fn=train_model,
+                inputs=[
+                    custom_model,
+                    version,
+                    lang,
+                    train_csv,
+                    eval_csv,
+                    num_epochs,
+                    batch_size,
+                    grad_acumm,
+                    out_path,
+                    max_audio_length,
+                ],
+                outputs=[progress_train, xtts_config, xtts_vocab, xtts_checkpoint,xtts_speaker, speaker_reference_audio],
+            )
+
+            optimize_model_btn.click(
+                fn=optimize_model,
+                inputs=[
+                    out_path,
+                    clear_train_data
+                ],
+                outputs=[progress_train,xtts_checkpoint],
+            )
+            
+            load_btn.click(
+                fn=load_model,
+                inputs=[
+                    xtts_checkpoint,
+                    xtts_config,
+                    xtts_vocab,
+                    xtts_speaker
+                ],
+                outputs=[progress_load],
+            )
+
+            tts_btn.click(
+                fn=run_tts,
+                inputs=[
+                    tts_language,
+                    tts_text,
+                    speaker_reference_audio,
+                    temperature,
+                    length_penalty,
+                    repetition_penalty,
+                    top_k,
+                    top_p,
+                    sentence_split,
+                    use_config
+                ],
+                outputs=[progress_gen, tts_output_audio,reference_audio],
+            )
+
+            load_params_tts_btn.click(
+                fn=load_params_tts,
+                inputs=[
+                    out_path,
+                    version
+                    ],
+                outputs=[progress_load,xtts_checkpoint,xtts_config,xtts_vocab,xtts_speaker,speaker_reference_audio],
+            )
+
+    demo.launch(
+        share=args.share,
+        debug=False,
+        server_port=args.port,
+        # inweb=True,
+        # server_name="localhost"
+    )