Improving docker build time by downloading the TTS model during build time, thereby enhancing firstrun.py

Dockerfile

@@ -18,10 +18,46 @@ ENV WITH_UI=true
 
 WORKDIR ${ALLTALK_DIR}
 
+##############################################################################
+# Download TTS models:
+##############################################################################
+COPY system/tts_engines/piper/available_models.json system/tts_engines/piper/
+COPY system/tts_engines/vits/available_models.json system/tts_engines/vits/
+COPY system/tts_engines/xtts/available_models.json system/tts_engines/xtts/
+RUN <<EOR
+    available_models="system/tts_engines/${TTS_MODEL}/available_models.json"
+    first_start_model=$(cat ${available_models} | jq -r '.first_start_model')
+    model=$(cat ${available_models} | jq ".models[] | select(.model_name==\"${first_start_model}\")")
+    folder_path=$( echo "${model}" | jq -r '.folder_path // empty' )
+    files_to_download=$( echo "${model}" | jq '.files_to_download[]?' )
+    github_rls_url=$( echo "${model}" | jq '.github_rls_url // empty' )
+
+    # Special case for vits: firstrun.py expects the ZIP download directly in the model folder.
+    # When firstrun.py unzips the file, the correct folder will be created automatically.
+    if [[ "$TTS_MODEL" == "vits" ]]; then
+      folder_path=""
+    fi
+
+    # Merging both fields:
+    files_to_download="${files_to_download} ${github_rls_url}"
+    target_dir="models/${TTS_MODEL}/${folder_path}"
+
+    echo "$files_to_download" | xargs -n 1 curl --create-dirs --output-dir "${target_dir}" -LO
+EOR
+
+##############################################################################
+# Download all RVC models:
+##############################################################################
+COPY system/tts_engines/rvc_files.json system/tts_engines/
+RUN <<EOR
+  rvc_files=$( jq -r '.[]' system/tts_engines/rvc_files.json )
+  echo "$rvc_files" | xargs -n 1 curl --create-dirs --output-dir models/rvc_base -LO
+EOR
+
 ##############################################################################
 # Install python dependencies (cannot use --no-deps because requirements are not complete)
 ##############################################################################
