app.py

import gradio as gr
from argparse import Namespace
import os
import shutil
import sys
import yt_dlp  # для працы з YouTube
import glob  # Import glob for file searching

GEMINI_API_KEY = os.getenv("gem2")


from sub_tools.media.converter import hls_to_media, media_to_signature, video_to_audio
from sub_tools.media.segmenter import segment_audio
from sub_tools.subtitles.combiner import combine_subtitles
from sub_tools.system.directory import change_directory
from sub_tools.system.console import header, success, error
from sub_tools.transcribe import transcribe

# Вызначэнне функцый для дэактывацыі процілеглага поля
def update_on_audio_change(audio):
    if audio is not None:
        return gr.update(interactive=False)
    return gr.update(interactive=True)

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

def main_logic(args: Namespace) -> tuple:
    """
    Асноўная логіка прыкладання: ад загрузкі відэа/аўдыё да зліцця субтытраў.
    Пасля зліцця субтытраў вяртае (тэкст субтытраў, шлях да SRT‑файла для спампоўкі).
    """
    output_str = ""
    subtitles_text = ""
    srt_file_path = None

    try:
        change_directory(args.output_path)
        step = 1

        if "video" in args.tasks:
            if not args.hls_url:
                output_str += f"{step}. Download Video: No video file uploaded\n"
                raise Exception("No video file uploaded")
            header(f"{step}. Download Video")
            output_str += f"{step}. Download Video: Started\n"
            hls_to_media(args.hls_url, args.video_file, False, args.overwrite)
            success("Done!")
            output_str += "Done!\n"
            step += 1

        if "audio" in args.tasks:
            header(f"{step}. Video to Audio")
            output_str += f"{step}. Video to Audio: Started\n"
            video_to_audio(args.video_file, args.audio_file, args.overwrite)
            success("Done!")
            output_str += "Done!\n"
            step += 1

        if "signature" in args.tasks:
            header(f"{step}. Audio to Signature")
            output_str += f"{step}. Audio to Signature: Started\n"
            media_to_signature(args.audio_file, args.signature_file, args.overwrite)
            success("Done!")
            output_str += "Done!\n"
            step += 1

        if "segment" in args.tasks:
            header(f"{step}. Segment Audio")
            output_str += f"{step}. Segment Audio: Started\n"

            print(f"Current working directory before segment_audio: {os.getcwd()}")  # Debug CWD
            print(f"Audio file path for segmentation: {args.audio_file}")  # Debug audio file path
            if not os.path.exists(args.audio_file):  # Critical file existence check
                error_msg = f"Error: Audio file not found before segmentation: {args.audio_file}"
                error(error_msg)
                return (error_msg, None)
            print(f"Audio file exists before segmentation: {args.audio_file}")  # Debug file existence

            segment_audio(args.audio_file, args.audio_segment_prefix, args.audio_segment_format, args.audio_segment_length, args.overwrite)
            success("Done!")
            output_str += "Done!\n"
            step += 1

        if "transcribe" in args.tasks:
            if not (args.gemini_api_key and args.gemini_api_key.strip()):
                output_str += f"{step}. Transcribe Audio: No Gemini API Key provided\n"
                raise Exception("No Gemini API Key provided")
            header(f"{step}. Transcribe Audio")
            output_str += f"{step}. Transcribe Audio: Started\n"
            transcribe(args)
            success("Done!")
            output_str += "Done!\n"
            step += 1

        if "combine" in args.tasks:
            header(f"{step}. Combine Subtitles")
            output_str += f"{step}. Combine Subtitles: Started\n"
            combine_subtitles(args.languages, args.audio_segment_prefix, args.audio_segment_format)
            success("Done!")
            output_str += "Done!\n"
            if args.languages:
                language = args.languages[0]
                srt_file_path = os.path.join(os.getcwd(), f"{language}.srt")
                try:
                    with open(srt_file_path, "r", encoding="utf-8") as f:
                        subtitles_text = f.read()
                except Exception as e:
                    subtitles_text = f"Error reading subtitles file: {str(e)}"
            else:
                subtitles_text = "No language specified"
            step += 1

        return (subtitles_text, srt_file_path)

    except Exception as e:
        error_msg = f"Error: {str(e)}"
        error(error_msg)
        return (error_msg, None)

