Command-Line Arguments

7. **Command-Line Arguments**:
    - Accepts a `--tts_model` argument to bypass the interactive menu and directly set up a specific TTS engine
      on the first time start-up:
        - `piper`: Sets up the Piper TTS engine.
        - `vits`: Sets up the VITS TTS engine.
        - `xtts`: Sets up the XTTS TTS engine.
        - `none`: Skips model setup entirely.

system/config/firstrun.py

@@ -1,4 +1,62 @@
+"""
+firstrun.py
+
+This script handles the first-time setup for the Alltalk TTS system. It allows the user to download and configure 
+the initial Text-to-Speech (TTS) engine and associated models. It provides automatic detection for Docker or Google 
+Colab environments to simplify the setup process, as well as an interactive menu or command-line arguments for manual configuration.
+
+Features:
+---------
+1. **Automatic Environment Detection**:
+    - Detects if the script is running in a Docker container or Google Colab environment.
+    - Automatically sets up the Piper TTS engine in these cases.
+
+2. **Command-Line Arguments**:
+    - Accepts a `--tts_model` argument to bypass the interactive menu and directly set up a specific TTS engine:
+        - `piper`: Sets up the Piper TTS engine.
+        - `vits`: Sets up the VITS TTS engine.
+        - `xtts`: Sets up the XTTS TTS engine.
+        - `none`: Skips model setup entirely.
+
+3. **Interactive Menu**:
+    - If no command-line argument is provided and the environment is not Docker/Colab, it presents an interactive 
+      menu for the user to select a TTS model.
+
+4. **Configuration Management**:
+    - Updates the Alltalk TTS configuration files to reflect the selected TTS engine and marks the initial setup as complete.
+
+Usage:
+------
+Command-Line:
+    python firstrun.py --tts_model [piper|vits|xtts|none]
+
+Interactive Mode:
+    python firstrun.py
+
+In Docker/Colab:
+    The script automatically selects and sets up the Piper TTS engine without user interaction.
+
+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.
+- **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.
+
+Notes:
+------
+- Ensure network access is available for downloading models during setup.
+- Configuration files are automatically updated and backed up if needed.
+
+"""
+import argparse
 import json
+from pathlib import Path
 import os
 import shutil
 import time
@@ -7,12 +65,10 @@ import zipfile
 import requests
 from tqdm import tqdm
 from inputimeout import inputimeout, TimeoutOccurred
-from pathlib import Path
 
 # Setting the module search path:
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.resolve()))
-
+from config import AlltalkTTSEnginesConfig, AlltalkConfig #pylint: disable=wrong-import-position
-from config import AlltalkTTSEnginesConfig, AlltalkConfig
 
 this_dir = Path(__file__).parent.resolve().parent.resolve().parent.resolve()
 
@@ -28,7 +84,7 @@ def is_running_in_docker():
     """Test for google colab"""
     path = '/proc/self/cgroup'
     return os.path.exists('/.dockerenv') or (
-        os.path.exists(path) and any('docker' in line for line in open(path))
+        os.path.exists(path) and any('docker' in line for line in open(path, encoding='utf-8'))
     )
 
 def download_file(url: str, dest_path: str, timeout: int = 30, retries: int = 3) -> bool:
@@ -36,51 +92,78 @@ def download_file(url: str, dest_path: str, timeout: int = 30, retries: int = 3)
     Download a file from a URL to a specified destination path with progress bar.
     
     Args:
-        url (str): URL of the file to download
+        url (str): URL of the file to download.
-        dest_path (str): Destination path where the file will be saved
+        dest_path (str): Destination path where the file will be saved.
         timeout (int, optional): Connection and read timeout in seconds. Defaults to 30.
         retries (int, optional): Number of retry attempts for failed downloads. Defaults to 3.