-COPY system/config system/config
+COPY system/config/*.whl system/config/
 COPY system/requirements/requirements_standalone.txt system/requirements/requirements_standalone.txt
 COPY system/requirements/requirements_parler.txt system/requirements/requirements_parler.txt
 ENV PIP_CACHE_DIR=${ALLTALK_DIR}/pip_cache
@@ -142,15 +178,6 @@ RUN mkdir -p /root/.triton/autotune
 ##############################################################################
 RUN find . -name model_settings.json -exec sed -i -e 's/"deepspeed_enabled": false/"deepspeed_enabled": true/g' {} \;
 
-##############################################################################
-# Download all RVC models:
-##############################################################################
-RUN <<EOR
-  jq -r '.[]' system/tts_engines/rvc_files.json > /tmp/rvc_files.txt
-  xargs -n 1 curl --create-dirs --output-dir models/rvc_base -LO < /tmp/rvc_files.txt
-  rm -f /tmp/rvc_files.txt
-EOR
-
 ##############################################################################
 # Start alltalk:
 ##############################################################################

system/config/firstrun.py

@@ -41,9 +41,8 @@ Functions:
 - **is_running_in_colab**: Detects if the script is running in Google Colab.
 - **is_running_in_docker**: Detects if the script is running in a Docker container.
 - **download_file**: Downloads a file with a progress bar and retry logic.
-- **setup_piper**: Sets up the Piper TTS engine and downloads associated files.
-- **setup_vits**: Sets up the VITS TTS engine and downloads associated files.
-- **setup_xtts**: Sets up the XTTS TTS engine and downloads associated files.
+- **setup_tts**: Sets up the TTS engine passed as argument.
+- **download_tts_model**: Downloads the TTS engine's associated files.
 - **update_tts_engines**: Updates the configuration to reflect the selected TTS engine.
 - **set_firstrun_model_false**: Marks the initial setup as complete in the configuration.
 - **warning_message**: Displays post-setup warnings or tips for the user.
@@ -65,6 +64,7 @@ import zipfile
 import requests
 from tqdm import tqdm
 from inputimeout import inputimeout, TimeoutOccurred
+from urllib.parse import urlsplit
 
 # Setting the module search path:
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.resolve()))
@@ -103,6 +103,11 @@ def download_file(url: str, dest_path: str, timeout: int = 30, retries: int = 3)
     Raises:
         ValueError: If the downloaded file size does not match the expected size.
     """
+    # Don't re-download if file already exists (which is the case when using Docker):
+    if Path(dest_path).exists():
+        print(f"File already exists at: {dest_path}. Skipping download.")
+        return True
+
     progress_bar = None
     for attempt in range(1, retries + 1):
         try:
@@ -166,41 +171,44 @@ def download_file(url: str, dest_path: str, timeout: int = 30, retries: int = 3)
     # Return False if all attempts fail
     return False
 
-def setup_piper():
-    """Download base piper files"""
-    os.makedirs(this_dir / 'models/piper', exist_ok=True)
-    os.makedirs(this_dir / 'models/xtts', exist_ok=True)
-    download_file("https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/ljspeech/high/en_US-ljspeech-high.onnx?download=true",
-                  this_dir / "models/piper/en_US-ljspeech-high.onnx")
-    download_file("https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/ljspeech/high/en_US-ljspeech-high.onnx.json?download=true.json",
-                  this_dir / "models/piper/en_US-ljspeech-high.onnx.json")
+def unzip_compressed_files(downloaded_files: []):
+    for file in downloaded_files:
+        _, file_extension = os.path.splitext(file)
+        if file_extension == '.zip':
+            with zipfile.ZipFile(file, 'r') as zip_ref:
+                zip_ref.extractall(file.parent)
+            os.remove(file)
 
-def setup_vits():
-    """Download base vits files"""
-    os.makedirs(this_dir / 'models/vits', exist_ok=True)
-    os.makedirs(this_dir / 'models/xtts', exist_ok=True)
-    zip_path = this_dir / 'models/vits/vits.zip'
-    download_file("https://coqui.gateway.scarf.sh/v0.6.1_models/tts_models--en--vctk--vits.zip", zip_path)
-    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-        zip_ref.extractall(this_dir / 'models/vits')
-    os.remove(zip_path)
+def setup_tts(tts_model: str, on_download_completed = unzip_compressed_files):
+    download_tts_model(tts_model, on_download_completed)
+    update_tts_engines(tts_model)
+    set_firstrun_model_false()
 
-def setup_xtts():
-    """Download base xtts files"""
-    os.makedirs(this_dir / 'models/xtts/xttsv2_2.0.3', exist_ok=True)
-    base_url = "https://huggingface.co/coqui/XTTS-v2/resolve/v2.0.3/"
-    files = [
-        "LICENSE.txt",
-        "README.md",
-        "config.json",
-        "model.pth",
-        "dvae.pth",
-        "mel_stats.pth",
-        "speakers_xtts.pth",
-        "vocab.json"
-    ]
-    for each_file in files:
-        download_file(base_url + each_file + "?download=true", this_dir / f"models/xtts/xttsv2_2.0.3/{each_file}")
+def download_tts_model(tts_model: str, on_download_completed = None):
+    with open(os.path.join(this_dir, "system", "tts_engines", tts_model, "available_models.json")) as f:
+        available_models = json.load(f)
+        first_start_model = available_models["first_start_model"]
+        model = next(m for m in available_models["models"] if m["model_name"] == first_start_model)
+
+        # folder_path: Special case for vits: firstrun.py expects the ZIP download directly in the model folder.
+        # When firstrun.py unzips the file, the correct folder will be created automatically.
+        folder_path = "" if tts_model == "vits" else model.get("folder_path", "")
+        files_to_download = model.get("files_to_download", []) or [model.get("github_rls_url", None)]
+        target_dir = this_dir / f"models/{tts_model}/{folder_path}"
+
+        if type(files_to_download) == dict:
+            files_to_download = list(files_to_download.values())
+
+        downloaded_files = []
+        for file_download in files_to_download:
+            urlpath = urlsplit(file_download).path
+            file_name = os.path.basename(urlpath)
+            target_file = target_dir / file_name
+            download_file(file_download, target_file)
+            downloaded_files.append(target_file)
+
+        if on_download_completed:
+            on_download_completed(downloaded_files)
 
 def update_tts_engines(engine):
     """Set engine to users choice that matches the download"""
@@ -282,18 +290,9 @@ args = parser.parse_args()
 if config.firstrun_model:
     # Handle command-line argument for tts_model
     if args.tts_model:
-        if args.tts_model == 'piper':
-            setup_piper()
-            update_tts_engines('piper')
-            set_firstrun_model_false()
-        elif args.tts_model == 'vits':
-            setup_vits()
-            update_tts_engines('vits')
-            set_firstrun_model_false()
-        elif args.tts_model == 'xtts':
-            setup_xtts()
-            update_tts_engines('xtts')
-            set_firstrun_model_false()
+        if args.tts_model in {'piper', 'vits', 'xtts'}:
+            print(f"[{branding}TTS] Setting up model '{args.tts_model}'.")
+            setup_tts(args.tts_model)
         elif args.tts_model == 'none':
             print(f"[{branding}TTS] No TTS model setup requested.")
         print(f"[{branding}TTS] Setup completed for {args.tts_model}. Exiting.")
@@ -302,9 +301,7 @@ if config.firstrun_model:
     # Check for Colab or Docker environment first
     if is_running_in_colab() or is_running_in_docker():
         print(f"[{branding}TTS] Detected Colab/Docker environment - automatically selecting Piper")
-        setup_piper()
-        update_tts_engines('piper')
-        set_firstrun_model_false()
+        setup_tts('piper')
         warning_message()
         sys.exit()
 
@@ -363,20 +360,10 @@ if config.firstrun_model:
         warning_message()
 
     if user_choice is None or user_choice.lower() != 'own':
-        if selected_model['name'] == 'piper':
-            setup_piper()
-            update_tts_engines('piper')
-        elif selected_model['name'] == 'vits':
-            setup_vits()
-            update_tts_engines('vits')
-        elif selected_model['name'] == 'xtts':
-            setup_xtts()
-            update_tts_engines('xtts')
-
-        print(f"[{branding}TTS] {selected_model['name']} model downloaded and configuration updated successfully.")
-        # Set firstrun_model to false
-        set_firstrun_model_false()
-        warning_message()
-        sys.exit()
+        if selected_model['name'] in {'piper', 'vits', 'xtts'}:
+            setup_tts(selected_model['name'])
+            print(f"[{branding}TTS] {selected_model['name']} model downloaded and configuration updated successfully.")
+            warning_message()
+            sys.exit()
 else:
     sys.exit()

system/tts_engines/piper/available_models.json

@@ -1,5 +1,5 @@
 {
-  "first_start_model": "en_GB-jenny_dioco-medium",
+  "first_start_model": "en_US-ljspeech-high",
   "models": [
     {
       "model_name": "ar_JO-kareem-low",