# AUTOGENERATED! DO NOT EDIT! File to edit: app.ipynb.

# %% auto 0
__all__ = ['secret_import_failed', 'TEMP', 'TEMP_DIR', 'OPENAI_CLIENT_TTS_THREADS', 'CARTESIAAI_CLIENT_TTS_THREADS',
           'DEFAULT_PROVIDER', 'DEFAULT_MODEL', 'DEFAULT_VOICE', 'providers', 'clean_text_prompt', 'launch_kwargs',
           'queue_kwargs', 'verify_authorization', 'split_text', 'concatenate_audio', 'create_speech_openai',
           'create_speech_cartesiaai', 'create_speech', 'get_input_text_len', 'get_generation_cost',
           'get_model_choices', 'update_model_choices', 'get_voice_choices', 'update_voice_choices', 'split_text_as_md']

# %% app.ipynb 4
import os
secret_import_failed = False
try:
    # don't need the openai api key in a variable
    _ = os.environ['OPENAI_API_KEY']
    print('OPENAI_API_KEY environment variable was found.')
except:
    print('OPENAI_API_KEY environment variable was not found.')
    secret_import_failed = True
try:
    CARTESIA_API_KEY = os.environ['CARTESIA_API_KEY']
    print('CARTESIA_API_KEY environment variable was found.')
except:
    print('CARTESIA_API_KEY environment variable was not found.')
    secret_import_failed = True
try:
    temp_ALLOWED_OAUTH_PROFILE_USERNAMES = os.environ['ALLOWED_OAUTH_PROFILE_USERNAMES']
    ALLOWED_OAUTH_PROFILE_USERNAMES = tuple([o for o in temp_ALLOWED_OAUTH_PROFILE_USERNAMES.split(',') if o not in ('','None')])
    del temp_ALLOWED_OAUTH_PROFILE_USERNAMES
    print(f'ALLOWED_OAUTH_PROFILE_USERNAMES environment variable was found. {ALLOWED_OAUTH_PROFILE_USERNAMES}')
except:
    print('ALLOWED_OAUTH_PROFILE_USERNAMES environment variable was not found.')
    secret_import_failed = True

if secret_import_failed == True:
    import tts_openai_secrets
    _ = os.environ['OPENAI_API_KEY']
    CARTESIA_API_KEY = os.environ['CARTESIA_API_KEY']
    ALLOWED_OAUTH_PROFILE_USERNAMES = os.environ['ALLOWED_OAUTH_PROFILE_USERNAMES']
    print('import tts_openai_secrets succeeded')

# %% app.ipynb 5
# If REQUIRE_AUTH environemnt variable is set to 'false' (from secrets) and HF_SPACE != 1 then we
# are running locally and don't require authentication and authorization, otherwise we do.
# We are using paid API's so don't want anybody/everybody to be able to use our paid services.
if os.environ.get("REQUIRE_AUTH",'true') == 'false' and os.environ.get('HF_SPACE',0) != 1:
    REQUIRE_AUTH = False
else:
    REQUIRE_AUTH = True
print('REQUIRE_AUTH:',REQUIRE_AUTH)

# %% app.ipynb 8
import os
import gradio as gr
import openai
from pydub import AudioSegment
import io
from datetime import datetime
from math import ceil
from multiprocessing.pool import ThreadPool
from functools import partial
from pathlib import Path
import uuid
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)  # for exponential backoff
import traceback
# from cartesia.tts import CartesiaTTS
import cartesia

# %% app.ipynb 11
TEMP = os.environ.get('GRADIO_TEMP_DIR','/tmp/')
TEMP_DIR = Path(TEMP)
print('TEMP Dir:', TEMP_DIR)

# %% app.ipynb 12
#Number of threads created PER USER REQUEST. This throttels the # of API requests PER USER request. This is in ADDITION to the Gradio threads.
OPENAI_CLIENT_TTS_THREADS = 10 
CARTESIAAI_CLIENT_TTS_THREADS = 3

DEFAULT_PROVIDER = 'openai'
DEFAULT_MODEL = 'tts-1'
DEFAULT_VOICE = 'alloy'

# %% app.ipynb 13
providers = dict()

# %% app.ipynb 14
# Add OpenAI as a provider
try:
    providers['openai'] = {
        'name': 'Open AI',
        'models': {o.id: o.id for o in openai.models.list().data if 'tts' in o.id},
        'voices': {o:{'id':o,'name':o.title()} for o in ['alloy', 'echo', 'fable', 'onyx', 'nova', 'shimmer']},
        'settings': {'max_chunk_size': 4000, 'chunk_processing_time': 60, 
                     'threads': OPENAI_CLIENT_TTS_THREADS,
                     'audio_file_conversion_kwargs':{'format': 'mp3'}},
    }
    print('Successfully added OpenAI as Provider')
