# @title Notebook Setup
# @markdown Please, run this cell first. You'll be prompted for your OpenAI API Keys.
# @markdown https://platform.openai.com/api-keys
!pip install openai tiktoken > /dev/null 2>&1

from getpass import getpass
from openai import OpenAI

# Prompt for API key and create a client
api_key = getpass('OpenAI API Key: ')
client = OpenAI(api_key=api_key)

# @markdown This cell can be run multiple times to generate audio samples. All results will be saved to your Google Drive in the OpenAI_TTS_Samples folder.

import io
import os
import time

from IPython.display import Audio, display
from google.colab import drive

# Mount Google Drive to /content/gdrive path.
drive.mount('/content/gdrive')

# Output folder is OpenAI_TTS_Samples, it will be created
# if it does not exist.
out_path = "/content/gdrive/MyDrive/OpenAI_TTS_Samples/"
try:
    os.mkdir(out_path)
except FileExistsError:
   pass

voices = ("alloy", "echo", "fable", "onyx", "nova", "shimmer")
input="""
Today is a wonderful day to learn Vim!
"""

# Loop through voices and generate a sample for each voice
for voice in voices:
  response = client.audio.speech.create(
    model="tts-1-hd",
    voice=voice,
    input=input
  )

  # Prepare file name - YYYY-MM-DD-HH-MM-SS_voice.mp3
  timestr = time.strftime("%Y-%m-%d-%H-%M-%S")
  speech_file_name = f"{timestr}_{voice}.mp3"
  speech_file_path = out_path + speech_file_name

  # Save generated audio to the file.
  response.stream_to_file(speech_file_path)

  # Display files in the notebook as <audio> element
  display(speech_file_name, Audio(speech_file_path, autoplay=False))