services/audio_model.py

import os
import requests
from tqdm import tqdm
from google.cloud import storage


def _download_starting_files() -> None:
    """
    Downloads the embeddings from a bucket
    """
    # Initialise a client
    credentials = os.getenv('GOOGLE_APPLICATION_CREDENTIALS')
    storage_client = storage.Client.from_service_account_json(credentials)
    bucket = storage_client.get_bucket('embeddings-bella')

    # Get both embeddings
    blob = bucket.blob("gpt_cond_latent.npy")
    blob.download_to_filename('assets/gpt_cond_latent.npy')
    blob = bucket.blob("speaker_embedding.npy")
    blob.download_to_filename('assets/speaker_embedding.npy')


def _download_file(url, destination):
    response = requests.get(url, stream=True)
    total_size_in_bytes = int(response.headers.get('content-length', 0))
    block_size = 1024

    progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)

    with open(destination, 'wb') as file:
        for data in response.iter_content(block_size):
            progress_bar.update(len(data))
            file.write(data)

    progress_bar.close()


def download_model():
    # Define files and their corresponding URLs
    files_to_download = {
        'LICENSE.txt': 'https://huggingface.co/coqui/XTTS-v2/resolve/v2.0.2/LICENSE.txt?download=true',
        'README.md': 'https://huggingface.co/coqui/XTTS-v2/resolve/v2.0.2/README.md?download=true',
        'config.json': 'https://huggingface.co/coqui/XTTS-v2/resolve/v2.0.2/config.json?download=true',
        'model.pth': 'https://huggingface.co/coqui/XTTS-v2/resolve/v2.0.2/model.pth?download=true',
        'vocab.json': 'https://huggingface.co/coqui/XTTS-v2/resolve/v2.0.2/vocab.json?download=true',
    }

    if not os.path.exists("tts_model"):
        os.makedirs("tts_model")

    # Download files if they don't exist
    print("[COQUI TTS] STARTUP: Checking Model is Downloaded.")
    for filename, url in files_to_download.items():
        destination = f'tts_model/{filename}'
        print(f"[COQUI TTS] STARTUP: Downloading {filename}...")
        _download_file(url, destination)

    # Downloads the embeddings from GCP
    _download_starting_files()