app.py

import os
import gradio as gr
from google import genai
from google.genai import types
import mimetypes
from pydub import AudioSegment
import asyncio
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Атрыманне ключоў і мадэляў з пераменных асяроддзя
GEMINI_API_KEY = os.getenv("gemini")
MODEL_NAME_TH = os.getenv("modTH")  # Мадэль для транскрыпцыі
MODEL_NAME = os.getenv("mod")  # Мадэль для выпраўлення фармату і перакладу
PROMPT_TRANSCRIBE = os.getenv("p")  # Промпт для транскрыпцыі

# Стварэнне кліента
client = genai.Client(api_key=GEMINI_API_KEY)

def transcribe_audio(audio_file: str) -> str:
    """Транскрыбуе аўдыяфайл з дапамогай Google GenAI."""
    try:
        if not audio_file or not os.path.exists(audio_file):
            return "Памылка: Файл не знойдзены або не указаны."
            
        mime_type, _ = mimetypes.guess_type(audio_file)
        if not mime_type or not mime_type.startswith("audio"):
            return (
                "Немагчыма вызначыць тып файла або файл не з'яўляецца аўдыяфайлам. "
                "Падтрымліваюцца толькі аўдыяфайлы."
            )

        with open(audio_file, "rb") as f:
            audio_data = f.read()

        # Стварэнне contents для перадачы ў API
        contents = [
            {"text": PROMPT_TRANSCRIBE},
            {"mime_type": mime_type, "data": audio_data},
        ]

        # Выклік API
        response = client.models.generate_content(
            model=MODEL_NAME_TH,
            contents=contents,
            config={"temperature": 0.2}
        )
        
        # Атрыманне адказу
        return response.text.strip()

    except FileNotFoundError:
        logger.error(f"Файл не знойдзены: {audio_file}")
        return "Памылка: Файл не знойдзены."
    except Exception as e:
        logger.error(f"Памылка пры транскрыпцыі: {str(e)}", exc_info=True)
        return f"Нечаканая памылка: {str(e)}"


def fix_subtitles_format(transcript: str) -> str:
    """Выпраўляе фармат часу ў субцітрах."""
    if not transcript:
        return ""
        
    prompt_fix = (
        "Не змяняй тэксты, выправі толькі часовы фармат у субцітрах на правільны, "
        "вось прыклад 00:00:01,589. \nУ адказ напішы толькі субцітры:\n"
        f"{transcript}"
    )
    try:
        response_fix = client.models.generate_content(
            model=MODEL_NAME,
            contents=[{"text": prompt_fix}],
            config={"temperature": 0.2}
        )
        return response_fix.text.strip()
    except Exception as e:
        logger.error(f"Памылка пры выпраўленні субцітраў: {str(e)}", exc_info=True)
        return transcript


def create_srt(transcript: str, filename: str = "subtitles.srt") -> tuple[str, str]:
    """Стварае SRT-файл з транскрыпцыі."""
    if not transcript:
        return "", ""
        
    try:
        with open(filename, "w", encoding="utf-8") as f:
            f.write(transcript)
        return transcript, filename
    except Exception as e:
        logger.error(f"Памылка пры запісе SRT-файла: {str(e)}", exc_info=True)
        return f"Памылка пры запісе SRT-файла: {str(e)}", ""


def process_audio(audio_path: str) -> tuple[str, str]:
    """Апрацоўвае аўдыёфайл: транскрыбуе і стварае SRT."""
    if not audio_path:
        return "Не указаны шлях да аўдыёфайла.", ""
        
    transcript = transcribe_audio(audio_path)
    if transcript.startswith("Памылка") or transcript.startswith("Немагчыма") or transcript.startswith("Нечаканая"):
        return transcript, ""

    fixed_transcript = fix_subtitles_format(transcript)
    text, srt_file = create_srt(fixed_transcript)
    return text, srt_file


def extract_audio_from_video(video_file: str) -> tuple[str, str]:
    """Выдзяляе аўдыёдарожку з відэафайла."""
    if not video_file or not os.path.exists(video_file):
        return "", "Памылка: Відэафайл не знойдзены або не указаны."
        
    try:
        audio = AudioSegment.from_file(video_file)
        audio_path = "extracted_audio.mp3"
        audio.export(audio_path, format="mp3")
        return audio_path, ""
    except Exception as e:
        logger.error(f"Памылка пры выдзяленні аўдыё з відэафайла: {str(e)}", exc_info=True)
        return "", f"Памылка пры выдзяленні аўдыё з відэафайла: {str(e)}"


