Update app.py

app.py

@@ -7,6 +7,8 @@ import tempfile
 import torch
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 import spaces
+import pytube
+import librosa
 
 def extract_audio(video_path):
     video = mp.VideoFileClip(video_path)
@@ -75,19 +77,33 @@ def process_video(video_path, target_language):
 
     return output_path
 
-def gradio_interface(video, target_language):
-    output_video = process_video(video.name, target_language)
+def download_youtube_video(youtube_link):
+    yt = pytube.YouTube(youtube_link)
+    video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
+    video_path = video.download(output_path=tempfile.gettempdir())
+    return video_path
+
+def gradio_interface(video, yt_link, target_language):
+    if video is not None:
+        video_path = video.name
+    elif yt_link:
+        video_path = download_youtube_video(yt_link)
+    else:
+        raise ValueError("Please provide either a video file or a YouTube link.")
+    
+    output_video = process_video(video_path, target_language)
     return output_video
 
 iface = gr.Interface(
     fn=gradio_interface,
     inputs=[
         gr.Video(label="Upload Video"),
+        gr.Textbox(label="YouTube Link"),
         gr.Dropdown(choices=["es", "fr", "de", "it", "ja", "ko", "zh-CN"], label="Target Language")
     ],
     outputs=gr.Video(label="Processed Video"),
     title="Video Subtitle Translator",
-    description="Upload a video, and get it back with translated subtitles!"
+    description="Upload a video or provide a YouTube link, and get it back with translated subtitles!"
 )
 
 iface.launch()