except Exception as e:
    print(f"""Error: Failed to add OpenAI as a provider.\nException: {repr(e)}\nTRACEBACK:\n""",traceback.format_exc())
# providers

# %% app.ipynb 15
# Add Cartesia AI as a provider
try:
    providers['cartesiaai'] = {
        'name': 'Cartesia AI',
        'models': {'upbeat-moon': 'Sonic Turbo English'},
        'voices': {v['id']:v for k,v in cartesia.tts.CartesiaTTS().get_voices().items()},
        'settings': {'max_chunk_size': 500, 'chunk_processing_time': 20, 
                     'threads': CARTESIAAI_CLIENT_TTS_THREADS,
                     'audio_file_conversion_kwargs':{'format': 'raw', 'frame_rate': 44100, 
                                                     'channels': 1, 'sample_width': 2}},
    }
    print('Successfully added Cartesia AI as Provider')
except Exception as e:
    print(f"""Error: Failed to add Cartesia AI as a provider.\nException: {repr(e)}\nTRACEBACK:\n""",traceback.format_exc())
# providers

# %% app.ipynb 19
clean_text_prompt = """Your job is to clean up text that is going to be fed into a text to speech (TTS) model. You must remove parts of the text that would not normally be spoken such as reference marks `[1]`, spurious citations such as `(Reddy et al., 2021; Wu et al., 2022; Chang et al., 2022; Kondratyuk et al., 2023)` and any other part of the text that is not normally spoken. Please also clean up sections and headers so they are on new lines with proper numbering. You must also clean up any math formulas that are salvageable from being copied from a scientific paper. If they are garbled and do not make sense then remove them. You must carefully perform the text cleanup so it is translated into speech that is easy to listen to however you must not modify the text otherwise. It is critical that you repeat all of the text without modifications except for the cleanup activities you've been instructed to do. Also you must clean all of the text you are given, you may not omit any of it or stop the cleanup task early."""


# %% app.ipynb 21
def verify_authorization(profile: gr.OAuthProfile=None) -> str:
    print('Profile:', profile)
    if REQUIRE_AUTH == False:
        return 'WARNING_NO_AUTH_REQUIRED_LOCAL'
    elif profile is not None and profile.username in ALLOWED_OAUTH_PROFILE_USERNAMES:
        return f"{profile.username}"
    else:
        # print('Unauthorized',profile)
        raise PermissionError(f'Your huggingface username ({profile}) is not authorized. Must be set in ALLOWED_OAUTH_PROFILE_USERNAMES environment variable.')
        return None

# %% app.ipynb 22
def split_text(input_text, provider):
    settings = providers[provider]['settings']
    max_length = settings['max_chunk_size']
    lookback = max_length // 4
        
    # If the text is shorter than the max_length, return it as is
    if len(input_text) <= max_length:
        return [input_text]

    chunks = []
    while input_text:
        # Check if the remaining text is shorter than the max_length
        if len(input_text) <= max_length:
            chunks.append(input_text)
            break

        # Define the split point, initially set to max_length
        split_point = max_length

        # Look for a newline in the last 'lookback' characters
        newline_index = input_text.rfind('\n', max_length-lookback, max_length)
        if newline_index != -1:
            split_point = newline_index + 1  # Include the newline in the current chunk

        # If no newline, look for a period followed by space
        elif '. ' in input_text[max_length-lookback:max_length]:
            # Find the last '. ' in the lookback range
            period_index = input_text.rfind('. ', max_length-lookback, max_length)
            split_point = period_index + 2  # Split after the space

        # Split the text and update the input_text
        chunks.append(input_text[:split_point])
        input_text = input_text[split_point:]

    return chunks

# %% app.ipynb 23
def concatenate_audio(files:list, **kwargs):

    # Initialize an empty AudioSegment object for concatenation
    combined = AudioSegment.empty()

    # Loop through the list of mp3 binary data
    for data in files:
        # Convert binary data to an audio segment
        audio_segment = AudioSegment.from_file(io.BytesIO(data), **kwargs)
        
        # Concatenate this segment to the combined segment
        combined += audio_segment

    #### Return Bytes Method
    # # Export the combined segment to a new mp3 file
    # # Use a BytesIO object to handle this in memory
    # combined_mp3 = io.BytesIO()
    # combined.export(combined_mp3, format="mp3")

    # # Seek to the start so it's ready for reading
    # combined_mp3.seek(0)

    # return combined_mp3.getvalue()

    #### Return Filepath Method
    filepath = TEMP_DIR/(str(uuid.uuid4())+'.mp3')
    combined.export(filepath, format="mp3")
    print('Saving mp3 file to temp directory: ', filepath)
    return str(filepath)

