app.py

import os
import tempfile
import gradio as gr
import numpy as np

# Lazy imports to make startup faster on Spaces CPU
ASR_MODEL = None
EN2HI = None
HI2EN = None
TTS_ENGINE = None
XTTS_AVAILABLE = False

# --- Utilities ---

def lazy_import_asr():
    global ASR_MODEL
    if ASR_MODEL is None:
        import whisper
        # small is ok for 15–60s clips on CPU. change to "medium" if GPU.
        ASR_MODEL = whisper.load_model("small")
    return ASR_MODEL


def lazy_import_mt():
    global EN2HI, HI2EN
    if EN2HI is None or HI2EN is None:
        from transformers import pipeline
        # MarianMT models (offline, no API). Good enough for EN<>HI.
        EN2HI = pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi")
        HI2EN = pipeline("translation", model="Helsinki-NLP/opus-mt-hi-en")
    return EN2HI, HI2EN


def lazy_import_tts(prefer_xtts=True):
    """
    Try XTTS v2 (multi-lingual + cloning) if available, else fall back to a small CPU TTS.
    Fallback voices are non-cloned.
    """
    global TTS_ENGINE, XTTS_AVAILABLE
    if TTS_ENGINE is not None:
        return TTS_ENGINE, XTTS_AVAILABLE

    try:
        # Coqui TTS
        from TTS.api import TTS
        if prefer_xtts:
            # XTTS v2 supports speaker embedding + multilingual
            TTS_ENGINE = TTS("tts_models/multilingual/multi-dataset/xtts_v2")
            XTTS_AVAILABLE = True
        else:
            raise RuntimeError("skip xtts")
    except Exception:
        try:
            # Lightweight English-only fallback (CPU friendly)
            from TTS.api import TTS
            TTS_ENGINE = TTS("tts_models/en/vctk/vits")
            XTTS_AVAILABLE = False
        except Exception as e:
            raise RuntimeError(f"No TTS backend could be loaded: {e}")

    return TTS_ENGINE, XTTS_AVAILABLE


def extract_audio(video_path, out_wav):
    import moviepy.editor as mp
    clip = mp.VideoFileClip(video_path)
    audio = clip.audio
    audio.write_audiofile(out_wav, fps=16000, codec="pcm_s16le")
    clip.close()
    return out_wav


def transcribe(video_path, language_hint="auto"):
    """ASR using Whisper. Returns (text, lang_code)."""
    asr = lazy_import_asr()
    with tempfile.TemporaryDirectory() as td:
        wav = os.path.join(td, "audio.wav")
        extract_audio(video_path, wav)
        result = asr.transcribe(wav, language=None if language_hint == "auto" else language_hint)
        text = result.get("text", "").strip()
        lang = result.get("language", "unknown")
    return text, lang


def translate_text(text, target_lang):
    """Use MarianMT for EN<>HI only. If source==target, return original."""
    if not text.strip():
        return ""
    en2hi, hi2en = lazy_import_mt()
    tl = target_lang.lower()

    # naive detection: if Devanagari present, treat as Hindi
    is_devanagari = any(0x0900 <= ord(ch) <= 0x097F for ch in text)

    if tl in ["hi", "hindi"]:
        if is_devanagari:
            return text
        # assume english -> hindi
        out = en2hi(text)[0]["translation_text"]
        return out
    elif tl in ["en", "english"]:
        if not is_devanagari:
            return text
        out = hi2en(text)[0]["translation_text"]
        return out
    else:
        # unsupported target -> return original
        return text


def synthesize_tts(text, lang="en", ref_audio=None):
    """
    text -> wav_path. If XTTS available and ref_audio provided + consent, clone voice.
    """
    tts, xtts = lazy_import_tts(prefer_xtts=True)

    tmp_wav = tempfile.mktemp(suffix=".wav")

    if xtts:
        # XTTS supports many langs: 'en', 'hi', etc.
        speaker_wav = ref_audio if (ref_audio and os.path.exists(ref_audio)) else None
        tts.tts_to_file(text=text, language=lang, speaker_wav=speaker_wav, file_path=tmp_wav)
    else:
        # fallback is English-only; if lang not en, we just attempt anyway
        tts.tts_to_file(text=text, file_path=tmp_wav)

    return tmp_wav