def run_subtools(
    tasks,
    hls_url,
    video_file,
    audio_file,
    signature_file,
    output_path,
    languages,
    overwrite,
    retry,
    gemini_api_key,
    debug,
    audio_segment_prefix,
    audio_segment_format,
    audio_segment_length
):
    """
    Падрыхтоўка каталога вываду і запуск асноўнай логікі.
    """
    if os.path.exists(output_path):
        shutil.rmtree(output_path)
    os.makedirs(output_path, exist_ok=True)

    if isinstance(languages, str):
        languages = [lang.strip() for lang in languages.split(",") if lang.strip()]

    args = Namespace(
        tasks=tasks,
        hls_url=hls_url,
        video_file=video_file,
        audio_file=audio_file,
        signature_file=signature_file,
        output_path=output_path,
        languages=languages,
        overwrite=overwrite,
        retry=retry,
        gemini_api_key=gemini_api_key,
        debug=debug,
        audio_segment_prefix=audio_segment_prefix,
        audio_segment_format=audio_segment_format,
        audio_segment_length=audio_segment_length,
    )

    return main_logic(args)


def transcribe_youtube(youtube_url: str) -> tuple:
    """
    Спампоўвае аўдыё з відэа YouTube праз yt_dlp і вяртае паведамленне і шлях да часовага аўдыёфайла.
    Выкарыстоўвае chromewebstore.google.com_cookies.txt, калі ён ёсць.
    """
    if not youtube_url:
        return "Не ўведзена спасылка", None

    output_dir = "output_youtube_downloads"  # Define a specific output directory
    os.makedirs(output_dir, exist_ok=True)  # Create the directory if it doesn't exist
    outtmpl = os.path.join(output_dir, "temp_youtube_audio.%(ext)s")  # Path within the output dir

    try:
        ydl_opts = {
            'format': 'bestaudio/best',
            'outtmpl': outtmpl,
            #'quiet': True,  # Прыбіраем quiet для адладкі
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',  # Download as MP3
                'preferredquality': '0',  # лепшая якасць
            }],
        }

        default_cookies_file = "chromewebstore.google.com_cookies.txt"
        if os.path.exists(default_cookies_file):
            ydl_opts['cookiefile'] = default_cookies_file
            print(f"Выкарыстоўваецца файл cookie па змаўчанні: {default_cookies_file}")

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([youtube_url])

        # Шукаем створаны аўдыёфайл у output_dir
        audio_files = glob.glob(os.path.join(output_dir, "temp_youtube_audio.*"))
        print(f"Знойдзеныя аўдыёфайлы ў {output_dir}: {audio_files}")  # Дадаем вывад для адладкі
        if not audio_files:
            return ("Памылка: файл аўдыё не створаны", None)

        audio_file_rel = None  # Relative path
        for f in audio_files:
            if f.endswith(".mp3"):
                audio_file_rel = f
                break
        if audio_file_rel is None:
            audio_file_rel = audio_files[0]

        audio_file_abs = os.path.abspath(audio_file_rel)  # Get absolute path here!

        print(f"Абраны аўдыёфайл (relative): {audio_file_rel}")  # Дадаем вывад для адладкі
        print(f"Абраны аўдыёфайл (absolute): {audio_file_abs}")  # Дадаем вывад для адладкі
        return ("Спампоўка YouTube аўдыё завершана", audio_file_abs)  # Return absolute path!
    except Exception as e:
        err_msg = f"Памылка пры апрацоўцы YouTube (yt_dlp): {e}. Калі ласка, пераканайцеся, што спасылка сапраўдная, і праверце, ці даступныя відэа."
        print(err_msg)
        return (err_msg, None)