+    
+    Returns:
+        bool: True if the file was successfully downloaded, False otherwise.
+    
+    Raises:
+        ValueError: If the downloaded file size does not match the expected size.
     """
     progress_bar = None
-    for attempt in range(retries):
+    for attempt in range(1, retries + 1):
         try:
-            response = requests.get(url, stream=True, 
+            print(f"Attempting to download: {url} (Attempt {attempt}/{retries})")
-                                 timeout=(timeout, timeout))
+            response = requests.get(url, stream=True, timeout=(timeout, timeout))
             response.raise_for_status()
+
             total_size = int(response.headers.get('content-length', 0))
             block_size = 1024  # 1 Kibibyte
+
             # Initialize progress bar
             progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True,
-                              desc=f"Downloading {Path(dest_path).name} (Attempt {attempt + 1}/{retries})")
+                                desc=f"Downloading {Path(dest_path).name} (Attempt {attempt}/{retries})")
             Path(dest_path).parent.mkdir(parents=True, exist_ok=True)
+
             with open(dest_path, 'wb') as dwn_file:
                 for data in response.iter_content(block_size):
                     if data:
                         dwn_file.write(data)
                         progress_bar.update(len(data))
+
             if progress_bar:
                 progress_bar.close()
+
+            # Verify file size
             if progress_bar.n != total_size and total_size > 0:
-                raise requests.exceptions.RequestException(
+                raise ValueError(
-                    f"Downloaded size ({progress_bar.n} bytes) does not match expected size ({total_size} bytes)"
+                    f"Downloaded size ({progress_bar.n} bytes) does not match expected size ({total_size} bytes)."
                 )
+
+            print(f"Download completed: {dest_path}")
             return True
+
+        except requests.Timeout:
+            print(f"Timeout error during download: {url}")
+        except requests.ConnectionError:
+            print(f"Connection error during download: {url}")
         except requests.exceptions.RequestException as e:
-            if progress_bar:
+            print(f"HTTP error during download: {url}. Details: {str(e)}")
-                progress_bar.close()
+        except ValueError as e:
+            print(f"File validation error: {str(e)}")
             if Path(dest_path).exists():
-                Path(dest_path).unlink()
+                Path(dest_path).unlink()  # Remove incomplete file
-            if attempt == retries - 1:
+            break  # Stop retrying on file validation errors
-                raise requests.exceptions.RequestException(
-                    f"Failed to download {url} after {retries} attempts: {str(e)}"
-                )
-            print(f"Download attempt {attempt + 1}/{retries} failed: {str(e)}")
-            print("Retrying in 5 seconds...")
-            time.sleep(5)
         except Exception as e:
+            print(f"Unexpected error during download: {str(e)}")
+            break  # Stop retrying on unexpected errors
+        finally:
             if progress_bar:
                 progress_bar.close()
-            raise Exception(f"Unexpected error downloading {url}: {str(e)}")
+
+        if Path(dest_path).exists():
+            Path(dest_path).unlink()  # Remove incomplete file
+
+        if attempt < retries:
+            print(f"Retrying in 5 seconds... (Attempt {attempt + 1}/{retries})")
+            time.sleep(5)
+        else:
+            print(f"Exhausted all retries for {url}.")
+
+    # Return False if all attempts fail
     return False
 
 def setup_piper():
@@ -117,7 +200,7 @@ def setup_xtts():
         "vocab.json"
     ]
     for each_file in files:
-        download_file(base_url + file + "?download=true", this_dir / f"models/xtts/xttsv2_2.0.3/{each_file}")
+        download_file(base_url + each_file + "?download=true", this_dir / f"models/xtts/xttsv2_2.0.3/{each_file}")
 
 def update_tts_engines(engine):
     """Set engine to users choice that matches the download"""
@@ -129,6 +212,13 @@ def update_tts_engines(engine):
         print(f"[{branding}TTS] Error updating TTS engine configuration: {str(e)}")
 
 def set_firstrun_model_false():
+    """
+    Marks the first-run setup as complete in the Alltalk configuration.
+
+    This function updates the `firstrun_model` flag in the configuration file to `False`,
+    indicating that the initial setup has been completed. It attempts to save the updated 
+    configuration and verifies the changes, retrying with a direct file write if necessary.
+    """
     try:
         # Get a fresh instance of the config
         config_fr = AlltalkConfig.get_instance()
@@ -151,6 +241,7 @@ def set_firstrun_model_false():
         print(f"[{branding}TTS] Error saving firstrun configuration: {str(e)}")
 
 def warning_message():
+    """AllTalk V1 Warning message re updating"""
     print(f"[{branding}TTS]")
     print(f"[{branding}TTS] \033[93mIf you have you have UPGRADED from v1 ensure you have re-installed\033[0m")
     print(f"[{branding}TTS] \033[93mthe requirements. Otherwise you will get failures and errors!\033[0m")
@@ -181,22 +272,47 @@ if "ip_address" in config_json:
 
 config = AlltalkConfig.get_instance()
 branding = config.branding
+# Argument parser setup
+parser = argparse.ArgumentParser(description="TTS Setup Script")
+parser.add_argument('--tts_model', type=str, choices=['piper', 'vits', 'xtts', 'none'],
+                    help="Specify TTS model to set up (piper, vits, xtts, or none)")
+args = parser.parse_args()
 
 # Check if firstrun_model is true
 if config.firstrun_model:
-    # Check for Colab or Docker environment
+    # 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()
+        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.")
+        sys.exit()
+
+    # 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()
         warning_message()
-        exit()
+        sys.exit()
 
+    # Present the menu if no argument is passed and not in Colab/Docker
     print(f"[{branding}TTS]")
-    print(f"[{branding}TTS] \033[92mThis is the first time startup.. Please download a start TTS model. Other TTS engines\033[0m")
+    print(f"[{branding}TTS] \033[92mThis is the first-time startup.. Please download a start TTS model. Other TTS engines\033[0m")
     print(f"[{branding}TTS] \033[92mand TTS models can be downloaded/managed in the Gradio Interface `TTS Engines Settings`\033[0m")
-    print(f"[{branding}TTS] \033[92mtab after inital setup.\033[0m")
+    print(f"[{branding}TTS] \033[92mtab after initial setup.\033[0m")
     print(f"[{branding}TTS]")
 
     # List of available models
@@ -228,12 +344,15 @@ if config.firstrun_model:
         print(f"[{branding}TTS] No input received. Proceeding with the default model (piper).")
     elif user_choice.lower() == 'own' or (user_choice.isdigit() and int(user_choice) == len(models) + 1):
         print(f"[{branding}TTS]")
-        print(f"[{branding}TTS] If you have your own XTTS models, please move them into the")
+        print(f"[{branding}TTS] Please use the Gradio interface TTS Engine Settings > [Engine Name] > Model Download")
-        print(f"[{branding}TTS] \033[93m/models/xtts/\033[0m folder")
+        print(f"[{branding}TTS] To download models for your selected TTS Engine. Or use the [Engine Name] help sections")
+        print(f"[{branding}TTS] for instructions on using your own TTS models.")
         print(f"[{branding}TTS]")
         os.makedirs('models/piper', exist_ok=True)
         os.makedirs('models/vits', exist_ok=True)
         os.makedirs('models/xtts/', exist_ok=True)
+        os.makedirs('models/f5-tts/', exist_ok=True)
+        os.makedirs('models/rvc_voices/', exist_ok=True)
         update_tts_engines('piper')
         set_firstrun_model_false()
         warning_message()