Update app.py

app.py

@@ -6,6 +6,10 @@ import pandas as pd
 import requests
 import io
 from transformers import MarianMTModel, MarianTokenizer
+from gradio_client import Client
+
+# Initialize Gradio Client for Whisper JAX
+client = Client(src="sanchit-gandhi/whisper-jax")
 
 def fetch_languages(url):
     response = requests.get(url)
@@ -19,70 +23,15 @@ def fetch_languages(url):
     else:
         return []
 
-def text_to_srt(text):
-    lines = text.split('\n')
-    srt_content = ""
-    for i, line in enumerate(lines):
-        if line.strip() == "":
-            continue
-        try:
-            times, content = line.split(']', 1)
-            start, end = times[1:].split(' -> ')
-            if start.count(":") == 1:
-                start = "00:" + start
-            if end.count(":") == 1:
-                end = "00:" + end
-            srt_content += f"{i+1}\n{start.replace('.', ',')} --> {end.replace('.', ',')}\n{content.strip()}\n\n"
-        except ValueError:
-            continue
-    temp_file_path = '/tmp/output.srt'
-    with open(temp_file_path, 'w', encoding='utf-8') as file:
-        file.write(srt_content)
-    return temp_file_path
-
-def translate_text(text, source_language_code, target_language_code):
-    model_name = f"Helsinki-NLP/opus-mt-{source_language_code}-{target_language_code}"
-    try:
-        tokenizer = MarianTokenizer.from_pretrained(model_name)
-        model = MarianMTModel.from_pretrained(model_name)
-    except Exception as e:
-        return f"Failed to load model for {source_language_code} to {target_language_code}: {str(e)}"
-    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512))
-    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
-    return translated_text
-
-def translate_srt(input_file, source_language_code, target_language_code):
-    subs = pysrt.open(input_file)
-    for sub in subs:
-        sub.text = translate_text(sub.text, source_language_code, target_language_code)
-    translated_srt_path = input_file.replace(".srt", f"_{target_language_code}.srt")
-    subs.save(translated_srt_path)
-    return translated_srt_path
-
-def add_subtitle_to_video(input_video, subtitle_file, soft_subtitle=True):
-    video_input_stream = ffmpeg.input(input_video)
-    subtitle_input_stream = ffmpeg.input(subtitle_file)
-    input_video_name = os.path.splitext(os.path.basename(input_video))[0]
-    output_video = f"/tmp/{input_video_name}_subtitled.mp4"
-
-    if soft_subtitle:
-        stream = ffmpeg.output(
-            video_input_stream, subtitle_input_stream, output_video,
-            **{"c": "copy", "c:s": "mov_text"}
-        )
-    else:
-        stream = ffmpeg.output(
-            video_input_stream, output_video,
-            vf=f"subtitles={subtitle_file}"
-        )
-
-    ffmpeg.run(stream, overwrite_output=True)
-    return output_video
+def transcript_audio(audio_file, task, return_timestamps, api_name="/predict_1"):
+    prediction = client.predict(audio_file=audio_file, task=task, return_timestamps=return_timestamps, api_name=api_name)
+    return prediction['transcription'], prediction['transcription_time_s']
 
-def process_video(input_video, text_transcription, video_language, target_language):
-    srt_path = text_to_srt(text_transcription)
+def process_video(input_video, video_language, target_language):
+    transcription, _ = transcript_audio(input_video, "transcribe", True)
+    srt_path = text_to_srt(transcription)
     translated_srt_path = translate_srt(srt_path, video_language, target_language)
-    output_video = add_subtitle_to_video(input_video.name, translated_srt_path)
+    output_video = add_subtitle_to_video(input_video, translated_srt_path)
     return output_video
 
 language_url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
@@ -91,11 +40,11 @@ video_language_options = fetch_languages(language_url)
 with gr.Blocks() as app:
     with gr.Row():
         input_video = gr.Video(label="Video File")
-        text_transcription = gr.TextArea(label="Text Transcription")
         video_language = gr.Dropdown(choices=video_language_options, label="Language of the Video")
         target_language = gr.Dropdown(choices=video_language_options, label="Language Translated")
+        process_btn = gr.Button("Process Video")
     output_video = gr.Video(label="Video with Translated Subtitles")
 
-    input_video.change(fn=process_video, inputs=[input_video, text_transcription, video_language, target_language], outputs=output_video)
+    process_btn.click(fn=process_video, inputs=[input_video, video_language, target_language], outputs=output_video)
 
 app.launch()