def process_youtube_url(youtube_url):
    """
    Апрацоўвае YouTube спасылку: спампоўвае аўдыё, адлюстроўвае аўдыёфайл і запускае ланцуг апрацоўкі для стварэння субтытраў.
    """
    download_msg, audio_file = transcribe_youtube(youtube_url)
    if not audio_file:
        return download_msg, None  # Return None for other outputs if download fails

    full_audio_file_path = audio_file
    print(f"Поўны шлях да аўдыёфайла: {full_audio_file_path}")

    if not os.path.exists(full_audio_file_path):
        return f"Памылка: Аўдыёфайл не знойдзены па шляху: {full_audio_file_path}", None
    print(f"Файл існуе: {full_audio_file_path}")

    output_text_file, output_file_file = process_uploaded_file(audio=full_audio_file_path, video=None)

    return output_text_file, output_file_file  # Correct return order: text, file_path


def process_uploaded_file(audio, video):
    """
    Выбірае, які файл загружаны, і фармуе параметры для апрацоўкі.
    Калі загружаны аўдыёфайл – запускае апрацоўку для аўдыё,
    калі відэафайл – запускае поўны ланцуг апрацоўкі.
    """
    if audio is not None and video is None:
        tasks = ["signature", "segment", "transcribe", "combine"]
        video_file = ""  # відэа не выкарыстоўваецца
        audio_file = audio
        hls_url = ""  # не патрабуецца
    elif video is not None and audio is None:
        tasks = ["video", "audio", "signature", "segment", "transcribe", "combine"]
        video_file = video
        audio_file = "audio.mp3"  # прызначаем імя для аўдыёфайла
        hls_url = "dummy"  # задаём няпустое значэнне для праверкі
    else:
        return "Error: Загрузіце толькі АЎДЫЁ або ВІДЭАфайл, а не абодва.", None

    return run_subtools(
        tasks=tasks,
        hls_url=hls_url,
        video_file=video_file,
        audio_file=audio_file,
        signature_file="message.shazamsignature",
        output_path="output",
        languages="be",
        overwrite=False,
        retry=50,
        gemini_api_key=GEMINI_API_KEY,
        debug=True,
        audio_segment_prefix="audio_segment",
        audio_segment_format="mp3",
        audio_segment_length=300000
    )

# --------------------- Gradio UI ---------------------

with gr.Blocks() as demo:
    gr.Markdown("# пакуль не працуе! Транскрыпцыя аўдыя для беларускай мовы. Перарабляем на больш якаснае разпазнаванне для доўгіх відэа")
    gr.Markdown(
        """
## [Можна выкарыстаць сэрвіс для кароткіх відэа тут] (https://huggingface.co/spaces/archivartaunik/SubtitlesBE)
[Ёсць пытанні ці прапановы? Далучайцеся да беларускаймоўнай суполкі штучнага інтэлекту](https://t.me/belarusai)
**Хочаце каб сэрвіс працаваў? Налівайце каву! :** [Buy me a coffee](https://buymeacoffee.com/tuteishygpt)
**Агучце беларускую мову тут :** [Беларуская мадэль маўлення](https://huggingface.co/spaces/archivartaunik/Bextts)
        """
    )
    with gr.Tabs():
        with gr.Tab("Файл"):
            with gr.Row():
                audio_input = gr.Audio(type="filepath", label="Аўдыёфайл")  # Keep type="filepath"
                video_input = gr.Video(label="Відэафайл", interactive=True)  # Reset to interactive True initially

            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)
            submit_btn_file = gr.Button("Submit")
            output_text_file = gr.Textbox(label="Тэкст субтытраў")
            output_file_file = gr.File(label="Спампаваць SRT файл")

            submit_btn_file.click(fn=process_uploaded_file, inputs=[audio_input, video_input], outputs=[output_text_file, output_file_file])

        with gr.Tab("YouTube"):
            youtube_url_input = gr.Textbox(label="YouTube URL", placeholder="Устаўце спасылку на відэа YouTube")
            submit_btn_youtube = gr.Button("Submit")
            output_text_youtube = gr.Textbox(label="Тэкст субтытраў")
            output_file_youtube = gr.File(label="Спампаваць SRT файл")

            submit_btn_youtube.click(
                fn=process_youtube_url,
                inputs=youtube_url_input,
                outputs=[output_text_youtube, output_file_youtube]  # Correct output order
            )

demo.launch(allowed_paths=["output_youtube_downloads"])