Update app.py

app.py

@@ -1,47 +1,43 @@
 import gradio as gr
-from pytube import YouTube
 from moviepy.audio.io.AudioFileClip import AudioFileClip
+import yt_dlp
 import os
 
 def download_youtube(url, output_format, resolution, audio_bitrate):
     try:
-        yt = YouTube(url)
-        title = yt.title.replace(" ", "_").replace("/", "_")
-        
-        # Set output filename and path
-        output_filename = f"{title}.{output_format.lower()}"
-        output_path = f"downloads/{output_filename}"
-
-        # Get stream
-        if output_format == "MP4":
-            stream = yt.streams.filter(progressive=True, file_extension="mp4", res=resolution).first()
-        else:  # Audio
-            stream = yt.streams.filter(only_audio=True).order_by("abr").desc().first()
-
-        if not stream:
-            return "Selected quality not available", None
-
-        # Download
-        downloaded_path = stream.download(output_path="downloads", filename=title + ".mp4")
-
-        if output_format in ["MP3", "WAV"]:
-            audio = AudioFileClip(downloaded_path)
-            if output_format == "MP3":
-                final_path = f"downloads/{title}.mp3"
-                audio.write_audiofile(final_path, bitrate=audio_bitrate)
-            else:
-                final_path = f"downloads/{title}.wav"
-                audio.write_audiofile(final_path)
-            audio.close()
-            os.remove(downloaded_path)
-        else:
-            final_path = downloaded_path
-
-        return f"Downloaded: {output_format}", final_path
+        output_dir = "downloads"
+        os.makedirs(output_dir, exist_ok=True)
+        output_template = f"{output_dir}/%(title)s.%(ext)s"
+
+        ydl_opts = {
+            "format": "bestaudio/best" if output_format in ["MP3", "WAV"] else f"bestvideo[height<={resolution[:-1]}]+bestaudio/best",
+            "outtmpl": output_template,
+            "postprocessors": [],
+        }
+
+        if output_format == "MP3":
+            ydl_opts["postprocessors"].append({
+                'key': 'FFmpegExtractAudio',
+                'preferredcodec': 'mp3',
+                'preferredquality': audio_bitrate.replace("k", ""),
+            })
+        elif output_format == "WAV":
+            ydl_opts["postprocessors"].append({
+                'key': 'FFmpegExtractAudio',
+                'preferredcodec': 'wav',
+            })
+
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info_dict = ydl.extract_info(url, download=True)
+            filename = ydl.prepare_filename(info_dict)
+            if output_format in ["MP3", "WAV"]:
+                filename = os.path.splitext(filename)[0] + f".{output_format.lower()}"
+            return "Download complete", filename
 
     except Exception as e:
         return f"Error: {str(e)}", None
 
+
 formats = ["MP3", "WAV", "MP4"]
 resolutions = ["360p", "480p", "720p", "1080p"]
 bitrates = ["64k", "128k", "192k", "256k", "320k"]