app.py

"""
=========================================================================
MONGOL DUB — HF Spaces дээр ажиллах Gradio апп
=========================================================================
Энэ файлыг шинэ HF Space (Gradio SDK, CPU Basic — үнэгүй) дээр байрлуулна.
Хэрэглэгч зөвхөн browser-оороо орж видео upload хийнэ — тооцоолол бүгд
HF-ийн серверт (эсвэл холбогдсон cloud API-д) хийгдэнэ. Локал компьютер
дээр юу ч суулгах, ажиллуулах шаардлагагүй.

Space Settings > Variables and secrets хэсэгт дараах Secret-үүдийг нэм:
  GROQ_API_KEY
  HF_TOKEN
=========================================================================
"""

import os
import json
import time
import math
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path

import gradio as gr

GROQ_API_KEY = (os.environ.get("GROQ_API_KEY") or "").strip()
GEMINI_API_KEY = (os.environ.get("GEMINI_API_KEY") or "").strip()
HF_TOKEN = (os.environ.get("HF_TOKEN") or "").strip()

HF_SPACE = "ezensait/OmniVoice"  # Чиний өөрийн (CPU, quota-гүй) Space


def _check_keys():
    """Space эхлэхэд Logs дээр key-үүд зөв ачаалагдсан эсэхийг харуулна
    (утгыг нь бус, зөвхөн сүүлийн 4 тэмдэгтийг харуулна — аюулгүй)."""
    for name, val in [("GROQ_API_KEY", GROQ_API_KEY),
                       ("HF_TOKEN", HF_TOKEN)]:
        if not val or not val.strip():
            print(f"[ТУЛГАМДСАН АСУУДАЛ] {name} хоосон эсвэл олдсонгүй! "
                  f"Settings > Variables and secrets хэсэгт шалгаад Restart хий.")
        else:
            print(f"[OK] {name} ачаалагдлаа (төгсгөл: ...{val.strip()[-4:]}, урт: {len(val.strip())})")

    if GEMINI_API_KEY:
        print(f"[OK] GEMINI_API_KEY ачаалагдлаа (сайн чанарын орчуулга идэвхтэй)")
    else:
        print(f"[МЭДЭЭЛЭЛ] GEMINI_API_KEY байхгүй — зөвхөн Groq-оор орчуулна (ажиллана, чанар арай доогуур)")


_check_keys()


@dataclass
class Segment:
    start: float
    end: float
    text_src: str
    text_mn: str = ""
    audio_path: str = ""

    @property
    def duration(self):
        return self.end - self.start


# ---------------------------------------------------------------------
# Pipeline функцууд
# ---------------------------------------------------------------------

def extract_audio(video_path: str, out_path: str) -> str:
    # 48kbps mono mp3 — яриа таних чанарт хангалттай, файлын хэмжээг эрс багасгана
    # (30 минут ~55MB WAV байсныг ~11MB болгоно — Groq-ийн 25MB хязгаарт багтана)
    subprocess.run([
        "ffmpeg", "-y", "-i", video_path,
        "-ar", "16000", "-ac", "1", "-b:a", "48k", "-vn", out_path
    ], check=True, capture_output=True)
    return out_path


GROQ_MAX_BYTES = 24 * 1024 * 1024  # 24MB — Groq-ийн 25MB хязгаараас бага зэрэг доор, аюулгүй зай


def _transcribe_chunk(client, audio_path: str, time_offset: float) -> list[Segment]:
    with open(audio_path, "rb") as f:
        result = client.audio.transcriptions.create(
            file=(audio_path, f.read()),
            model="whisper-large-v3",
            response_format="verbose_json",
            timestamp_granularities=["segment"],
        )
    return [Segment(start=s["start"] + time_offset, end=s["end"] + time_offset,
                     text_src=s["text"].strip())
            for s in result.segments]