def merge_audio(video_path, dubbed_wav, out_mp4):
    import moviepy.editor as mp
    video = mp.VideoFileClip(video_path)
    audio = mp.AudioFileClip(dubbed_wav).set_duration(video.duration)
    final = video.set_audio(audio)
    # write with AAC audio so browsers play fine
    final.write_videofile(out_mp4, audio_codec="aac", codec="libx264")
    video.close()
    audio.close()
    final.close()
    return out_mp4


# --- Gradio UI callbacks ---

def full_dub(video, target_lang, language_hint, use_cloning, consent, ref_voice=None):
    """
    End-to-end pipeline: ASR -> (optional) Translate -> TTS -> Merge
    Returns: transcribed_text, translated_text, dubbed_video
    """
    if not consent:
        raise gr.Error("Please confirm consent for dubbing/cloning.")

    if video is None:
        raise gr.Error("Please upload a video (15–60 seconds recommended).")

    # 1) ASR
    text, detected = transcribe(video, language_hint=language_hint)

    # 2) Translate if needed
    translated = translate_text(text, target_lang)

    # 3) TTS
    # choose TTS language code
    lang_code = "en" if target_lang.lower().startswith("en") else "hi"

    # If no translation occurred (unsupported), speak original text
    tts_text = translated if translated.strip() else text

    # Only pass ref_voice if user ticked cloning
    ref_path = None
    if use_cloning and ref_voice is not None:
        # Coqui XTTS expects a file path
        ref_path = ref_voice

    out_wav = synthesize_tts(tts_text, lang=lang_code, ref_audio=ref_path)

    # 4) Merge
    out_mp4 = tempfile.mktemp(suffix=".mp4")
    merge_audio(video, out_wav, out_mp4)

    # Small watermark text for transparency
    mark = f"Dubbed ({'cloned' if use_cloning else 'synthetic'}) | Lang: {target_lang.upper()}"

    return text, translated, out_mp4, mark


# --- Build UI ---
with gr.Blocks(theme=gr.themes.Soft(), css=".small {font-size: 12px}") as demo:
    gr.Markdown("""
    # 🎬 Safe Video Dubbing (EN ⇄ HI) — CPU Friendly
    **Hinglish UI**: Neeche steps follow karo. 15–60s video pe best results.

    - ASR: Whisper-small (offline)
    - Translation: MarianMT (offline)
    - TTS: Tries **XTTS v2** (voice cloning) → agar available nahin hua toh English-only fallback TTS use karega.
    - Output: MP4 with replaced audio

    **Note (Safety)**: Sirf wahi awaaz clone/ dub karein jiske liye aapke paas clear consent ho. App ek chhota watermark lagata hai transparency ke liye.
    """)

    with gr.Row():
        video_in = gr.Video(label="Input Video (mp4/mov)")
        ref_voice = gr.Audio(label="(Optional) Reference Voice for Cloning — same person ki 5–10s WAV/MP3", type="filepath")

    with gr.Row():
        target_lang = gr.Radio(["en", "hi"], value="hi", label="Target Language (Output)")
        language_hint = gr.Dropdown(["auto", "en", "hi"], value="auto", label="ASR Language Hint")

    with gr.Row():
        use_cloning = gr.Checkbox(False, label="Enable Voice Cloning (XTTS v2)")
        consent = gr.Checkbox(False, label="I confirm consent for dubbing / cloning (required)")

    with gr.Row():
        btn = gr.Button("▶️ Run Dubbing")

    with gr.Row():
        transcribed = gr.Textbox(label="Transcribed Text")
        translated = gr.Textbox(label="Translated Text")

    dubbed = gr.Video(label="Dubbed Output (MP4)")
    watermark = gr.Markdown(visible=True)

    btn.click(fn=full_dub, inputs=[video_in, target_lang, language_hint, use_cloning, consent, ref_voice],
              outputs=[transcribed, translated, dubbed, watermark])

    gr.Markdown("""
    ---
    **Tips (Hinglish):**
    - Agar Space CPU pe slow hai, 15–30 sec clips try karo.
    - XTTS cloning ke liye GPU recommended. HF Space hardware settings me GPU select kar sakte ho.
    - Hindi ke liye XTTS best; fallback TTS English-only ho sakta hai.
    - Clean audio = better results. Background noise kam rakho.
    """)

if __name__ == "__main__":
    demo.launch()