audiobook/backend/download_piper_voices.py

#!/usr/bin/env python3
"""
Simple script to download Piper TTS voice models.
"""
import os
import sys
import urllib.request
import ssl

# Create SSL context that doesn't verify (for development)
ssl_context = ssl.create_default_context()
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE

# Models directory
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(SCRIPT_DIR, "piper_models")
os.makedirs(MODELS_DIR, exist_ok=True)

# Available voices with direct download links
VOICES = {
    "1": {
        "name": "en_US-lessac-medium",
        "description": "Male, American, Medium Quality",
        "model_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/en_US-lessac-medium.onnx",
        "config_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/en_US-lessac-medium.onnx.json",
    },
    "2": {
        "name": "en_US-amy-medium",
        "description": "Female, American, Medium Quality",
        "model_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/en_US-amy-medium.onnx",
        "config_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/en_US-amy-medium.onnx.json",
    },
    "3": {
        "name": "en_US-libritts-high",
        "description": "Neutral, American, High Quality",
        "model_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/en_US-libritts-high.onnx",
        "config_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_US/en_US-libritts-high.onnx.json",
    },
    "4": {
        "name": "en_GB-alan-medium",
        "description": "Male, British, Medium Quality",
        "model_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_GB/en_GB-alan-medium.onnx",
        "config_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_GB/en_GB-alan-medium.onnx.json",
    },
    "5": {
        "name": "en_GB-alba-medium",
        "description": "Female, British, Medium Quality",
        "model_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_GB/en_GB-alba-medium.onnx",
        "config_url": "https://huggingface.co/rhasspy/piper-voices/resolve/v1.0.0/en/en_GB/en_GB-alba-medium.onnx.json",
    },
}


def download_file(url, dest_path):
    """Download a file with progress indication."""
    try:
        print(f"    Downloading from: {url}")
        
        # Use urllib with SSL context
        with urllib.request.urlopen(url, context=ssl_context, timeout=120) as response:
            total_size = int(response.headers.get('content-length', 0))
            
            with open(dest_path, 'wb') as f:
                downloaded = 0
                chunk_size = 8192
                
                while True:
                    chunk = response.read(chunk_size)
                    if not chunk:
                        break
                    
                    f.write(chunk)
                    downloaded += len(chunk)
                    
                    if total_size > 0:
                        percent = (downloaded / total_size) * 100
                        print(f"\r    Progress: {percent:.1f}% ({downloaded}/{total_size} bytes)", end='')
        
        print()  # New line after progress
        
        # Verify file size
        file_size = os.path.getsize(dest_path)
        if file_size < 1000:
            print(f"    ✗ Download failed (file too small: {file_size} bytes)")
            os.remove(dest_path)
            return False
        
        print(f"    ✓ Downloaded: {file_size:,} bytes")
        return True
        
    except Exception as e:
        print(f"    ✗ Error: {e}")
        if os.path.exists(dest_path):
            os.remove(dest_path)
        return False


def download_voice(voice_key):
    """Download a voice model and config."""
    voice = VOICES[voice_key]
    name = voice["name"]
    
    model_path = os.path.join(MODELS_DIR, f"{name}.onnx")
    config_path = os.path.join(MODELS_DIR, f"{name}.onnx.json")
    
    # Check if already exists
    if os.path.exists(model_path) and os.path.exists(config_path):
        print(f"  ✓ {name} - Already downloaded")
        return True
    
    print(f"\n  Downloading: {name}")
    print(f"  Description: {voice['description']}")
    
    # Download model
    print("  [1/2] Model file (.onnx)...")
    if not download_file(voice["model_url"], model_path):
        return False
    
    # Download config
    print("  [2/2] Config file (.onnx.json)...")
    if not download_file(voice["config_url"], config_path):
        return False
    
    print(f"  ✓ {name} - Complete!")
    return True


def main():
    print("=" * 60)
    print("Piper TTS Voice Model Downloader")
    print("=" * 60)
    print(f"\nModels directory: {MODELS_DIR}\n")
    
    print("Available voices:\n")
    for key, voice in VOICES.items():
        print(f"  {key}) {voice['name']}")
        print(f"     {voice['description']}")
        print()
    
    print("Quick options:")
    print("  a) Download all voices (recommended)")
    print("  s) Download starter pack (1, 2, 4)")
    print("  q) Quit")
    print()
    
    choice = input("Enter your choice (1-5, a, s, or q): ").strip().lower()
    
    if choice == 'q':
        print("Exiting...")
        return
    
    print("\n" + "=" * 60)
    print("Downloading...")
    print("=" * 60)
    
    success_count = 0
    total_count = 0
    
    if choice == 'a':
        # Download all
        for key in VOICES.keys():
            total_count += 1
            if download_voice(key):
                success_count += 1
    elif choice == 's':
        # Starter pack
        for key in ['1', '2', '4']:
            total_count += 1
            if download_voice(key):
                success_count += 1
    elif choice in VOICES:
        # Single voice
        total_count += 1
        if download_voice(choice):
            success_count += 1
    else:
        print(f"Invalid choice: {choice}")
        return
    
    print("\n" + "=" * 60)
    print("Download Summary")
    print("=" * 60)
    print(f"\nSuccessfully downloaded: {success_count}/{total_count} voices")
    
    # List downloaded files
    print(f"\nFiles in {MODELS_DIR}:")
    try:
        files = sorted(os.listdir(MODELS_DIR))
        if files:
            for f in files:
                size = os.path.getsize(os.path.join(MODELS_DIR, f))
                print(f"  {f} ({size:,} bytes)")
        else:
            print("  (empty)")
    except Exception as e:
        print(f"  Error listing files: {e}")
    
    print("\n" + "=" * 60)
    print("Next Steps")
    print("=" * 60)
    print("\n1. Test Piper installation:")
    print("   python test_piper_direct.py")
    print("\n2. Start the backend:")
    print("   uvicorn main:app --reload")
    print("\n3. In the UI, click 'Piper TTS' button")
    print("\n4. Select a Piper voice and generate!")
    print()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\nDownload cancelled by user.")
        sys.exit(1)
    except Exception as e:
        print(f"\n\nError: {e}")
        sys.exit(1)