def transcribe_with_timestamps(audio_path: str) -> list[Segment]:
    from groq import Groq
    client = Groq(api_key=GROQ_API_KEY)

    file_size = os.path.getsize(audio_path)
    if file_size <= GROQ_MAX_BYTES:
        segments = _transcribe_chunk(client, audio_path, time_offset=0.0)
        print(f"[Transcribe] {len(segments)} сегмент танигдлаа (1 дор)")
        return segments

    # Хэт том файл — 15 минутын хэсгүүдэд хуваагаад тус тусад нь таниулна
    print(f"[Transcribe] Файл {file_size/1e6:.1f}MB > 24MB — хэсэг хэсгээр таних горимд шилжинэ")
    from pydub import AudioSegment
    full_audio = AudioSegment.from_file(audio_path)
    total_ms = len(full_audio)
    chunk_ms = 15 * 60 * 1000

    all_segments = []
    for start_ms in range(0, total_ms, chunk_ms):
        end_ms = min(start_ms + chunk_ms, total_ms)
        chunk_path = audio_path + f".chunk_{start_ms}.mp3"
        full_audio[start_ms:end_ms].export(chunk_path, format="mp3", bitrate="48k")
        chunk_segments = _transcribe_chunk(client, chunk_path, time_offset=start_ms / 1000.0)
        all_segments.extend(chunk_segments)
        print(f"[Chunk {start_ms/1000:.0f}-{end_ms/1000:.0f}s] {len(chunk_segments)} сегмент танигдлаа")

    return all_segments


# Session-ийн турш Gemini-ийн эрх дүүрсэн эсэхийг санаж, дахин дэмий хандахгүй байх
_gemini_exhausted = False


