Updated firstrun better docker/colab support

In Docker or Google colab environments where people may not have access to a keyboard, just download piper (as its small) and dont wait the 60 seconds for a keypress.

system/config/firstrun.py

@@ -1,6 +1,7 @@
 import json
 import os
 import shutil
+import time
 import sys
 import zipfile
 import requests
@@ -15,18 +16,75 @@ from config import AlltalkTTSEnginesConfig, AlltalkConfig
 
 this_dir = Path(__file__).parent.resolve().parent.resolve().parent.resolve()
 
-def download_file(url, dest_path):
+def is_running_in_colab():
-    response = requests.get(url, stream=True)
+    """Test for google colab"""
-    total_size = int(response.headers.get('content-length', 0))
+    try:
-    block_size = 1024  # 1 Kibibyte
+        import google.colab
-    t = tqdm(total=total_size, unit='iB', unit_scale=True, desc=str(dest_path))
+        return True
-    with open(dest_path, 'wb') as file:
+    except ImportError:
-        for data in response.iter_content(block_size):
+        return False
-            t.update(len(data))
+
-            file.write(data)
+def is_running_in_docker():
-    t.close()
+    """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))
+    )
+
+def download_file(url: str, dest_path: str, timeout: int = 30, retries: int = 3) -> bool:
+    """
+    Download a file from a URL to a specified destination path with progress bar.
+    
+    Args:
+        url (str): URL of the file to download
+        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.
+    """
+    progress_bar = None
+    for attempt in range(retries):
+        try:
+            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})")
+            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()
+            if progress_bar.n != total_size and total_size > 0:
+                raise requests.exceptions.RequestException(
+                    f"Downloaded size ({progress_bar.n} bytes) does not match expected size ({total_size} bytes)"
+                )    
+            return True
+        except requests.exceptions.RequestException as e:
+            if progress_bar:
+                progress_bar.close()
+            if Path(dest_path).exists():
+                Path(dest_path).unlink()
+            if attempt == retries - 1:
+                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:
+            if progress_bar:
+                progress_bar.close()
+            raise Exception(f"Unexpected error downloading {url}: {str(e)}")
+    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",
@@ -35,6 +93,7 @@ def setup_piper():
                   this_dir / "models/piper/en_US-ljspeech-high.onnx.json")
 
 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'
@@ -44,6 +103,7 @@ def setup_vits():
     os.remove(zip_path)
 
 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 = [
@@ -56,17 +116,39 @@ def setup_xtts():
         "speakers_xtts.pth",
         "vocab.json"
     ]
-    for file in files:
+    for each_file in files:
-        download_file(base_url + file + "?download=true", this_dir / f"models/xtts/xttsv2_2.0.3/{file}")
+        download_file(base_url + file + "?download=true", this_dir / f"models/xtts/xttsv2_2.0.3/{each_file}")
-
-
 
 def update_tts_engines(engine):
-    AlltalkTTSEnginesConfig.get_instance().change_engine(engine).save()
+    """Set engine to users choice that matches the download"""
+    try:
+        config_instance = AlltalkTTSEnginesConfig.get_instance()
+        config_instance.change_engine(engine)
+        config_instance.save()
+    except Exception as e:
+        print(f"[{branding}TTS] Error updating TTS engine configuration: {str(e)}")
 
 def set_firstrun_model_false():
-    config.firstrun_model = False
+    try:
-    config.save()
+        # Get a fresh instance of the config
+        config_fr = AlltalkConfig.get_instance()
+        config_fr.firstrun_model = False
+
+        # Force a flush to disk immediately
+        config_fr.save()
+
+        # Verify the save worked
+        config_path = AlltalkConfig.default_config_path()
+        with open(config_path, 'r', encoding="utf-8") as f:
+            saved_config = json.load(f)
+            if saved_config.get('firstrun_model', True):
+                print(f"[{branding}TTS] Warning: Failed to save firstrun_model=False")
+                # Try one more time with direct file write
+                saved_config['firstrun_model'] = False
+                with open(config_path, 'w', encoding="utf-8") as f:
+                    json.dump(saved_config, f, indent=2)
+    except Exception as e:
+        print(f"[{branding}TTS] Error saving firstrun configuration: {str(e)}")
 
 def warning_message():
     print(f"[{branding}TTS]")
@@ -77,7 +159,7 @@ def warning_message():
     print(f"[{branding}TTS]")
 
 # Load confignew.json (might be version 1 or 2):
-with open(AlltalkConfig.default_config_path(), 'r') as file:
+with open(AlltalkConfig.default_config_path(), 'r', encoding="utf-8") as file:
     config_json = json.load(file)
 
 branding = config_json['branding']
@@ -102,6 +184,15 @@ branding = config.branding
 
 # Check if firstrun_model is true
 if config.firstrun_model:
+    # Check for Colab or Docker environment
+    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()
+
     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[92mand TTS models can be downloaded/managed in the Gradio Interface `TTS Engines Settings`\033[0m")
@@ -125,7 +216,7 @@ if config.firstrun_model:
     print(f"[{branding}TTS]    \033[94mIn \033[91m60 seconds\033[0m \033[94ma Piper model will be \033[91mdownloaded automatically.\033[0m")
     print(f"[{branding}TTS]")
 
-    # Auto-select model after 30 seconds
+    # Auto-select model after 60 seconds
     selected_model = models[0]  # Default to the first model (piper)
     try:
         user_choice = inputimeout(prompt=f"[{branding}TTS]    \033[92mEnter your choice 1-4: \033[0m ", timeout=60)
@@ -167,6 +258,6 @@ if config.firstrun_model:
         # Set firstrun_model to false
         set_firstrun_model_false()
         warning_message()
-        exit()
+        sys.exit()
 else:
-    exit()
+    sys.exit()

system/config/firstrun_tgwui.py

@@ -2,11 +2,13 @@ import json
 import os
 import shutil
 import sys
+import time
 import zipfile
 import requests
 from tqdm import tqdm
 from inputimeout import inputimeout, TimeoutOccurred
 from pathlib import Path
+from typing import Optional, Dict, List, Any
 
 # Setting the module search path:
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.resolve()))
@@ -15,156 +17,271 @@ from config import AlltalkTTSEnginesConfig, AlltalkConfig
 
 this_dir = Path(__file__).parent.resolve().parent.resolve().parent.resolve()
 
-def download_file(url, dest_path):
+def is_running_in_colab() -> bool:
-    response = requests.get(url, stream=True)
+    """Check if the script is running in Google Colab."""
-    total_size = int(response.headers.get('content-length', 0))
+    try:
-    block_size = 1024  # 1 Kibibyte
+        import google.colab
-    t = tqdm(total=total_size, unit='iB', unit_scale=True, desc=str(dest_path))
+        return True
-    with open(dest_path, 'wb') as file:
+    except ImportError:
-        for data in response.iter_content(block_size):
+        return False
-            t.update(len(data))
-            file.write(data)
-    t.close()
 
-def setup_piper():
+def is_running_in_docker() -> bool:
-    os.makedirs(this_dir / 'models/piper', exist_ok=True)
+    """Check if the script is running in a Docker container."""
-    os.makedirs(this_dir / 'models/xtts', exist_ok=True)
+    return os.path.exists('/.dockerenv') or (
-    download_file("https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/ljspeech/high/en_US-ljspeech-high.onnx?download=true",
+        os.path.exists('/proc/self/cgroup') and 
-                  this_dir / "models/piper/en_US-ljspeech-high.onnx")
+        any('docker' in line for line in open('/proc/self/cgroup'))
-    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 setup_vits():
+def download_file(url: str, dest_path: str, timeout: int = 30) -> bool:
-    os.makedirs(this_dir / 'models/vits', exist_ok=True)
+    """
-    os.makedirs(this_dir / 'models/xtts', exist_ok=True)
+    Download a file from a URL with progress bar and error handling.
-    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_xtts():
+    Args:
-    os.makedirs(this_dir / 'models/xtts/xttsv2_2.0.3', exist_ok=True)
+        url: URL to download from
-    base_url = "https://huggingface.co/coqui/XTTS-v2/resolve/v2.0.3/"
+        dest_path: Where to save the file
-    files = [
+        timeout: Connection timeout in seconds
-        "LICENSE.txt",
-        "README.md",
-        "config.json",
-        "model.pth",
-        "dvae.pth",
-        "mel_stats.pth",
-        "speakers_xtts.pth",
-        "vocab.json"
-    ]
-    for file in files:
-        download_file(base_url + file + "?download=true", this_dir / f"models/xtts/xttsv2_2.0.3/{file}")
     
