app.py

import os
import gradio as gr
import google.generativeai as genai
import mimetypes
from pydub import AudioSegment

GEMINI_API_KEY = os.getenv("gemini")
MODEL_NAME_TH = os.getenv("modTH")
MODEL_NAME = os.getenv("mod")

genai.configure(api_key=GEMINI_API_KEY)

def transcribe_audio(audio_file):
    try:
        mime_type, _ = mimetypes.guess_type(audio_file)
        if mime_type is None:
            return "Немагчыма вызначыць тып файла. Падтрымліваюцца толькі аўдыяфайлы."
        with open(audio_file, "rb") as f:
            audio_data = f.read()
        prompt_text = (
            """The user wants me to transcribe the audio into subtitles in SRT format, with a maximum of three words per one subtitle.
I need to listen to the audio and create subtitles with timestamps. Check time format hours:minutes:seconds,milliseconds (00:00:00,000) after creating all.  Example SRT format:
1
00:00:01,670 --> 00:00:02,030
За мяжою, за мяжою,
2
00:00:02,270 --> 00:00:03,850
ні сваё і не чужое.
3
00:00:04,240 --> 00:00:05,760
Хоць спявай ты, хоць ты грай,
4
00:00:05,770 --> 00:00:06,250
а навокал іншы край.
            """
        )
        model = genai.GenerativeModel(MODEL_NAME_TH)
        response = model.generate_content(
            [prompt_text, {"mime_type": mime_type, "data": audio_data}]
        )
        if response.text:
            transcript = response.text.strip()
        else:
            transcript = "Не атрымалася транскрыбаваць аўдыя. Магчыма, памылка з API."
        return transcript
    except FileNotFoundError:
        return "Памылка: Файл не знойдзены."
    except genai.APIError as e:
        return f"Памылка API: {str(e)}"
    except Exception as e:
        return f"Нечаканая памылка: {str(e)}"

def fix_subtitles_format(transcript):
    """
    Дадатковы запыт да мадэлі, які выпраўляе фармат часоў у субцітрах.
    """
    try:
        prompt_fix = (
            f"Не змяняй тэксты, выправі толькі часовы фармат у субцітрах на правільны, вось прыклад 00:00:01,589 \n"
            f" У адказ напішы толькі субцітры: {transcript}"
        )
        model = genai.GenerativeModel(MODEL_NAME)
        response_fix = model.generate_content(prompt_fix)
        if response_fix.text:
            fixed_transcript = response_fix.text.strip()
        else:
            fixed_transcript = transcript
        return fixed_transcript
    except Exception as e:
        return transcript

def create_srt(transcript, filename="subtitles.srt"):
    try:
        with open(filename, "w", encoding="utf-8") as f:
            f.write(transcript)
        return transcript, filename
    except Exception as e:
        return f"Памылка пры запісе SRT-файла: {str(e)}", None

def process_audio(audio):
    transcript = transcribe_audio(audio)
    if transcript.startswith("Памылка"):
        return transcript, None
    # Дадаем другі запыт для выпраўлення фармату часоў у субцітрах
    fixed_transcript = fix_subtitles_format(transcript)
    text, srt_file = create_srt(fixed_transcript)
    return text, srt_file

def extract_audio_from_video(video_file):
    try:
        audio = AudioSegment.from_file(video_file)
        audio_path = "extracted_audio.mp3"
        audio.export(audio_path, format="mp3")
        return audio_path, None
    except Exception as e:
        return None, f"Памылка пры выдзяленні аўдыі з відэафайла: {str(e)}"

def process_video(video):
    audio_path, error = extract_audio_from_video(video)
    if error:
        return error, None
    return process_audio(audio_path)

def process_file(audio, video):
    if audio is not None:
        return process_audio(audio)
    elif video is not None:
        return process_video(video)
    else:
        return "Няма файла для апрацоўкі.", None

def update_on_audio_change(audio):
    if audio is not None:
        return gr.update(value=None, interactive=False)
    else:
        return gr.update(interactive=True)

def update_on_video_change(video):
    if video is not None:
        return gr.update(value=None, interactive=False)
    else:
        return gr.update(interactive=True)

def translate_transcript(transcript, target_language):
    try:
        prompt_text = (
            f"перакладзі толькі тэксты субцітраў на {target_language} мову. Астатняя пакінь як ёсць."
            f"Тэкст:\n{transcript}"
        )
        model = genai.GenerativeModel(MODEL_NAME)
        response = model.generate_content(prompt_text)
        if response.text:
            translated = response.text.strip()
        else:
            translated = "Не атрымалася перакласці тэкст. Магчыма, памылка з API."
        translated_srt_filename = "translated_subtitles.srt"
        with open(translated_srt_filename, "w", encoding="utf-8") as f:
            f.write(translated)
        return translated, translated_srt_filename
    except Exception as e:
        return f"Памылка пры перакладзе: {str(e)}", None

with gr.Blocks() as demo:
    gr.Markdown("# Транскрыпцыя аўдыя для беларускай мовы")
    gr.Markdown(
        """
## Загрузіце аўдыёфайл або відэафайл да 15 хвілін.  Калі загружаны аўдыёфайл, відэа неактыўна, і наадварот.
Субцітры будуць аўтаматычна згенераваны разам з файлам субцітраў.
[Далучайцеся да беларускаймоўнай суполкі ў ТГ](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()