import gradio as gr
from pytube import YouTube
import whisper
from transformers import pipeline
from bark import generate_audio, preload_models
from moviepy import VideoFileClip
import os
import tempfile

# Load models
asr_model = whisper.load_model("base")
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi")
preload_models()

def process_video(link):
    try:
        # Download audio from YouTube
        yt = YouTube(link)
        audio_stream = yt.streams.filter(only_audio=True).first()
        temp_dir = tempfile.mkdtemp()
        audio_path = os.path.join(temp_dir, "audio.mp4")
        audio_stream.download(filename=audio_path)

        # Step 2: Transcribe using Whisper
        result = asr_model.transcribe(audio_path)
        english_text = result["text"]

        # Step 3: Translate to Hindi
        hindi_text = translator(english_text)[0]['translation_text']

        # Step 4: Generate Hindi audio using Bark
        hindi_audio = generate_audio(hindi_text)
        hindi_audio_path = os.path.join(temp_dir, "hindi.wav")
        with open(hindi_audio_path, "wb") as f:
            f.write(hindi_audio)

        # Step 5: Merge audio with video
        audio_clip = mp.AudioFileClip(hindi_audio_path)
        video_clip = mp.ColorClip(size=(1280, 720), color=(0, 0, 0), duration=audio_clip.duration)
        final_video = video_clip.set_audio(audio_clip)
        output_path = os.path.join(temp_dir, "final_output.mp4")
        final_video.write_videofile(output_path, fps=24)

        return hindi_text, output_path

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

gr.Interface(
    fn=process_video,
    inputs=gr.Textbox(label="YouTube Video Link"),
    outputs=[
        gr.Textbox(label="Hindi Translation"),
        gr.Video(label="Dubbed Hindi Video")
    ],
    title="🎙️ YouTube Hindi Dubber AI",
    description="Paste a YouTube video link. This app transcribes it, translates it to Hindi, and dubs it using Bark AI."
).launch()