def translate_segments(segments: list[Segment], progress) -> list[Segment]:
    """Gemini-г үндсэн орчуулагчаар (чанар сайн), Groq-г нөөц (fallback) болгож ашиглана.
    GEMINI_API_KEY байхгүй/quota дүүрсэн бол автоматаар Groq руу шилжинэ,
    видео зогсохгүй. Gemini нэг удаа quota дүүрвэл, энэ Space дахин
    Restart хийх хүртэл дахин оролдохгүй — цаг хэмнэнэ."""
    global _gemini_exhausted
    from groq import Groq
    groq_client = Groq(api_key=GROQ_API_KEY)

    gemini_client = None
    if GEMINI_API_KEY and not _gemini_exhausted:
        try:
            from google import genai as google_genai
            gemini_client = google_genai.Client(api_key=GEMINI_API_KEY)
        except Exception as e:
            print(f"[Gemini clientинициализаци амжилтгүй] {e}")

    BATCH_SIZE = 25  # том багц — давтагдах prompt-ийн overhead-ийг багасгаж токен хэмнэнэ
    done = 0

    def try_groq_model(model_name: str, prompt: str, max_retries: int = 3):
        """Тухайн Groq загвараар орчуулна. Өдрийн token хязгаар (TPD) цохигдвол
        дахин оролдохгүй, шууд None буцааж дараагийн загвар руу шилжинэ
        (учир нь TPD хэдэн секундэд арилдаггүй, хүлээх нь дэмий цаг үрнэ)."""
        for attempt in range(max_retries):
            try:
                response = groq_client.chat.completions.create(
                    model=model_name,
                    messages=[{"role": "user", "content": prompt}],
                    response_format={"type": "json_object"},
                )
                text = response.choices[0].message.content.strip()
                parsed = json.loads(text)
                return parsed["items"] if isinstance(parsed, dict) else parsed
            except Exception as e:
                err_str = str(e)
                if "per day" in err_str or "TPD" in err_str:
                    print(f"[Groq {model_name}] Өдрийн token хязгаар дүүрсэн, дараагийн загвар руу шилжинэ: {e}")
                    return None
                wait = 10 * (attempt + 1)
                print(f"[Groq {model_name} дахин оролдлого {attempt+1}/{max_retries}] {e} — {wait} сек хүлээнэ")
                time.sleep(wait)
        return None

    for batch_start in range(0, len(segments), BATCH_SIZE):
        batch = segments[batch_start: batch_start + BATCH_SIZE]
        items = [
            {"id": i, "duration": round(seg.duration, 1), "text": seg.text_src}
            for i, seg in enumerate(batch)
        ]
        prompt = (
            "You are a professional Mongolian dubbing translator working on a movie/video. "
            "Translate each item's \"text\" into natural, idiomatic SPOKEN Mongolian (Cyrillic). "
            "Rules:\n"
            "1. Never translate word-for-word — rephrase so it sounds like something a Mongolian "
            "person would actually say out loud, not a mechanical translation.\n"
            "2. Preserve the speaker's tone and emotion (casual, excited, urgent, sarcastic, etc.).\n"
            "3. For names, slang, acronyms, or culture-specific terms with no direct Mongolian "
            "equivalent, use a natural, commonly understood Mongolian rendering — never invent "
            "a garbled or meaningless word.\n"
            "4. Every translation must be fully understandable on its own to a Mongolian listener.\n"
            "5. Prioritize natural, well-composed Mongolian phrasing. Do not shorten or pad the "
            "translation to match any particular length — say it the way a Mongolian speaker "
            "naturally would, regardless of how long or short that is.\n\n"
            "Respond with a JSON object of the form {\"items\": [{\"id\": ..., \"mn\": ...}, ...]}, "
            "same order, no other text, no markdown fences.\n\n"
            f"{json.dumps(items, ensure_ascii=False)}"
        )

        results = None

        # 1-р оролдлого: Gemini (чанар сайн)
        if gemini_client:
            try:
                response = gemini_client.models.generate_content(
                    model="gemini-2.5-flash", contents=prompt,
                )
                text = response.text.strip().replace("```json", "").replace("```", "").strip()
                parsed = json.loads(text)
                results = parsed["items"] if isinstance(parsed, dict) else parsed
                print(f"[Батч {batch_start}] Gemini-ээр орчуулагдлаа")
            except Exception as e:
                err_str = str(e)
                if "RESOURCE_EXHAUSTED" in err_str or "429" in err_str:
                    _gemini_exhausted = True
                    print(f"[Gemini quota дүүрлээ] Энэ session-д цаашид Gemini-г алгасна: {e}")
                else:
                    print(f"[Gemini амжилтгүй, Groq руу шилжинэ] {e}")

        # 2-р оролдлого: Groq том загвар
        if results is None:
            results = try_groq_model("llama-3.3-70b-versatile", prompt)
            if results is not None:
                print(f"[Батч {batch_start}] Groq (70b)-оор орчуулагдлаа")

        # 3-р оролдлого: Groq жижиг загвар (тусдаа token квот, эцсийн нөөц)
        if results is None:
            results = try_groq_model("llama-3.1-8b-instant", prompt)
            if results is not None:
                print(f"[Батч {batch_start}] Groq (8b, нөөц)-оор орчуулагдлаа")

        if results is None:
            raise gr.Error("Орчуулга бүрэн амжилтгүй боллоо (Gemini, Groq 70b, Groq 8b гурав)."
                            " Өнөөдрийн үнэгүй эрх дүүрсэн байж магадгүй, маргааш дахин турьж үзнэ үү.")

        for item in results:
            batch[item["id"]].text_mn = item["mn"].strip()

        done += len(batch)
        progress(0.25 + 0.3 * (done / len(segments)), desc=f"Орчуулж байна {done}/{len(segments)}")

    return segments


