# 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
DEFAULT_TASK = "transcribe"
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
}
'''

# Functions
@spaces.GPU(duration=10)
def transcribe(input=None, task=DEFAULT_TASK):
    print(input)
    if input is None: raise gr.Error("Invalid input.")
    output = repo(input, 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")
        task = gr.Radio(["transcribe", "translate"], label="Task", value=DEFAULT_TASK)
        submit = gr.Button("▶")
        maintain = gr.Button("☁️")

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

main.launch(show_api=True)