app.py

# -*- coding: utf-8 -*-
import torch
import gradio as gr
import yt_dlp as youtube_dl
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
from huggingface_hub import login

import tempfile
import uuid
import os
import time
import yt_dlp

# --- Конфигурация ---
HF_TOKEN = os.getenv("HF_ACCESS_TOKEN")
MODEL_NAME = "artyomboyko/whisper-small-ru-v4"
BATCH_SIZE = 8
# FILE_LIMIT_MB = 1000

# --- Инициализация ---
try:
    if HF_TOKEN:
        login(token=HF_TOKEN)
        print("Успешный вход в Hugging Face Hub.")
    else:
        print("HF_ACCESS_TOKEN не найден. Запуск в оффлайн-режиме или с публичными моделями.")
except Exception as e:
    print(f"Не удалось войти в Hugging Face Hub: {e}")

# Определение устройства для вычислений (GPU или CPU)
device = 0 if torch.cuda.is_available() else "cpu"
if device == 0:
    print("GPU доступен, вычисления будут производиться на GPU.")
else:
    print("GPU не доступен, вычисления будут производиться на CPU.")

# Загрузка модели распознавания речи
try:
    pipe = pipeline(
        task="automatic-speech-recognition",
        model=MODEL_NAME,
        chunk_length_s=30,
        device=device,
    )
    print(f"Модель {MODEL_NAME} успешно загружена.")
except Exception as e:
    print(f"Ошибка при загрузке модели {MODEL_NAME}: {e}")
    exit()


# --- Функции ---
def transcribe(filepath, task):
    """
    Транскрибирует аудиофайл с помощью загруженной модели Whisper.

    Args:
        filepath (str): Путь к аудиофайлу.
        task (str): Задача для модели ("transcribe" или "translate").

    Returns:
        str: Распознанный текст.
    """
    if filepath is None:
        # Используем gr.Warning вместо gr.Error для неблокирующего уведомления
        gr.Warning("Аудиофайл не предоставлен! Пожалуйста, загрузите или запишите аудио перед отправкой запроса.")
        return "Ошибка: Аудиофайл не предоставлен." # Возвращаем строку ошибки

    print(f"Начало транскрибации файла: {filepath}, Задача: {task}")
    try:
        result = pipe(filepath, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
        text = result["text"]
        print("Транскрибация завершена успешно.")
        return text
    except Exception as e:
        print(f"Ошибка во время транскрибации: {e}")
        # Используем gr.Error для вывода серьезной ошибки
        raise gr.Error(f"Произошла ошибка при обработке аудио: {e}")


def download_audio(video_url: str, download_folder: str) -> str:

    if not video_url:
        gr.Warning("URL YouTube видео не предоставлен.")
        return None

    unique_filename = f"{uuid.uuid4()}.mp3"
    audio_filepath = os.path.join(download_folder, unique_filename)
    
    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': audio_filepath,
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'quiet': True,
    }

    try:
        os.makedirs(download_folder, exist_ok=True)
    
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([video_url])

    except Exception as e:
        print(f"Не удалось скачать аудио с YouTube: {e}")
        gr.Error(f"Не удалось скачать аудио с YouTube: {e}", duration=5)
        audio_filepath = None
    
    return audio_filepath

def transcribe_youtube(youtube_url, task):
    """
    Скачивает аудио с YouTube и транскрибирует его.

    Args:
        youtube_url (str): URL YouTube видео.
        task (str): Задача для модели ("transcribe" или "translate").

    Returns:
        str: Распознанный текст или сообщение об ошибке.
    """
    audio_filepath = download_audio(youtube_url, "./downloded_audio")
    if audio_filepath:
        try:
            transcription = transcribe(audio_filepath, task)
            # Удаляем временный файл и директорию после транскрибации
            os.remove(audio_filepath)
            os.rmdir(os.path.dirname(audio_filepath))
            return transcription
        except Exception as e:
            # Удаляем временный файл и директорию даже если транскрибация не удалась
            if os.path.exists(audio_filepath):
                 os.remove(audio_filepath)
                 os.rmdir(os.path.dirname(audio_filepath))
            raise e # Перебрасываем ошибку для обработки Gradio
    else:
        return "Не удалось скачать аудио."

# --- Создание интерфейса Gradio ---
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="neutral")) as demo:
    gr.Markdown(
        f"""
        # Whisper: Транскрибация аудио
        Транскрибируйте длинные аудиозаписи (с микрофона или из файла) одним нажатием кнопки!
        Демо использует модель OpenAI Whisper [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME})
        через библиотеку 🤗 Transformers для распознавания речи.
        """
    )

    # Создание вкладок
    with gr.Tabs():
        
        # Вкладка 1: Транскрибация с микрофона
        with gr.TabItem("Микрофон"):
            with gr.Row():
                mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Запись с микрофона")
                task_mic = gr.Radio(["transcribe", "translate"], label="Задача", value="transcribe")
            output_mic = gr.Textbox(label="Результат", lines=7) # Увеличили количество строк для текста
            mic_button = gr.Button("Транскрибировать с микрофона")

        # Вкладка 2: Транскрибация из файла
        with gr.TabItem("Аудиофайл"):
            with gr.Row():
                file_input = gr.Audio(sources=["upload"], type="filepath", label="Загрузить аудиофайл")
                task_file = gr.Radio(["transcribe", "translate"], label="Задача", value="transcribe")
            output_file = gr.Textbox(label="Результат", lines=7)
            file_button = gr.Button("Транскрибировать файл")

        # Вкладка 3: Транскрибация с YouTube URL
        with gr.TabItem("YouTube URL"):
            with gr.Row():
                url_input = gr.Textbox(label="Введите URL YouTube видео")
                task_youtube = gr.Radio(["transcribe", "translate"], label="Задача", value="transcribe")
            output_youtube = gr.Textbox(label="Результат", lines=7)
            youtube_button = gr.Button("Транскрибировать с YouTube")


    # --- Привязка функций к кнопкам ---
    
    # Связываем кнопку "Транскрибировать с микрофона" с функцией transcrib
    mic_button.click(
        fn=transcribe,                     
        inputs=[mic_input, task_mic],      
        outputs=output_mic                 
    )

    # Связываем кнопку "Транскрибировать файл" с функцией transcribe
    file_button.click(
        fn=transcribe,
        inputs=[file_input, task_file],
        outputs=output_file
    )

    # Связываем кнопку "Транскрибировать с YouTube" с функцией transcribe_youtube
    youtube_button.click(
        fn=transcribe_youtube,
        inputs=[url_input, task_youtube],
        outputs=output_youtube
    )

# --- Запуск приложения ---
print("Запуск интерфейса Gradio...")
demo.queue().launch(debug=True, share=False)