Update app.py

app.py

@@ -319,8 +319,35 @@ button:hover {
 """
 
 def download_audio(url, cookie_file=None):
+    """
+    Downloads audio from YouTube or Google Drive and converts it to WAV format.
+    
+    Args:
+        url (str): URL of the YouTube video or Google Drive file.
+        cookie_file (file object): File object containing YouTube cookies in Netscape format.
+    
+    Returns:
+        tuple: (file_path, message, (sample_rate, data)) or (None, error_message, None)
+    """
+    # Common output directory
+    os.makedirs('ytdl', exist_ok=True)
+    
+    # Validate cookie file
+    cookie_path = None
+    if cookie_file:
+        if not hasattr(cookie_file, 'name') or not os.path.exists(cookie_file.name):
+            return None, "Invalid or missing cookie file", None
+        cookie_path = cookie_file.name
+        logger.info(f"Using cookie file: {cookie_path}")
+    
+    if 'drive.google.com' in url:
+        return download_from_google_drive(url)
+    else:
+        return download_from_youtube(url, cookie_path)
+
+def download_from_youtube(url, cookie_path):
     ydl_opts = {
-        'format': 'bestaudio[ext=webm]/bestaudio[ext=m4a]/bestaudio[ext=opus]/bestaudio[ext=aac]/bestaudio -video',
+        'format': 'bestaudio/best',  # More flexible format selection
         'postprocessors': [{
             'key': 'FFmpegExtractAudio',
             'preferredcodec': 'wav',
@@ -332,52 +359,66 @@ def download_audio(url, cookie_file=None):
         'force_ipv4': True,
         'referer': 'https://www.youtube.com/',
         'noplaylist': True,
-        'cookiefile': cookie_file.name if cookie_file else None,
+        'cookiefile': cookie_path,
         'extractor_retries': 5,
         'ignoreerrors': False,
         'no_check_certificate': True,
         'verbose': True,
     }
-    temp_output_path = None
+    
     try:
-        if 'drive.google.com' in url:
-            os.makedirs('ytdl', exist_ok=True)
-            file_id = url.split('/d/')[1].split('/')[0]
-            download_url = f'https://drive.google.com/uc?id={file_id}'
-            temp_output_path = 'ytdl/gdrive_temp_audio'
-            gdown.download(download_url, temp_output_path, quiet=False)
-            if not os.path.exists(temp_output_path):
-                return None, "Downloaded file not found", None
-            output_path = 'ytdl/gdrive_audio.wav'
-            try:
-                audio = AudioSegment.from_file(temp_output_path)
-                audio.export(output_path, format="wav")
-            except Exception as e:
-                return None, f"Failed to process Google Drive file as audio: {str(e)}. Ensure the file contains audio (e.g., MP3, WAV, or video with audio track).", None
-            sample_rate, data = scipy.io.wavfile.read(output_path)
-            return output_path, "Download and audio conversion successful", (sample_rate, data)
-        else:
-            os.makedirs('ytdl', exist_ok=True)
-            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-                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 = scipy.io.wavfile.read(file_path)
-                return file_path, "Download successful", (sample_rate, data)
-    except yt_dlp.utils.ExtractorError as e:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            info_dict = ydl.extract_info(url, download=True)
+            # Get the actual file path after download
+            file_path = ydl.prepare_filename(info_dict).rsplit('.', 1)[0] + '.wav'
+            
+            if not os.path.exists(file_path):
+                return None, "Downloaded audio file not found", None
+            
+            sample_rate, data = scipy.io.wavfile.read(file_path)
+            return file_path, "YouTube audio download and conversion successful", (sample_rate, data)
+    
+    except yt_dlp.utils.DownloadError as e:
         if "Sign in to confirm you’re not a bot" in str(e):
             return None, "Authentication error. Please upload valid YouTube cookies: https://github.com/yt-dlp/yt-dlp/wiki/Extractors#exporting-youtube-cookies", None
-        return None, f"Download error: {str(e)}", None
+        return None, f"YouTube download error: {str(e)}", None
+    except yt_dlp.utils.GeoRestrictedError:
+        return None, "Video is geo-restricted in your region", None
     except Exception as e:
-        return None, f"Unexpected error: {str(e)}", None
+        return None, f"Unexpected error during YouTube download: {str(e)}", None
+
+def download_from_google_drive(url):
+    temp_output_path = 'ytdl/gdrive_temp_audio'
+    output_path = 'ytdl/gdrive_audio.wav'
+    
+    try:
+        # Extract file ID from URL
+        file_id = url.split('/d/')[1].split('/')[0]
+        download_url = f'https://drive.google.com/uc?id={file_id}'
+        
+        # Download file
+        gdown.download(download_url, temp_output_path, quiet=False)
+        
+        if not os.path.exists(temp_output_path):
+            return None, "Google Drive downloaded file not found", None
+        
+        # Convert to WAV
+        audio = AudioSegment.from_file(temp_output_path)
+        audio.export(output_path, format="wav")
+        
+        sample_rate, data = scipy.io.wavfile.read(output_path)
+        return output_path, "Google Drive audio download and conversion successful", (sample_rate, data)
+    
+    except Exception as e:
+        return None, f"Failed to process Google Drive file: {str(e)}. Ensure the file contains audio (e.g., MP3, WAV, or video with audio track).", None
+    
     finally:
-        if temp_output_path and os.path.exists(temp_output_path):
-            os.remove(temp_output_path)
-            logger.info(f"Temporary file deleted: {temp_output_path}")
+        if os.path.exists(temp_output_path):
+            try:
+                os.remove(temp_output_path)
+                logger.info(f"Temporary file deleted: {temp_output_path}")
+            except Exception as e:
+                logger.warning(f"Failed to delete temporary file {temp_output_path}: {str(e)}")
 
 @spaces.GPU(duration=60)
 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)):