Update app.py

app.py

@@ -1,209 +1,293 @@
+# gemini_srt_generator_improved.py
+"""A Gradio interface that generates SRT subtitles from audio or video files
+using the new `google.genai` SDK (Gemini models).
+
+Key improvements over the previous version
+-----------------------------------------
+1. Switched from **google.generativeai** ➜ **google.genai** (new SDK).
+2. File validation (size / MIME‑type) **before** uploading to Gemini → cheaper &
+   safer.
+3. Robust *retry* wrapper with exponential back‑off + global request timeout.
+4. Automatic timestamped history folder (`transcripts/`) so results are never
+   overwritten.
+5. Cleaner UI: • progress bar (gr.Progress) • automatic scroll to bottom in the
+   live status textbox.
+6. Configurable model + token limit via UI dropdown.
+7. Minor refactor & type hints.
+"""
+from __future__ import annotations
+
+import json
+import mimetypes
 import os
+import threading
+import time
+import uuid
+from datetime import datetime
+from functools import wraps
+from pathlib import Path
+from typing import Callable, List, Tuple
+
 import gradio as gr
-import google.generativeai as genai
-import mimetypes
+import google.genai as genai  # NEW SDK ✅
 from pydub import AudioSegment
 
-GEMINI_API_KEY = os.getenv("gembeh")
-MODEL_NAME_TH = os.getenv("modTH")
-MODEL_NAME = os.getenv("mod")
+# -----------------------
+# CONSTANTS & CONFIG
+# -----------------------
+MAX_FILE_SIZE_MB = 200  # Hard limit to prevent huge uploads
+ALLOWED_AUDIO_PREFIXES = ("audio/",)
+ALLOWED_VIDEO_PREFIXES = ("video/",)
+HISTORY_DIR = Path("transcripts")
+HISTORY_DIR.mkdir(exist_ok=True)
 
-genai.configure(api_key=GEMINI_API_KEY)
+DEFAULT_MODEL = "gemini-2.5-flash-preview-04-17"
+FALLBACK_MODEL = "gemini-2.5-flash"
 