def generate_voice_segments(segments: list[Segment], work_dir: Path,
                             use_clone: bool, ref_audio_path: str, ref_text: str,
                             voice_prompt: str, voice_type: str, progress) -> list[Segment]:
    from gradio_client import Client, handle_file
    from concurrent.futures import ThreadPoolExecutor, as_completed
    import threading
    import shutil

    # Хоолойны төрлийг Gender(param_9) + Age(param_10) + Pitch(param_11) хослуулж
    # тогтмол нэг профайл болгоно — зөвхөн Gender-ийг бэхлээд Age/Pitch-ийг Auto
    # орхивол тэдгээр нь сегмент бүрд санамсаргүй өөрчлөгдөж "холимог хоолой" гардаг
    # байсныг chi screenshot-оор баталгаажуулсан тул, 3-г хамт бэхэлнэ.
    voice_map = {
        "Автомат": {"gender": "Auto", "age": "Auto", "pitch": "Auto"},
        "Эрэгтэй (гүн дуутай)": {"gender": "Male", "age": "Middle-aged", "pitch": "Low Pitch"},
        "Эрэгтэй (нарийн)": {"gender": "Male", "age": "Young Adult", "pitch": "High Pitch"},
        "Эмэгтэй": {"gender": "Female", "age": "Young Adult", "pitch": "Moderate Pitch"},
    }
    voice_profile = voice_map.get(voice_type, voice_map["Автомат"])
    gender_val = voice_profile["gender"]
    age_val = voice_profile["age"]
    pitch_val = voice_profile["pitch"]

    client = Client(HF_SPACE, token=HF_TOKEN)

    # DIAGNOSTIC: энэ Space-ийн бодит API endpoint-уудыг Logs дээр хэвлэнэ.
    try:
        api_info = client.view_api(print_info=False, return_format="dict")
        named = list(api_info.get("named_endpoints", {}).keys())
        print(f"[OmniVoice бодит endpoint-ууд] {named}")
        for ep_name, ep_info in api_info.get("named_endpoints", {}).items():
            param_names = [p.get("parameter_name", p.get("label", "?"))
                            for p in ep_info.get("parameters", [])]
            print(f"[{ep_name} параметрүүд] {param_names}")
    except Exception as e:
        named = []
        print(f"[Endpoint шалгалт амжилтгүй] {e}")

    lock = threading.Lock()
    counter = {"done": 0}

    def process_one(i: int, seg: Segment):
        clean_text = (seg.text_mn or "").strip()
        if not clean_text or clean_text.startswith(("[", "(")):
            print(f"[Segment {i} алгаслаа] хоосон/дуут бус текст: {clean_text!r}")
            return

        out_path = str(work_dir / f"seg_{i:04d}.wav")
        # Хугацааг ЭХ хэлнээс биш, МОНГОЛ орчуулгын үгийн тооноос тооцоолно —
        # ингэснээр загвар байгалиан, тогтмол хэмнэлээр ярина (хурдасгах/удаашруулах шаардлагагүй)
        word_count = len(clean_text.split())
        natural_du = max(1, round(word_count / 2.2))  # ~2.2 үг/сек — байгалиан ярианы хурд

        try:
            if use_clone and ref_audio_path:
                result = client.predict(
                    text=seg.text_mn,
                    lang="Mongolian",
                    ref_aud=handle_file(ref_audio_path),
                    ref_text=ref_text,
                    ns=10,
                    du=natural_du,
                    api_name="/_clone_fn",
                )
            else:
                result = client.predict(
                    text=seg.text_mn,
                    lang="Mongolian",
                    ns=10,
                    du=natural_du,
                    param_9=gender_val,
                    param_10=age_val,
                    param_11=pitch_val,
                    api_name="/_design_fn",
                )
        except ValueError as e:
            if "api_name" in str(e):
                raise gr.Error(f"api_name таарсангүй. Бодит endpoint-ууд: {named}")
            print(f"[Segment {i} алдаа, алгаслаа] {e}")
            return
        except Exception as e:
            print(f"[Segment {i} загварын алдаа, алгаслаа] {e}")
            return

        generated_audio = result[0] if isinstance(result, (list, tuple)) else result
        if not generated_audio:
            status_msg = result[1] if isinstance(result, (list, tuple)) and len(result) > 1 else "тодорхойгүй"
            print(f"[Segment {i} аудиогүй, алгаслаа] Загварын мессеж: {status_msg}")
            return

        shutil.copy(generated_audio, out_path)
        seg.audio_path = out_path

        with lock:
            counter["done"] += 1
            progress(counter["done"] / len(segments),
                      desc=f"Дуу хоолой үүсгэж байна {counter['done']}/{len(segments)}")

    # 2 vCPU-д тохируулж 2 сегментийг зэрэг боловсруулна
    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = [executor.submit(process_one, i, seg) for i, seg in enumerate(segments)]
        for f in as_completed(futures):
            f.result()  # api_name гэх мэт критик алдааг дахин шидэж гаргана

    return segments


