# stt_gradio.py
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import gradio as gr

# -------------------
# 1️⃣ Detect GPU
# -------------------
use_cuda = torch.cuda.is_available()
dtype = torch.float16 if use_cuda else torch.float32

print(f"🌟 Using {'GPU' if use_cuda else 'CPU'}, dtype={dtype}")

# -------------------
# 2️⃣ Load Whisper model
# -------------------
hub_id = "Muhammadidrees/WispherVOICE"

print("⏳ Loading model...")
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    hub_id,
    torch_dtype=dtype,
    device_map="auto",          # accelerate handles device placement
    trust_remote_code=True
)

processor = AutoProcessor.from_pretrained(
    hub_id,
    trust_remote_code=True
)

# -------------------
# 3️⃣ Create pipeline (no device argument!)
# -------------------
pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor
)

print("🎧 Whisper pipeline ready.")

# -------------------
# 4️⃣ Transcription Function
# -------------------
def transcribe(audio):
    # Gradio audio input returns a file path
    if audio is None:
        return "No audio provided."
    result = pipe(audio)
    return result["text"]

# -------------------
# 5️⃣ Gradio Interface
# -------------------
demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(sources=["microphone", "upload"], type="filepath"),
    outputs="text",
    title="🎤 Whisper Speech-to-Text",
    description="Record or upload audio and get real-time transcription using Whisper."
)

if __name__ == "__main__":
    demo.launch()