-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 = os.getenv("p")
-        model = genai.GenerativeModel(MODEL_NAME_TH)
-        response = model.generate_content(
-            [prompt_text, {"mime_type": mime_type, "data": audio_data}]
+GENERATION_BASE_CONFIG = {
+    "temperature": 0.35,
+    "top_p": 0.95,
+    "top_k": 64,
+    "response_mime_type": "application/json",
+}
+
+# -----------------------
+#  Utils
+# -----------------------
+
+def retry(retries: int = 3, delay: float = 3.0, backoff: float = 2.0):
+    """Simple exponential‑backoff retry decorator."""
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            _delay = delay
+            last_exc = None
+            for attempt in range(1, retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except Exception as exc:  # noqa: BLE001
+                    last_exc = exc
+                    if attempt == retries:
+                        break
+                    time.sleep(_delay)
+                    _delay *= backoff
+            raise last_exc  # Re‑raise after exhausting retries
+
+        return wrapper
+
+    return decorator
+
+
+def seconds_to_timestamp(sec: float) -> str:
+    h, remainder = divmod(sec, 3600)
+    m, remainder = divmod(remainder, 60)
+    s = int(remainder)
+    ms = int(round((remainder - s) * 1000))
+    return f"{int(h):02d}:{int(m):02d}:{s:02d},{ms:03d}"
+
+
+# -----------------------
+#  Validation helpers
+# -----------------------
+
+def _validate_file(path: str, allowed_prefixes: tuple[str, ...]) -> None:
+    if not path or not os.path.isfile(path):
+        raise ValueError("Файл не знойдзены.")
+
+    size_mb = os.path.getsize(path) / 1024 / 1024
+    if size_mb > MAX_FILE_SIZE_MB:
+        raise ValueError(
+            f"Файл занадта вялікі: {size_mb:.1f} MB > {MAX_FILE_SIZE_MB} MB."
         )
-        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):
+
+    mime, _ = mimetypes.guess_type(path)
+    if not mime or not mime.startswith(allowed_prefixes):
+        raise ValueError(f"Непадтрыманы тып файла: {mime or 'невядомы'}.")
+
+
+# -----------------------
+#  Gemini helpers
+# -----------------------
+
+def _configure_genai(api_key: str) -> None:
+    if not api_key:
+        raise ValueError("Не знойдзены API‑ключ для Gemini (env var `GEMINI_API_KEY`).")
+    genai.configure(api_key=api_key, request_timeout=90)  # global 90 s timeout
+
+
+def _get_model(name: str):
+    return genai.GenerativeModel(model_name=name, generation_config=GENERATION_BASE_CONFIG)
+
+
+@retry(retries=3)
+def _upload_to_gemini(path: str, status_callback: Callable[[str], None]):
+    mime_type, _ = mimetypes.guess_type(path)
+    status_callback("📤 Загружаем файл у Gemini …")
+    file_obj = genai.upload_file(path, mime_type=mime_type)
+    status_callback("✅ Файл загружаны.")
+    return file_obj
+
+
+@retry(retries=3)
+def _transcribe(file_obj, model, status_callback: Callable[[str], None]):
+    status_callback("🔍 Пачынаем транскрыпцыю …")
+    chat = model.start_chat(history=[])
+    return chat.send_message(file_obj)
+
+
+# -----------------------
+#  Core processing
+# -----------------------
+
+def transcribe_audio(audio_path: str, model_name: str, status_callback: Callable[[str], None]):
+    _validate_file(audio_path, ALLOWED_AUDIO_PREFIXES)
+    file_obj = _upload_to_gemini(audio_path, status_callback)
+
+    stop_event = threading.Event()
+
+    def _progress():
+        frames = ["⏳", "⏳.", "⏳..", "⏳..."]
+        while not stop_event.is_set():
+            for frame in frames:
+                if stop_event.is_set():
+                    break
+                status_callback(f"Транскрыпцыя ідзе {frame}")
+                time.sleep(0.6)
+
+    thread = threading.Thread(target=_progress)
+    thread.start()
+
     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()
+        model = _get_model(model_name)
+        response = _transcribe(file_obj, model, status_callback)
+    finally:
+        stop_event.set()
+        thread.join()
+
+    if not response.text:
+        raise RuntimeError("❌ Пусты адказ ад мадэлі.")
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    raw_json_path = HISTORY_DIR / f"response_{timestamp}.json"
+    raw_json_path.write_text(response.text, encoding="utf-8")
+
+    status_callback("📥 Апрацоўка транскрыпцыі …")
+    return json.loads(response.text)
+
+
+def transcripts_to_srt(transcripts: List[dict]) -> Tuple[str, Path]:
+    srt_lines: list[str] = []
+    for idx, seg in enumerate(transcripts, start=1):
+        start_ts = seconds_to_timestamp(seg["start"])
+        end_ts = seconds_to_timestamp(seg["end"])
+        srt_lines.append(f"{idx}\n{start_ts} --> {end_ts}\n{seg['text']}\n")
+    content = "\n".join(srt_lines)
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = HISTORY_DIR / f"subtitles_{timestamp}.srt"
+    filename.write_text(content, encoding="utf-8")
+    return content, filename
+
+
+def extract_audio_from_video(video_file: str, status_callback: Callable[[str], None]) -> str:
+    _validate_file(video_file, ALLOWED_VIDEO_PREFIXES)
+    status_callback("🎞 Вылучаем аўдыё з відэа …")
+    audio = AudioSegment.from_file(video_file)
+    path = f"extracted_{uuid.uuid4().hex}.mp3"
+    audio.export(path, format="mp3")
+    status_callback("✅ Аўдыё вылучана.")
+    return path
+
+
+def process_audio(audio_path: str, model_name: str, status_callback):
+    transcripts = transcribe_audio(audio_path, model_name, status_callback)
+    status_callback("📝 Канвертацыя ў SRT …")
+    return transcripts_to_srt(transcripts)
+
+
+def process_video(video_path: str, model_name: str, status_callback):
+    audio_path = extract_audio_from_video(video_path, status_callback)
+    return process_audio(audio_path, model_name, status_callback)
+
+
+def process_file(audio: str | None, video: str | None, model_name: str, progress: gr.Progress, status_callback):
+    status_callback("🔄 Пачатак апрацоўкі …")
+    result: Tuple[str, Path]
+    with progress:
+        if audio:
+            result = process_audio(audio, model_name, status_callback)
+        elif video:
+            result = process_video(video, model_name, status_callback)
         else:
-            fixed_transcript = transcript
-        return fixed_transcript
-    except Exception as e:
-        return transcript
+            raise ValueError("Ні адзін файл не загружаны.")
+    status_callback("✅ Гатова!")
+    return result
 
-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 check_audio_length(audio):
-    if audio is not None:
-        try:
-            audio_seg = AudioSegment.from_file(audio)
-            if audio_seg.duration_seconds > 600:
-                return "Памылка: Аўдыёфайл даўжэй за 10 хвілін."
-            else:
-                return ""
-        except Exception as e:
-            return f"Памылка пры праверцы аўдыё: {str(e)}"
-    return ""
-
-def check_video_length(video):
-    if video is not None:
-        try:
-            audio_seg = AudioSegment.from_file(video)
-            if audio_seg.duration_seconds > 600:
-                return "Памылка: Відэафайл даўжэй за 10 хвілін."
-            else:
-                return ""
-        except Exception as e:
-            return f"Памылка пры праверцы відэа: {str(e)}"
-    return ""
-
-def process_file(audio, video):
-    if audio is not None:
-        error = check_audio_length(audio)
-        if error:
-            return error, None
-        return process_audio(audio)
-    elif video is not None:
-        error = check_video_length(video)
-        if error:
-            return error, None
-        return process_video(video)
-    else:
-        return "Няма файла для апрацоўкі.", None
-
-def on_audio_change(audio):
-    # Калі загружаны аўдыёфайл, адключаем відэафайл і правяраем працягласць
-    if audio is not None:
-        error_msg = check_audio_length(audio)
-        return gr.update(value=None, interactive=False), error_msg
-    else:
-        return gr.update(interactive=True), ""
-
-def on_video_change(video):
-    # Калі загружаны відэафайл, адключаем аўдыёфайл і правяраем працягласць
-    if video is not None:
-        error_msg = check_video_length(video)
-        return gr.update(value=None, interactive=False), error_msg
-    else:
-        return gr.update(interactive=True), ""
-
-def translate_transcript(transcript, target_language):
-    try:
-        prompt_text = (
-            f"перакладзі толькі тэксты субцітраў на {target_language} мову. Астатняя пакінь як ёсць."
-            f"Тэкст:\n{transcript}"
+
+# -----------------------
+#  Gradio UI
+# -----------------------
+
+def build_ui():
+    api_key_default = os.getenv("GEMINI_API_KEY", "")
+
+    with gr.Blocks(title="Gemini SRT Generator (Belarusian Edition)") as demo:
+        gr.Markdown(
+            """
+            ## Загрузіце аўдыё- ці відэафайл — атрымайце субцітры SRT
+            [Суполка беларускага ШІ](https://t.me/belarusai) • 
+            [Buy Me A Coffee](https://buymeacoffee.com/tuteishygpt)
+            """
         )
-        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:
-    # Дадаем Google Analytics код праз HTML-кампанент
-    gr.HTML("""
-    <!-- Google tag (gtag.js) -->
-    <script async src="https://www.googletagmanager.com/gtag/js?id=G-2QZ4X58TG6"></script>
-    <script>
-      window.dataLayer = window.dataLayer || [];
-      function gtag(){dataLayer.push(arguments);}
-      gtag('js', new Date());
-      gtag('config', 'G-2QZ4X58TG6');
-    </script>
-    """)
-
-    gr.Markdown("# Транскрыпцыя кароткіх аўдыя для беларускай мовы")
-    gr.Markdown(
-        """
-## Загрузіце аўдыёфайл або відэафайл да 10 хвілін. Субцітры з кароткімі тэкстамі будуць згенераваны разам з файлам субцітраў. 
-[Ёсць пытанні ці прапановы? Далучайцеся да беларускаймоўнай суполкі штучнага інтэлекту](https://t.me/belarusai)
-**Хочаце каб сэрвіс працаваў? Налівайце каву! :** [Buy me a coffee](https://buymeacoffee.com/tuteishygpt)
-**Агучце беларускую мову тут :** [Беларуская мадэль маўлення](https://huggingface.co/spaces/archivartaunik/Bextts)
-
-        """
-    )
-    with gr.Row():
-        audio_input = gr.Audio(type="filepath", label="Аўдыёфайл")
-        video_input = gr.Video(label="Відэафайл")
-    # Поле Транскрыпцыя для паказу памылак будзе агульным
-    transcript_output = gr.Textbox(label="Транскрыпцыя", lines=10)
-    # Пры загрузцы аўдыё або відэа запускаем праверку працягласці і абнаўляем адпаведна поле Транскрыпцыя
-    audio_input.change(fn=on_audio_change, inputs=audio_input, outputs=[video_input, transcript_output])
-    video_input.change(fn=on_video_change, inputs=video_input, outputs=[audio_input, transcript_output])
-    
-    btn = gr.Button("Апрацаваць")
-    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"
+
+        with gr.Row():
+            api_key_box = gr.Textbox(
+                label="🔑 Gemini API‑key (калі не ўсталяваны як env)",
+                type="password",
+                value=api_key_default,
+            )
+        with gr.Row():
+            model_dropdown = gr.Dropdown(
+                [DEFAULT_MODEL, FALLBACK_MODEL],
+                value=DEFAULT_MODEL,
+                label="🧠 Мадэль Gemini",
+            )
+        with gr.Row():
+            audio_input = gr.Audio(type="filepath", label="🎙 Аўдыёфайл")
+            video_input = gr.Video(label="🎥 Відэафайл")
+        btn = gr.Button("🚀 Апрацаваць")
+        with gr.Row():
+            transcript_output = gr.Textbox(
+                label="📄 SRT-транскрыпцыя", lines=12, autoscroll=True
+            )
+            file_output = gr.File(label="⬇️ SRT-файл")
+        status_output = gr.Textbox(label="🛠️ Статус", interactive=False, autoscroll=True)
+
+        def wrapped_process(audio, video, api_key, model_name, progress=gr.Progress()):
+            _configure_genai(api_key or api_key_default)
+
+            def update_status(text):
+                status_output.value = text
+                # Force scroll to bottom (JS hack)
+                status_output.scroll_to_end()
+
+            content, file_path = process_file(audio, video, model_name, progress, update_status)
+            return content, file_path
+
+        btn.click(
+            fn=wrapped_process,
+            inputs=[audio_input, video_input, api_key_box, model_dropdown],
+            outputs=[transcript_output, file_output],
         )
-        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()
+
+    return demo
+
+
+def main():
+    demo = build_ui()
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()