Create app.py

app.py

@@ -0,0 +1,287 @@
+import os
+import re
+import subprocess
+import tempfile
+from functools import lru_cache
+from pathlib import Path
+
+import gradio as gr
+import torch
+from transformers import pipeline
+
+# =========================
+# Settings
+# =========================
+WHISPER_MODEL = os.getenv("WHISPER_MODEL", "openai/whisper-small")
+NLLB_MODEL = os.getenv("NLLB_MODEL", "facebook/nllb-200-distilled-600M")
+TARGET_LANG = "mya_Mymr"  # Burmese (Myanmar)
+
+# Common NLLB language codes
+LANGS = {
+    "English": "eng_Latn",
+    "Myanmar": "mya_Mymr",
+    "Thai": "tha_Thai",
+    "Japanese": "jpn_Jpan",
+    "Korean": "kor_Hang",
+    "Chinese (Simplified)": "zho_Hans",
+    "Hindi": "hin_Deva",
+    "French": "fra_Latn",
+    "Spanish": "spa_Latn",
+    "German": "deu_Latn",
+    "Russian": "rus_Cyrl",
+    "Arabic": "arb_Arab",
+    "Indonesian": "ind_Latn",
+    "Vietnamese": "vie_Latn",
+}
+
+VIDEO_EXTS = {".mp4", ".mkv", ".mov", ".webm", ".avi", ".flv", ".m4v"}
+AUDIO_EXTS = {".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg", ".opus"}
+
+
+# =========================
+# Helpers
+# =========================
+def get_device():
+    return 0 if torch.cuda.is_available() else -1
+
+
+@lru_cache(maxsize=1)
+def get_asr():
+    device = get_device()
+    kwargs = {
+        "task": "automatic-speech-recognition",
+        "model": WHISPER_MODEL,
+        "device": device,
+    }
+    return pipeline(**kwargs)
+
+
+@lru_cache(maxsize=32)
+def get_translator(src_lang_code: str):
+    device = get_device()
+    kwargs = {
+        "task": "translation",
+        "model": NLLB_MODEL,
+        "src_lang": src_lang_code,
+        "tgt_lang": TARGET_LANG,
+        "device": device,
+    }
+    return pipeline(**kwargs)
+
+
+def ffmpeg_to_wav(input_path: str) -> str:
+    """
+    Convert any audio/video file to 16kHz mono WAV for stable ASR.
+    Requires ffmpeg installed.
+    """
+    input_path = str(input_path)
+    out_dir = tempfile.mkdtemp(prefix="audio_")
+    out_wav = str(Path(out_dir) / "audio.wav")
+
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        input_path,
+        "-vn",
+        "-ac",
+        "1",
+        "-ar",
+        "16000",
+        out_wav,
+    ]
+
+    try:
+        subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    except FileNotFoundError as e:
+        raise RuntimeError("ffmpeg မတွေ့ပါ။ apt.txt ထဲမှာ ffmpeg ထည့်ပါ။") from e
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(
+            "ဖိုင်ကို audio အဖြစ်ပြောင်းမရပါ။ video/audio file ကို ပြန်စစ်ပါ။"
+        ) from e
+
+    return out_wav
+
+
+def format_timestamp(seconds: float) -> str:
+    if seconds is None:
+        seconds = 0.0
+    ms = int(round((seconds - int(seconds)) * 1000))
+    total = int(seconds)
+    hh = total // 3600
+    mm = (total % 3600) // 60
+    ss = total % 60
+    return f"{hh:02d}:{mm:02d}:{ss:02d},{ms:03d}"
+
+
+def smart_two_line(text: str, width: int = 44) -> str:
+    text = re.sub(r"\s+", " ", text).strip()
+    if len(text) <= width:
+        return text
+
+    # Try to split near the middle on a space
+    words = text.split(" ")
+    if len(words) == 1:
+        mid = len(text) // 2
+        return text[:mid].rstrip() + "\n" + text[mid:].lstrip()
+
+    total_len = len(text)
+    best_idx = 1
+    best_diff = float("inf")
+    current = 0
+
+    for i, w in enumerate(words[:-1], start=1):
+        current += len(w) + 1
+        diff = abs(current - total_len / 2)
+        if diff < best_diff:
+            best_diff = diff
+            best_idx = i
+
+    line1 = " ".join(words[:best_idx]).strip()
+    line2 = " ".join(words[best_idx:]).strip()
+
+    if len(line1) > width and len(line2) > width:
+        mid = len(text) // 2
+        return text[:mid].rstrip() + "\n" + text[mid:].lstrip()
+
+    return line1 + "\n" + line2
+
+
+def translate_text(src_text: str, src_lang_code: str) -> str:
+    src_text = src_text.strip()
+    if not src_text:
+        return ""
+
+    # If already Burmese, keep as-is
+    if src_lang_code == TARGET_LANG:
+        return src_text
+
+    translator = get_translator(src_lang_code)
+    result = translator(src_text, max_new_tokens=256)
+
+    if isinstance(result, list) and result:
+        return result[0].get("translation_text", "").strip()
+    if isinstance(result, dict):
+        return result.get("translation_text", "").strip()
+    return str(result).strip()
+
+
+def normalize_chunks(asr_result):
+    chunks = asr_result.get("chunks")
+    if chunks and isinstance(chunks, list):
+        return chunks
+
+    # Fallback: single subtitle
+    text = asr_result.get("text", "").strip()
+    if not text:
+        return []
+    return [{"timestamp": (0.0, None), "text": text}]
+
+
+def build_srt(input_path: str, source_language_name: str):
+    if not input_path:
+        raise gr.Error("ဖိုင်တင်ပါ။")
+
+    if source_language_name not in LANGS:
+        raise gr.Error("Source language မမှန်ပါ။")
+
+    src_lang_code = LANGS[source_language_name]
+    audio_path = ffmpeg_to_wav(input_path)
+
+    asr = get_asr()
+    # Whisper pipeline supports timestamps for ASR. We keep them for SRT timing.
+    asr_result = asr(
+        audio_path,
+        return_timestamps=True,
+        generate_kwargs={
+            "task": "transcribe",
+            "language": src_lang_code,
+        },
+    )
+
+    chunks = normalize_chunks(asr_result)
+
+    subtitles = []
+    index = 1
+
+    for chunk in chunks:
+        text = (chunk.get("text") or "").strip()
+        if not text:
+            continue
+
+        ts = chunk.get("timestamp")
+        start = 0.0
+        end = None
+
+        if isinstance(ts, (tuple, list)) and len(ts) >= 2:
+            start = ts[0] if ts[0] is not None else 0.0
+            end = ts[1]
+
+        translated = translate_text(text, src_lang_code)
+        if not translated:
+            continue
+
+        if end is None:
+            # Safe fallback duration if the model gives only a start timestamp
+            end = start + max(2.0, min(6.0, len(translated) / 10.0))
+
+        subtitles.append(
+            f"{index}\n"
+            f"{format_timestamp(float(start))} --> {format_timestamp(float(end))}\n"
+            f"{smart_two_line(translated)}\n"
+        )
+        index += 1
+
+    if not subtitles:
+        raise gr.Error("Subtitle ထုတ်မရပါ။ အသံမရှင်းတာ သို့မဟုတ် ဖိုင်ပြဿနာရှိနိုင်တယ်။")
+
+    srt_text = "\n".join(subtitles).strip() + "\n"
+
+    out_dir = tempfile.mkdtemp(prefix="srt_")
+    out_path = str(Path(out_dir) / f"{Path(input_path).stem}_burmese.srt")
+    with open(out_path, "w", encoding="utf-8") as f:
+        f.write(srt_text)
+
+    preview = srt_text[:8000]
+    return out_path, preview
+
+
+# =========================
+# UI
+# =========================
+with gr.Blocks(title="Burmese SRT Generator") as demo:
+    gr.Markdown(
+        """
+# Burmese SRT Generator
+Video / Audio file တင်ပြီး မြန်မာလို subtitle `.srt` ထုတ်မယ်။
+"""
+    )
+
+    with gr.Row():
+        input_file = gr.File(
+            label="Video / Audio File",
+            file_types=[
+                ".mp4", ".mkv", ".mov", ".webm", ".avi", ".flv", ".m4v",
+                ".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg", ".opus",
+            ],
+            type="filepath",
+        )
+        source_lang = gr.Dropdown(
+            choices=list(LANGS.keys()),
+            value="English",
+            label="Source Language",
+        )
+
+    run_btn = gr.Button("Generate Burmese SRT", variant="primary")
+
+    output_file = gr.File(label="Download .srt")
+    preview_box = gr.Textbox(label="Preview", lines=18)
+
+    run_btn.click(
+        fn=build_srt,
+        inputs=[input_file, source_lang],
+        outputs=[output_file, preview_box],
+    )
+
+if __name__ == "__main__":
+    demo.launch()