def assemble_track(segments: list[Segment], total_duration: float, out_path: str) -> str:
    from pydub import AudioSegment

    # Урьдчилан бодит нийт урттай track бэлдэнэ (доор шаардлагатай бол сунгана)
    track = AudioSegment.silent(duration=int(total_duration * 1000))
    next_free_ms = 0  # өмнөх clip хаана дуусснаас хойш л дараагийнх эхэлнэ (давхцахгүй)

    for seg in segments:
        if not seg.audio_path:
            continue  # энэ сегмент алгассан (хоосон текст эсвэл загварын алдаа) — чимээгүй үлдэнэ
        clip = AudioSegment.from_wav(seg.audio_path)

        # Эх timestamp-ыг зөөлөн чиглүүлэгч болгож ашиглана, гэхдээ хэзээ ч
        # өмнөх clip дуусахаас өмнө эхлүүлэхгүй (давхцал үүсэхгүй) — хугацааг шахахгүй,
        # зөвхөн шаардлагатай бол дараагийн клипийг арай хожуу эхлүүлнэ.
        start_ms = max(int(seg.start * 1000), next_free_ms)

        if start_ms + len(clip) > len(track):
            track += AudioSegment.silent(duration=(start_ms + len(clip) - len(track)))

        track = track.overlay(clip, position=start_ms)
        next_free_ms = start_ms + len(clip)

    track.export(out_path, format="wav")
    return out_path


def merge_audio_to_video(video_path: str, audio_path: str, out_path: str, keep_background: bool):
    if keep_background:
        # Эх дууг бага түвшинд (дэвсгэр) хадгалж, шинэ хоолойтой холино
        subprocess.run([
            "ffmpeg", "-y", "-i", video_path, "-i", audio_path,
            "-filter_complex",
            "[0:a]volume=0.18[bg];[1:a]volume=1.0[voice];"
            "[bg][voice]amix=inputs=2:duration=longest:dropout_transition=0[aout]",
            "-map", "0:v:0", "-map", "[aout]",
            "-c:v", "copy", "-shortest", out_path
        ], check=True, capture_output=True)
    else:
        # Эх дууг бүрэн сольж, зөвхөн шинэ хоолойг ашиглана
        subprocess.run([
            "ffmpeg", "-y", "-i", video_path, "-i", audio_path,
            "-c:v", "copy", "-map", "0:v:0", "-map", "1:a:0",
            "-shortest", out_path
        ], check=True, capture_output=True)


# ---------------------------------------------------------------------
# Gradio-ийн гол функц
# ---------------------------------------------------------------------

def split_video(video_file, chunk_minutes):
    if video_file is None:
        raise gr.Error("Эхлээд видео файл оруулна уу.")
    if not chunk_minutes or chunk_minutes < 1:
        chunk_minutes = 8

    work_dir = Path(tempfile.mkdtemp(prefix="split_"))
    out_pattern = str(work_dir / "part_%03d.mp4")

    # -c copy: дахин encode хийхгүй тул маш хурдан (зөвхөн container таслана)
    subprocess.run([
        "ffmpeg", "-y", "-i", video_file,
        "-c", "copy", "-map", "0",
        "-f", "segment", "-segment_time", str(int(chunk_minutes * 60)),
        "-reset_timestamps", "1",
        out_pattern
    ], check=True, capture_output=True)

    parts = sorted(str(p) for p in work_dir.glob("part_*.mp4"))
    if not parts:
        raise gr.Error("Хуваахад алдаа гарлаа — видео формат тохирохгүй байж болзошгүй.")
    return parts, f"{len(parts)} хэсэг болгож хуваалаа. Тус бүрийг доороос татаж аваад, Dub tab-руу нэг нэгээр нь оруулна уу."


