import gradio as gr
import spaces
import torch
import json
from transformers import pipeline

@spaces.GPU(duration=30)
def transcribe_audio(audio_path):
    """Транскрибирует аудио с word timestamps"""
    try:
        pipe = pipeline(
            "automatic-speech-recognition",
            model="openai/whisper-large-v3",
            torch_dtype=torch.float16,
            device="cuda:0",
        )

        result = pipe(audio_path, return_timestamps="word")

        text = result["text"]
        chunks = result.get("chunks", [])

        formatted = f"📝 Текст:\n{text}\n\n⏱️ Слова:\n"
        for chunk in chunks[:20]:
            word = chunk["text"]
            start = chunk["timestamp"][0] or 0
            end = chunk["timestamp"][1] or 0
            formatted += f"{start:06.2f}s - {end:06.2f}s: {word}\n"

        return formatted, json.dumps(result, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ Ошибка: {str(e)}", f'{{"error": "{str(e)}"}}'

demo = gr.Interface(
    fn=transcribe_audio,
    inputs=gr.Audio(type="filepath", label="Аудио"),
    outputs=[
        gr.Textbox(label="Результат", lines=20),
        gr.JSON(label="JSON")
    ],
    title="🎤 Whisper Test",
)

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