# Imports
import gradio as gr
import spaces
import torch

from transformers import pipeline

# Pre-Initialize
DEVICE = "auto"
if DEVICE == "auto":
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[SYSTEM] | Using {DEVICE} type compute device.")

# Variables
BATCH_SIZE = 8

repo = pipeline(task="automatic-speech-recognition", model="openai/whisper-large-v3-turbo", chunk_length_s=30, device=DEVICE)

css = '''
.gradio-container{max-width: 560px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''

@spaces.GPU(15)
def transcribe(inputs, task):
    if inputs is None: raise gr.Error("Invalid input.")
    output = repo(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
    return output

def cloud():
    print("[CLOUD] | Space maintained.")

# Initialize
with gr.Blocks(css=css) as main:
    with gr.Column():
        gr.Markdown("🪄 Transcribe audio to text.")
        
    with gr.Column():
        input = gr.Audio(sources="upload", type="filepath", label="Input"),
        type = gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
        submit = gr.Button("▶")
        maintain = gr.Button("☁️")

    with gr.Column():
        output = gr.Textbox(lines=1, value=DEFAULT_INPUT, label="Output")
            
    submit.click(transcribe, inputs=[input, type], outputs=[output], queue=False)
    maintain.click(cloud, inputs=[], outputs=[], queue=False)

main.launch(show_api=True)