def dub_video(video_file, use_clone, ref_audio, ref_text, voice_prompt, voice_type, keep_background,
              progress=gr.Progress()):
    if video_file is None:
        raise gr.Error("Эхлээд видео файл оруулна уу.")

    work_dir = Path(tempfile.mkdtemp(prefix="mongol_dub_"))
    video_path = video_file

    progress(0.05, desc="Audio задалж байна...")
    audio_path = extract_audio(video_path, str(work_dir / "extracted.mp3"))

    from pydub.utils import mediainfo
    total_duration = float(mediainfo(video_path)["duration"])

    progress(0.15, desc="Эх хэлээр текст болгож байна...")
    segments = transcribe_with_timestamps(audio_path)
    if not segments:
        raise gr.Error("Ярианы текст олдсонгүй. Видеонд дуу байгаа эсэхийг шалгана уу.")

    progress(0.25, desc="Орчуулж эхэлж байна...")
    segments = translate_segments(segments, progress)

    progress(0.55, desc="Дуу хоолой үүсгэж эхэлж байна...")
    segments = generate_voice_segments(
        segments, work_dir, use_clone, ref_audio, ref_text, voice_prompt, voice_type, progress
    )

    progress(0.9, desc="Track угсарч байна...")
    final_audio = assemble_track(segments, total_duration, str(work_dir / "final_audio.wav"))

    progress(0.95, desc="Видео дээр нэгтгэж байна...")
    out_path = str(work_dir / "dubbed_output.mp4")
    merge_audio_to_video(video_path, final_audio, out_path, keep_background)

    script_text = "\n".join(f"[{s.start:.1f}-{s.end:.1f}] {s.text_mn}" for s in segments)

    progress(1.0, desc="Дууслаа!")
    return out_path, script_text


# ---------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------

with gr.Blocks(title="Монгол Дуб — KinoEzen") as demo:
    gr.Markdown("# 🎬 Монгол Дуб\nГадаад хэлтэй видеог автоматаар Монгол хэлээр дуб хийнэ.")

    with gr.Tabs():
        with gr.Tab("1️⃣ Видео хуваах (урт бол эхлээд энд)"):
            gr.Markdown(
                "Урт видео (10+ минут) бол эхлээд энд жижиг хэсэгт хуваа — "
                "тус бүр 10-20 минутанд дуусах тул холболт тасрах эрсдэл багасна."
            )
            with gr.Row():
                with gr.Column():
                    split_input = gr.File(label="Видео оруулах (mp4)", type="filepath",
                                           file_types=[".mp4", ".mkv", ".mov", ".avi"])
                    chunk_minutes = gr.Number(label="Хэдэн минут тутамд хуваах вэ?", value=8, minimum=1)
                    split_btn = gr.Button("Хуваах", variant="primary")
                with gr.Column():
                    split_output = gr.File(label="Хуваагдсан хэсгүүд (татаж авах)", file_count="multiple")
                    split_status = gr.Textbox(label="Төлөв")

            split_btn.click(
                fn=split_video,
                inputs=[split_input, chunk_minutes],
                outputs=[split_output, split_status],
            )

        with gr.Tab("2️⃣ Дуб хийх"):
            with gr.Row():
                with gr.Column():
                    video_input = gr.File(label="Видео оруулах (mp4)", type="filepath",
                                           file_types=[".mp4", ".mkv", ".mov", ".avi"])
                    use_clone = gr.Checkbox(label="Өөрийн клон хоолойгоор ярих уу?", value=False)
                    ref_audio = gr.File(label="Reference audio (клон ашиглах бол)", type="filepath",
                                         file_types=[".wav", ".mp3"])
                    ref_text = gr.Textbox(label="Reference audio-нд хэлсэн үг (клон ашиглах бол)")
                    voice_prompt = gr.Textbox(
                        label="Preset хоолойн тодорхойлолт (клон ашиглахгүй бол, одоогоор ашиглагдахгүй)",
                        value="male, adult, calm, neutral accent",
                        visible=False,
                    )
                    voice_type = gr.Radio(
                        label="Хоолойны төрөл (клон ашиглахгүй бол)",
                        choices=["Автомат", "Эрэгтэй (гүн дуутай)", "Эрэгтэй (нарийн)", "Эмэгтэй"],
                        value="Эрэгтэй (гүн дуутай)",
                    )
                    keep_background = gr.Checkbox(
                        label="Эх бичлэгийн дэвсгэр дуу/хөгжмийг хадгалах уу?",
                        value=True,
                    )
                    submit_btn = gr.Button("Дуб хийж эхлэх", variant="primary")

                with gr.Column():
                    video_output = gr.File(label="Бэлэн видео (татаж авах)")
                    script_output = gr.Textbox(label="Монгол скрипт", lines=15)

            submit_btn.click(
                fn=dub_video,
                inputs=[video_input, use_clone, ref_audio, ref_text, voice_prompt, voice_type, keep_background],
                outputs=[video_output, script_output],
            )

demo.launch()