def process_video(video_path: str) -> tuple[str, str]:
    """Апрацоўвае відэафайл: выдзяляе аўдыё, транскрыбуе і стварае SRT."""
    if not video_path:
        return "Не указаны шлях да відэафайла.", ""
        
    audio_path, error = extract_audio_from_video(video_path)
    if error:
        return error, ""
    return process_audio(audio_path)


def process_file(audio_path: str | None, video_path: str | None) -> tuple[str, str]:
    """Апрацоўвае файл (аўдыё або відэа)."""
    try:
        if audio_path:
            return process_audio(audio_path)
        elif video_path:
            return process_video(video_path)
        else:
            return "Няма файла для апрацоўкі.", ""
    except Exception as e:
        logger.error(f"Памылка пры апрацоўцы файла: {str(e)}", exc_info=True)
        return f"Памылка пры апрацоўцы файла: {str(e)}", ""


def update_on_audio_change(audio_path: str | None) -> gr.update:
    """Абнаўляе інтэрфейс пры змене аўдыёфайла."""
    try:
        return gr.update(value=None, interactive=not bool(audio_path))
    except Exception as e:
        logger.error(f"Памылка пры абнаўленні інтэрфейса: {str(e)}", exc_info=True)
        return gr.update(value=None, interactive=True)


def update_on_video_change(video_path: str | None) -> gr.update:
    """Абнаўляе інтэрфейс пры змене відэафайла."""
    try:
        return gr.update(value=None, interactive=not bool(video_path))
    except Exception as e:
        logger.error(f"Памылка пры абнаўленні інтэрфейса: {str(e)}", exc_info=True)
        return gr.update(value=None, interactive=True)


def translate_transcript(transcript: str, target_language: str) -> tuple[str, str]:
    """Перакладае транскрыпцыю на іншую мову і стварае SRT."""
    if not transcript:
        return "Няма тэксту для перакладу.", ""
        
    prompt_text = (
        f"Перакладзі толькі тэксты субцітраў на {target_language} мову. "
        "Астатняе пакінь як ёсць.\nТэкст:\n{transcript}"
    )
    try:
        response = client.models.generate_content(
            model=MODEL_NAME,
            contents=[{"text": prompt_text}],
            config={"temperature": 0.2}
        )
        translated = response.text.strip()
        translated_srt_filename = "translated_subtitles.srt"
        return create_srt(translated, translated_srt_filename)
    except Exception as e:
        logger.error(f"Памылка пры перакладзе: {str(e)}", exc_info=True)
        return f"Памылка пры перакладзе: {str(e)}", ""


try:
    with gr.Blocks() as demo:
        gr.Markdown("# Транскрыпцыя аўдыя для беларускай мовы")
        gr.Markdown(
            """
    ## Загрузіце аўдыёфайл або відэафайл да 15 хвілін. Калі загружаны аўдыёфайл, відэа неактыўна, і наадварот.
    Субцітры будуць аўтаматычна згенераваны разам з SRT-файлам.
    [Далучайцеся да беларускаймоўнай суполкі ў ТГ](https://t.me/belarusai)
    **Падтрымаць праект:** [Buy me a coffee](https://buymeacoffee.com/tuteishygpt)
            """
        )
        with gr.Row():
            audio_input = gr.Audio(type="filepath", label="Аўдыёфайл")
            video_input = gr.Video(label="Відэафайл")

        audio_input.change(
            fn=update_on_audio_change, inputs=audio_input, outputs=video_input
        )
        video_input.change(
            fn=update_on_video_change, inputs=video_input, outputs=audio_input
        )

        btn = gr.Button("Апрацаваць")
        transcript_output = gr.Textbox(label="Транскрыпцыя", lines=10)
        file_output = gr.File(label="SRT-файл")
        btn.click(
            fn=process_file, inputs=[audio_input, video_input], outputs=[transcript_output, file_output]
        )

        gr.Markdown("## Пераклад субцітраў")
        with gr.Row():
            language_dropdown = gr.Dropdown(
                choices=["English", "Руcкая", "Польская", "Літоўская", "Нямецкая"],
                label="Выберы мову перакладу",
                value="English",
            )
            translate_btn = gr.Button("Пераклад")
        translation_output = gr.Textbox(label="Пераклад", lines=10)
        translation_file_output = gr.File(label="Translated SRT-файл")
        translate_btn.click(
            fn=translate_transcript,
            inputs=[transcript_output, language_dropdown],
            outputs=[translation_output, translation_file_output],
        )

    # Запуск з дадатковымі параметрамі для стабільнасці
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False,
        show_error=True,
        prevent_thread_lock=True
    )
except Exception as e:
    logger.critical(f"Крытычная памылка пры запуску праграмы: {str(e)}", exc_info=True)
    print(f"Крытычная памылка: {str(e)}")