import spaces
import gradio as gr
import torch
from string import punctuation
import re


from parler_tts import ParlerTTSForConditionalGeneration
from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed

device = "cuda:0" if torch.cuda.is_available() else "cpu"

custom_repo_id = "AkhilTolani/parler-tts-music-200000"

model = ParlerTTSForConditionalGeneration.from_pretrained(custom_repo_id).to(device)
tokenizer = AutoTokenizer.from_pretrained("AkhilTolani/parler-tts-music-200000")
feature_extractor = AutoFeatureExtractor.from_pretrained("AkhilTolani/parler-tts-music-200000")

SAMPLE_RATE = feature_extractor.sampling_rate

SEED = 456

default_text = "i cant even find them getting your eyes with i can only think about checking your eyes your eyes your eyes"
default_description = "This hip hop track showcases a passionate male vocal with harmonizing background vocals, layered over shimmering cymbals, punchy kick and snare hits, synth pad, and wooden percussions. The song exudes a passionate, emotional, and groovy vibe that is sure to captivate listeners. This track would be perfect for setting the mood in a trendy urban nightclub or a stylish lounge setting."

examples = [
    [
        "i cant even find them getting your eyes with i can only think about checking your eyes your eyes your eyes",
        "This hip hop track showcases a passionate male vocal with harmonizing background vocals, layered over shimmering cymbals, punchy kick and snare hits, synth pad, and wooden percussions. The song exudes a passionate, emotional, and groovy vibe that is sure to captivate listeners. This track would be perfect for setting the mood in a trendy urban nightclub or a stylish lounge setting."
    ]
]

@spaces.GPU
def gen_tts(text, description):
    inputs = tokenizer(description, return_tensors="pt").to(device)
    prompt = tokenizer(text, return_tensors="pt").to(device)

    set_seed(SEED)
    generation = model.generate(input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids, min_length=20).to(torch.float32)
    audio_arr = generation.cpu().numpy().squeeze()

    return SAMPLE_RATE, audio_arr


css = """
        #share-btn-container {
            display: flex;
            padding-left: 0.5rem !important;
            padding-right: 0.5rem !important;
            background-color: #000000;
            justify-content: center;
            align-items: center;
            border-radius: 9999px !important; 
            width: 13rem;
            margin-top: 10px;
            margin-left: auto;
            flex: unset !important;
        }
        #share-btn {
            all: initial;
            color: #ffffff;
            font-weight: 600;
            cursor: pointer;
            font-family: 'IBM Plex Sans', sans-serif;
            margin-left: 0.5rem !important;
            padding-top: 0.25rem !important;
            padding-bottom: 0.25rem !important;
            right:0;
        }
        #share-btn * {
            all: unset !important;
        }
        #share-btn-container div:nth-child(-n+2){
            width: auto !important;
            min-height: 0px !important;
        }
        #share-btn-container .wrap {
            display: none !important;
        }
"""
with gr.Blocks(css=css) as block:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
              <div
                style="
                  display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;
                "
              >
                <h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
                  Parler-TTS + Vocals 🥰
                </h1>
              </div>
            </div>
        """
    )
    gr.HTML(
        f"""
        <p><a href="https://github.com/huggingface/parler-tts"> Parler-TTS + Vocals</a> is a finetuned model for generating high-quality vocals with features that can be controlled using a simple text prompt (e.g. gender, background noise, speaking rate, pitch and reverberation).</p>

        <p>Tips for ensuring good generation:
        <ul>
            <li>Include the term <b>"very clear audio"</b> to generate the highest quality audio, and "very noisy audio" for high levels of background noise</li>
            <li>Punctuation can be used to control the prosody of the generations, e.g. use commas to add small breaks in speech</li>
            <li>The remaining speech features (gender, speaking rate, pitch and reverberation) can be controlled directly through the prompt</li>
        </ul>
        </p>
        """
    )
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Input Text", lines=2, value=default_text, elem_id="input_text")
            description = gr.Textbox(label="Description", lines=2, value=default_description, elem_id="input_description")
            run_button = gr.Button("Generate Audio", variant="primary")
        with gr.Column():
            audio_out = gr.Audio(label="Parler-TTS + Vocals", type="numpy", elem_id="audio_out")

    inputs = [input_text, description]
    outputs = [audio_out]
    gr.Examples(examples=examples, fn=gen_tts, inputs=inputs, outputs=outputs, cache_examples=False)
    run_button.click(fn=gen_tts, inputs=inputs, outputs=outputs, queue=True)

block.queue()
block.launch(share=True)