# %% app.ipynb 24
def create_speech_openai(chunk_idx, input, model='tts-1', voice='alloy', speed=1.0, **kwargs):
    client = openai.OpenAI()
    
    @retry(wait=wait_random_exponential(min=1, max=180), stop=stop_after_attempt(6))
    def _create_speech_with_backoff(**kwargs):
        return client.audio.speech.create(**kwargs)
    
    response = _create_speech_with_backoff(input=input, model=model, voice=voice, speed=speed, **kwargs)
    client.close()
    return chunk_idx, response.content
if 'openai' in providers.keys():
    providers['openai']['settings']['create_speech_func'] = create_speech_openai
    print('Added create_speech_func for openai provider')

# %% app.ipynb 26
def create_speech_cartesiaai(chunk_idx, input, model='upbeat-moon', 
                             voice='248be419-c632-4f23-adf1-5324ed7dbf1d', #Hannah
                             websocket=False, 
                             output_format='pcm_44100', 
                             **kwargs):
    client = cartesia.tts.CartesiaTTS()
    
    # @retry(wait=wait_random_exponential(min=1, max=180), stop=stop_after_attempt(6))
    def _create_speech_with_backoff(**kwargs):
        return client.generate(**kwargs)
    
    response = _create_speech_with_backoff(transcript=input, model_id=model, 
                                           voice=client.get_voice_embedding(voice_id=voice), 
                                           websocket=websocket, 
                                           output_format=output_format, 
                                           **kwargs)
    client.close()
    return chunk_idx, response["audio"]
if 'cartesiaai' in providers.keys():
    providers['cartesiaai']['settings']['create_speech_func'] = create_speech_cartesiaai
    print('Added create_speech_func for create_speech_cartesiaai provider')

# %% app.ipynb 29
def create_speech(input_text, provider, model='tts-1', voice='alloy', 
                  profile: gr.OAuthProfile|None=None, # comment out of running locally
                  progress=gr.Progress(), **kwargs):

    #Verify auth if it is required. This is very important if this is in a HF space. DO NOT DELETE!!!
    if REQUIRE_AUTH: verify_authorization(profile)
    start = datetime.now()
    
    settings = providers[provider]['settings']
    create_speech_func = settings['create_speech_func']
    chunk_processing_time = settings['chunk_processing_time']
    threads = settings['threads']
    audio_file_conversion_kwargs = settings['audio_file_conversion_kwargs']
    
    # Split the input text into chunks
    chunks = split_text(input_text, provider=provider)

    # Initialize the progress bar
    progress(0, desc=f"Started processing {len(chunks)} text chunks using {threads} threads. ETA is ~{ceil(len(chunks)/threads)*chunk_processing_time/60.} min.")

    # Initialize a list to hold the audio data of each chunk
    audio_data = []

    # Process each chunk
    with ThreadPool(processes=threads) as pool:
        results = pool.starmap(
            partial(create_speech_func, model=model, voice=voice, **kwargs), 
            zip(range(len(chunks)),chunks)
        )
    audio_data = [o[1] for o in sorted(results)]

    # Progress
    progress(.9, desc=f"Merging audio chunks... {(datetime.now()-start).seconds} seconds to process.")
    
    # Concatenate the audio data from all chunks
    combined_audio = concatenate_audio(audio_data, **audio_file_conversion_kwargs)

    # Final update to the progress bar
    progress(1, desc=f"Processing completed... {(datetime.now()-start).seconds} seconds to process.")
    
    print(f"Processing time: {(datetime.now()-start).seconds} seconds.")

    return combined_audio


# %% app.ipynb 31
def get_input_text_len(input_text):
    return len(input_text)

# %% app.ipynb 32
def get_generation_cost(input_text, tts_model_dropdown, provider):
    text_len = len(input_text)
    if provider == 'openai':
        if tts_model_dropdown.endswith('-hd'):
            cost = text_len/1000 * 0.03
        else:
            cost = text_len/1000 * 0.015
    elif provider == 'cartesiaai':
        cost = text_len/1000 * 0.065
    else:
        raise ValueError(f'Invalid argument provider: {provider}')
    return "${:,.3f}".format(cost)

