Update app.py

app.py

@@ -326,9 +326,14 @@ button:hover {
 }
 """
 
-def download_audio(url, cookie_file):
+import os
+import yt_dlp
+from googledrivedownloader import GoogleDriveDownloader
+from scipy.io import wavfile
+
+def download_audio(url, cookie_file=None):
     ydl_opts = {
-        'format': 'bestaudio[ext=webm]/bestaudio[ext=m4a]/bestaudio[ext=opus]/bestaudio[ext=aac]/bestaudio -video',  # Prioritize audio-only formats, exclude video
+        'format': 'bestaudio[ext=webm]/bestaudio[ext=m4a]/bestaudio[ext=opus]/bestaudio[ext=aac]/bestaudio -video',
         'postprocessors': [{
             'key': 'FFmpegExtractAudio',
             'preferredcodec': 'wav',
@@ -347,29 +352,43 @@ def download_audio(url, cookie_file):
         'verbose': True,
     }
 
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+    # Check if it's a Google Drive link
+    if 'drive.google.com' in url or 'https://drive.google.com' in url:
         try:
-            info_dict = ydl.extract_info(url, download=True)
-            # Adjust file path for multiple possible extensions
-            base_file_path = ydl.prepare_filename(info_dict)
-            file_path = base_file_path
-            for ext in ['.webm', '.m4a', '.opus', '.aac']:
-                file_path = file_path.replace(ext, '.wav')
-            if not os.path.exists(file_path):
+            file_id = url.split('/d/')[1].split('/')[0]  # Extract file ID
+            output_path = 'ytdl/gdrive_audio.wav'
+            GoogleDriveDownloader.download_file_from_google_drive(file_id=file_id, dest_path=output_path, unzip=False)
+            if not os.path.exists(output_path):
                 return None, "Downloaded file not found", None
-            # Read the audio file for gr.Audio
-            sample_rate, data = scipy.io.wavfile.read(file_path)
+            sample_rate, data = wavfile.read(output_path)
             audio_data = (sample_rate, data)
-            return file_path, "Download successful", audio_data
-        except yt_dlp.utils.ExtractorError as e:
-            if "Sign in to confirm you’re not a bot" in str(e) or "The provided YouTube account cookies are no longer valid" in str(e):
-                return None, "Authentication failed. Please upload updated cookies from a logged-in browser session in the respective tab. See https://github.com/yt-dlp/yt-dlp/wiki/Extractors#exporting-youtube-cookies for instructions.", None
-            elif "HTTP Error 403: Forbidden" in str(e):
-                return None, "Download failed: HTTP Error 403. This format requires a GVS PO Token, or the cookies are invalid. Please upload fresh cookies. See https://github.com/yt-dlp/yt-dlp/wiki/PO-Token-Guide for advanced troubleshooting.", None
-            return None, f"Download failed: {str(e)}", None
+            return output_path, "Download successful", audio_data
         except Exception as e:
-            return None, f"Unexpected error: {str(e)}", None
+            return None, f"Google Drive download failed: {str(e)}", None
 
+    # Handle YouTube link
+    else:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            try:
+                info_dict = ydl.extract_info(url, download=True)
+                base_file_path = ydl.prepare_filename(info_dict)
+                file_path = base_file_path
+                for ext in ['.webm', '.m4a', '.opus', '.aac']:
+                    file_path = file_path.replace(ext, '.wav')
+                if not os.path.exists(file_path):
+                    return None, "Downloaded file not found", None
+                sample_rate, data = wavfile.read(file_path)
+                audio_data = (sample_rate, data)
+                return file_path, "Download successful", audio_data
+            except yt_dlp.utils.ExtractorError as e:
+                if "Sign in to confirm you’re not a bot" in str(e) or "The provided YouTube account cookies are no longer valid" in str(e):
+                    return None, "Authentication failed. Please upload updated cookies from a logged-in browser session in the respective tab. See https://github.com/yt-dlp/yt-dlp/wiki/Extractors#exporting-youtube-cookies for instructions.", None
+                elif "HTTP Error 403: Forbidden" in str(e):
+                    return None, "Download failed: HTTP Error 403. This format requires a GVS PO Token, or the cookies are invalid. Please upload fresh cookies. See https://github.com/yt-dlp/yt-dlp/wiki/PO-Token-Guide for advanced troubleshooting.", None
+                return None, f"Download failed: {str(e)}", None
+            except Exception as e:
+                return None, f"Unexpected error: {str(e)}", None
+                
 @spaces.GPU
 def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, pitch_shift, model_dir, output_dir, out_format, norm_thresh, amp_thresh, batch_size, exclude_stems="", progress=gr.Progress(track_tqdm=True)):
     """Separate audio into stems using a Roformer model."""