+    Returns:
+        bool: True if download was successful
+    """
+    progress_bar = None
+    try:
+        response = requests.get(url, stream=True, timeout=timeout)
+        response.raise_for_status()
         
+        total_size = int(response.headers.get('content-length', 0))
+        block_size = 1024
         
-def update_tts_engines(engine):
+        progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True, 
-    AlltalkTTSEnginesConfig.get_instance().change_engine(engine).save()
+                          desc=f"Downloading {Path(dest_path).name}")
         
-def set_firstrun_model_false():
+        Path(dest_path).parent.mkdir(parents=True, exist_ok=True)
-    config.firstrun_model = False
-    config.save()
         
-def warning_message():
+        with open(dest_path, 'wb') as file:
+            for data in response.iter_content(block_size):
+                if data:
+                    file.write(data)
+                    progress_bar.update(len(data))
+        
+        if progress_bar.n != total_size and total_size > 0:
+            raise requests.exceptions.RequestException(
+                f"Incomplete download: {progress_bar.n}/{total_size} bytes"
+            )
+        return True
+        
+    except Exception as e:
+        print(f"Error downloading {url}: {str(e)}")
+        if Path(dest_path).exists():
+            Path(dest_path).unlink()
+        return False
+        
+    finally:
+        if progress_bar:
+            progress_bar.close()
+
+def setup_model(model_name: str) -> bool:
+    """
+    Set up a specific TTS model.
+    
+    Args:
+        model_name: Name of the model to setup ('piper', 'vits', or 'xtts')
+    
+    Returns:
+        bool: True if setup was successful
+    """
+    try:
+        if model_name == "piper":
+            os.makedirs(this_dir / 'models/piper', exist_ok=True)
+            os.makedirs(this_dir / 'models/xtts', exist_ok=True)
+            return all([
+                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"
+                )
+            ])
+            
+        elif model_name == "vits":
+            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'
+            
+            if not download_file(
+                "https://coqui.gateway.scarf.sh/v0.6.1_models/tts_models--en--vctk--vits.zip",
+                zip_path
+            ):
+                return False
+                
+            try:
+                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+                    zip_ref.extractall(this_dir / 'models/vits')
+                os.remove(zip_path)
+                return True
+            except Exception as e:
+                print(f"Error extracting VITS model: {str(e)}")
+                return False
+                
+        elif model_name == "xtts":
+            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"
+            ]
+            return all(
+                download_file(
+                    base_url + file + "?download=true",
+                    this_dir / f"models/xtts/xttsv2_2.0.3/{file}"
+                ) for file in files
+            )
+            
+        return False
+        
+    except Exception as e:
+        print(f"Error setting up {model_name}: {str(e)}")
+        return False
+
+def update_tts_engines(engine: str) -> bool:
+    """Update TTS engine configuration with error handling."""
+    try:
+        AlltalkTTSEnginesConfig.get_instance().change_engine(engine).save()
+        return True
+    except Exception as e:
+        print(f"Error updating TTS engine configuration: {str(e)}")
+        return False
+
+def set_firstrun_model_false() -> bool:
+    """Set firstrun_model to false with verification."""
+    try:
+        # Get fresh config instance and update
+        config = AlltalkConfig.get_instance()
+        config.firstrun_model = False
+        config.save()
+        
+        # Verify the save worked
+        config_path = AlltalkConfig.default_config_path()
+        with open(config_path, 'r') as f:
+            saved_config = json.load(f)
+            if saved_config.get('firstrun_model', True):
+                print(f"Warning: Failed to save firstrun_model=False")
+                # Try direct file write as fallback
+                saved_config['firstrun_model'] = False
+                with open(config_path, 'w') as f:
+                    json.dump(saved_config, f, indent=2)
+        return True
+        
+    except Exception as e:
+        print(f"Error saving firstrun configuration: {str(e)}")
+        return False
+
+def warning_message(branding: str):
+    """Display warning messages about installation and compatibility."""
     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[93mIf 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")
     print(f"[{branding}TTS] \033[93mOn Linux ignore the \033[0m'sparse_attn requires a torch version' \033[93mand\033[0m")
     print(f"[{branding}TTS] \033[0m'using untested triton version' \033[93mmessages.\033[0m")
     print(f"[{branding}TTS]")
 
-# Load confignew.json (might be version 1 or 2):
+def main():
-with open(AlltalkConfig.default_config_path(), 'r') as file:
+    """Main execution function."""
-    config_json = json.load(file)
-
-branding = config_json['branding']
-
-# Check if the config file is version 1
-if "ip_address" in config_json:
-    print(f"[{branding}TTS] Detected version 1 configuration file.")
-    print(f"[{branding}TTS] Upgrading to version 2 configuration file...")
-
-    # Backup the old configuration file
-    backup_path = AlltalkConfig.default_config_path() + ".backup"
-    shutil.copyfile(AlltalkConfig.default_config_path(), backup_path)
-
-    # Copy the new configuration template to replace the old config file
-    new_config_template_path = this_dir / 'system/config/confignew.json'
-    shutil.copyfile(new_config_template_path, AlltalkConfig.default_config_path())
-
-    print(f"[{branding}TTS] Configuration file upgraded successfully.")
-
-config = AlltalkConfig.get_instance()
-branding = config.branding
-
-# Check if firstrun_model is true
-if config.firstrun_model:
-    print(f"[{branding}TTS]")
-    print(f"[{branding}TTS] \033[92mThis is the first time startup. Please download a start TTS model.\033[0m")
-    print(f"[{branding}TTS]")
-    
-    # List of available models
-    models = [
-        {"name": "piper", "model": "piper"},
-        {"name": "vits", "model": "tts_models--en--vctk--vits"},
-        {"name": "xtts", "model": "xttsv2_2.0.3"},
-    ]
-    
-    # Display models to the user
-    print(f"[{branding}TTS]    \033[94mAvailable First Time Start-up models:\033[0m")
-    print(f"[{branding}TTS]")
-    for idx, model in enumerate(models):
-        print(f"[{branding}TTS]    \033[93m{idx + 1}. \033[94m{model['name']} - {model['model']}\033[0m")
-    print(f"[{branding}TTS]    \033[93m{len(models) + 1}.\033[94m I have my own models already\033[0m")
-    print(f"[{branding}TTS]")
-    print(f"[{branding}TTS]    \033[94mIn \033[91m60 seconds\033[0m \033[94ma Piper model will be \033[91mdownloaded automatically.\033[0m")
-    print(f"[{branding}TTS]")
-    
-    # Auto-select model after 30 seconds
-    selected_model = models[0]  # Default to the first model (piper)
     try:
-        user_choice = inputimeout(prompt=f"[{branding}TTS]    \033[92mEnter your choice 1-4: \033[0m ", timeout=60)
+        # Load configuration
-    except TimeoutOccurred:
+        with open(AlltalkConfig.default_config_path(), 'r') as file:
-        user_choice = None
+            config_json = json.load(file)
 
-    if user_choice is None:
+        branding = config_json['branding']
-        print(f"[{branding}TTS]")
-        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] \033[93m/models/xtts/\033[0m folder")
-        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)
-        update_tts_engines('piper')
-        set_firstrun_model_false()
-        warning_message()
-        exit()
 
-    if user_choice is not None and user_choice.isdigit() and 1 <= int(user_choice) <= len(models):
+        # Handle version 1 config upgrade
-        selected_model = models[int(user_choice) - 1]
+        if "ip_address" in config_json:
-        warning_message()
+            print(f"[{branding}TTS] Detected version 1 configuration file.")
+            print(f"[{branding}TTS] Upgrading to version 2 configuration file...")
+
+            backup_path = AlltalkConfig.default_config_path() + ".backup"
+            shutil.copyfile(AlltalkConfig.default_config_path(), backup_path)
+
+            new_config_template_path = this_dir / 'system/config/confignew.json'
+            shutil.copyfile(new_config_template_path, AlltalkConfig.default_config_path())
+
+            print(f"[{branding}TTS] Configuration file upgraded successfully.")
+
+        config = AlltalkConfig.get_instance()
+        branding = config.branding
+
+        if not config.firstrun_model:
+            return
+
+        # Auto-select Piper for Colab/Docker
+        if is_running_in_colab() or is_running_in_docker():
+            print(f"[{branding}TTS] Detected Colab/Docker environment - automatically selecting Piper")
+            if setup_model('piper'):
+                update_tts_engines('piper')
+                set_firstrun_model_false()
+                warning_message(branding)
+            return
+
+        # Normal interactive setup
+        models = [
+            {"name": "piper", "model": "piper"},
+            {"name": "vits", "model": "tts_models--en--vctk--vits"},
+            {"name": "xtts", "model": "xttsv2_2.0.3"},
+        ]
+
+        print(f"[{branding}TTS]")
+        print(f"[{branding}TTS] \033[92mThis is the first time startup. Please download a start TTS model.\033[0m")
+        print(f"[{branding}TTS]")
+        print(f"[{branding}TTS]    \033[94mAvailable First Time Start-up models:\033[0m")
+        print(f"[{branding}TTS]")
+        
+        for idx, model in enumerate(models):
+            print(f"[{branding}TTS]    \033[93m{idx + 1}. \033[94m{model['name']} - {model['model']}\033[0m")
+        print(f"[{branding}TTS]    \033[93m{len(models) + 1}.\033[94m I have my own models already\033[0m")
+        print(f"[{branding}TTS]")
+        print(f"[{branding}TTS]    \033[94mIn \033[91m60 seconds\033[0m \033[94ma Piper model will be \033[91mdownloaded automatically.\033[0m")
+        print(f"[{branding}TTS]")
+
+        try:
+            user_choice = inputimeout(
+                prompt=f"[{branding}TTS]    \033[92mEnter your choice 1-4: \033[0m ",
+                timeout=60
+            )
+        except TimeoutOccurred:
+            user_choice = None
+
+        if user_choice is None:
+            print(f"[{branding}TTS] No input received. Proceeding with the default model (piper).")
+            selected_model = models[0]
+        elif user_choice.lower() == 'own' or (user_choice.isdigit() and int(user_choice) == len(models) + 1):
+            print(f"[{branding}TTS] If you have your own XTTS models, please move them into the")
+            print(f"[{branding}TTS] \033[93m/models/xtts/\033[0m folder")
+            
+            os.makedirs('models/piper', exist_ok=True)
+            os.makedirs('models/vits', exist_ok=True)
+            os.makedirs('models/xtts/', exist_ok=True)
             
-    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':
+            set_firstrun_model_false()
-            setup_vits()
+            warning_message(branding)
-            update_tts_engines('vits')
+            return
-        elif selected_model['name'] == 'xtts':
+        elif user_choice.isdigit() and 1 <= int(user_choice) <= len(models):
-            setup_xtts()
+            selected_model = models[int(user_choice) - 1]
-            update_tts_engines('xtts')
+            warning_message(branding)
+        else:
+            print(f"[{branding}TTS] Invalid choice. Proceeding with the default model (piper).")
+            selected_model = models[0]
 
-        print(f"[{branding}TTS] {selected_model['name']} model downloaded and configuration updated successfully.")
+        if setup_model(selected_model['name']):
-        # Set firstrun_model to false
+            update_tts_engines(selected_model['name'])
-        set_firstrun_model_false()
+            set_firstrun_model_false()
-        warning_message()
+            print(f"[{branding}TTS] {selected_model['name']} model downloaded and configuration updated successfully.")
-        exit()
+            warning_message(branding)
-else:
+
-    exit()
+    except Exception as e:
+        print(f"Unexpected error: {str(e)}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()