# %% app.ipynb 33
def get_model_choices(provider):
    return sorted([(v,k) for k,v in providers[provider]['models'].items()])

# %% app.ipynb 34
def update_model_choices(provider):
    choices = get_model_choices(provider)
    return gr.update(choices=choices,value=choices[0][1])

# %% app.ipynb 35
def get_voice_choices(provider, model):
    return sorted([(v['name'],v['id']) for v in providers[provider]['voices'].values()])

# %% app.ipynb 36
def update_voice_choices(provider, model):
    choices = get_voice_choices(provider, model)
    return gr.update(choices=choices,value=choices[0][1])

# %% app.ipynb 37
def split_text_as_md(*args, **kwargs):
    output = split_text(*args, **kwargs)
    return '# Text Splits:\n' + '<br>----------<br>'.join(output)

# %% app.ipynb 38
with gr.Blocks(title='TTS', head='TTS', delete_cache=(3600,3600)) as app:
    
    ### Define UI ###
    gr.Markdown("# TTS")
    gr.Markdown("""Start typing below and then click **Go** to create the speech from your text.
For requests longer than allowed by the API they will be broken into chunks automatically. [Spaces Link](https://matdmiller-tts-openai.hf.space/) | <a href="https://matdmiller-tts-openai.hf.space/" target="_blank">Spaces Link HTML</a>""")
    with gr.Row():
        input_text = gr.Textbox(max_lines=100, label="Enter text here")
    with gr.Row():
        tts_provider_dropdown = gr.Dropdown(value=DEFAULT_PROVIDER,
            choices=tuple([(v['name'],k) for k,v in providers.items()]), label='Provider', interactive=True)
        tts_model_dropdown = gr.Dropdown(value=DEFAULT_MODEL,choices=get_model_choices(DEFAULT_PROVIDER), 
                                         label='Model', interactive=True)
        tts_voice_dropdown = gr.Dropdown(value=DEFAULT_VOICE,choices=get_voice_choices(DEFAULT_PROVIDER, DEFAULT_MODEL),
                                         label='Voice', interactive=True)
        input_text_length = gr.Label(label="Number of characters")
        generation_cost = gr.Label(label="Generation cost")
    with gr.Row():
        output_audio = gr.Audio()
    go_btn = gr.Button("Go")
    clear_btn = gr.Button('Clear')
    if REQUIRE_AUTH:
        gr.LoginButton()
        auth_md = gr.Markdown('')
        
    chunks_md = gr.Markdown('',label='Chunks')
    

    ### Define UI Actions ###
    
    # input_text 
    input_text.input(fn=get_input_text_len, inputs=input_text, outputs=input_text_length)
    input_text.input(fn=get_generation_cost, 
                     inputs=[input_text,tts_model_dropdown,tts_provider_dropdown], 
                     outputs=generation_cost)
    input_text.input(fn=split_text_as_md, inputs=[input_text,tts_provider_dropdown], outputs=chunks_md)
    
    # tts_provider_dropdown
    tts_provider_dropdown.change(fn=update_model_choices, inputs=[tts_provider_dropdown], 
                                 outputs=tts_model_dropdown)
    tts_provider_dropdown.change(fn=update_voice_choices, inputs=[tts_provider_dropdown, tts_model_dropdown], 
                                 outputs=tts_voice_dropdown)
    tts_provider_dropdown.change(fn=split_text_as_md, inputs=[input_text,tts_provider_dropdown], outputs=chunks_md)
    
    # tts_model_dropdown
    tts_model_dropdown.change(fn=get_generation_cost, 
                              inputs=[input_text,tts_model_dropdown,tts_provider_dropdown], outputs=generation_cost)
    
    
    go_btn.click(fn=create_speech, 
                 inputs=[input_text, tts_provider_dropdown, tts_model_dropdown, tts_voice_dropdown], 
                 outputs=[output_audio])
    
    
    clear_btn.click(fn=lambda: '', outputs=input_text)

    if REQUIRE_AUTH:
        app.load(verify_authorization, None, auth_md)
    
    

# %% app.ipynb 39
# launch_kwargs = {'auth':('username',GRADIO_PASSWORD),
#                  'auth_message':'Please log in to Mat\'s TTS App with username: username and password.'}
launch_kwargs = {}
queue_kwargs = {'default_concurrency_limit':10}

# %% app.ipynb 41
#.py launch
if __name__ == "__main__":
    app.queue(**queue_kwargs)
    app.launch(**launch_kwargs)