Update run.py

run.py

@@ -1,33 +1,58 @@
 import gradio as gr
-import subprocess
+from yt_dlp import YoutubeDL
 import os
 
-def download_media(url):
-    # Determine if the URL is for audio or video
-    ydl_opts = {
-        'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',  # Prefer mp4 format
-        'outtmpl': 'downloaded_media.%(ext)s',  # Output file name
-    }
-
-    # Download using yt-dlp
-    subprocess.run(['yt-dlp', url, '--merge-output-format', 'mp4', '--recode-video', 'mp4'])
-
-    # Check if downloaded file exists
-    if os.path.exists('downloaded_media.mp4'):
-        return gr.outputs.Video("downloaded_media.mp4")
-    elif os.path.exists('downloaded_media.m4a'):
-        return gr.outputs.Audio("downloaded_media.m4a")
-    else:
-        return "Media could not be downloaded or unsupported format."
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=download_media,
-    inputs="text",
-    outputs="auto",
-    title="YouTube Downloader",
-    description="Enter a YouTube video URL to download its audio or video."
-)
-
-# Launch the interface
-iface.launch()
+def download_youtube(url, output_format):
+    ydl_opts = {}
+    file_path = ""
+    
+    if output_format == "audio":
+        ydl_opts = {
+            'format': 'bestaudio/best',
+            'postprocessors': [{
+                'key': 'FFmpegExtractAudio',
+                'preferredcodec': 'mp3',
+                'preferredquality': '192',
+            }],
+            'outtmpl': 'downloads/%(title)s.%(ext)s',
+        }
+    elif output_format == "video":
+        ydl_opts = {
+            'format': 'bestvideo+bestaudio/best',
+            'outtmpl': 'downloads/%(title)s.%(ext)s',
+        }
+
+    with YoutubeDL(ydl_opts) as ydl:
+        result = ydl.extract_info(url, download=True)
+        file_path = ydl.prepare_filename(result)
+
+        if output_format == "audio":
+            file_path = file_path.replace(result['ext'], 'mp3')
+
+    return file_path
+
+def show_media(file_path, output_format):
+    if output_format == "audio":
+        return gr.Audio.update(value=file_path)
+    elif output_format == "video":
+        return gr.Video.update(value=file_path)
+
+def yt_download(url, output_format):
+    file_path = download_youtube(url, output_format)
+    return show_media(file_path, output_format)
+
+with gr.Blocks() as demo:
+    gr.Markdown("## YouTube Downloader")
+    
+    with gr.Row():
+        url_input = gr.Textbox(label="YouTube URL")
+        format_input = gr.Radio(choices=["audio", "video"], label="Format")
+
+    download_btn = gr.Button("Download")
+
+    audio_output = gr.Audio(label="Audio Output")
+    video_output = gr.Video(label="Video Output")
+
+    download_btn.click(yt_download, inputs=[url_input, format_input], outputs=[audio_output, video_output])
+
+demo.launch()