support adding youtube

app.py

@@ -2,6 +2,11 @@ import gradio as gr
 import requests
 import uuid
 import os
+from typing import Optional
+import tempfile
+from pydub import AudioSegment
+import re
+
 ASR_API = "http://astarwiz.com:9998/asr"
 TTS_SPEAK_SERVICE = 'http://astarwiz.com:9603/speak'
 TTS_WAVE_SERVICE = 'http://astarwiz.com:9603/wave'
@@ -16,6 +21,70 @@ LANGUAGE_MAP = {
 # Add a password for developer mode
 DEVELOPER_PASSWORD = os.getenv("DEV_PWD")
 
+# Add this constant for the RapidAPI key
+RAPID_API_KEY = os.getenv("RAPID_API_KEY")
+
+def fetch_youtube_id(youtube_url: str) -> str:
+    if 'v=' in youtube_url:
+        return youtube_url.split("v=")[1]
+    elif 'shorts' in youtube_url:
+        return youtube_url.split("/")[-1]
+    else:
+        raise Exception("Unsupported URL format")
+
+def download_youtube_audio(youtube_url: str, output_dir: Optional[str] = None) -> Optional[str]:
+    video_id = fetch_youtube_id(youtube_url)
+    
+    if not video_id:
+        return None
+
+    if output_dir is None:
+        output_dir = tempfile.gettempdir()
+
+    output_filename = os.path.join(output_dir, f"{video_id}.mp3")
+    
+    if os.path.exists(output_filename):
+        return output_filename  # Return if the file already exists
+    
+    url = "https://youtube86.p.rapidapi.com/api/youtube/links"
+    headers = {
+        'Content-Type': 'application/json',
+        'x-rapidapi-host': 'youtube86.p.rapidapi.com',
+        'x-rapidapi-key': RAPID_API_KEY
+    }
+    data = {
+        "url": youtube_url
+    }
+    
+    response = requests.post(url, headers=headers, json=data)
+    print('Fetched audio links')
+    
+    if response.status_code == 200:
+        result = response.json()
+        for url in result[0]['urls']:
+            if url.get('isBundle'):
+                audio_url = url['url']
+                extension = url['extension']
+                audio_response = requests.get(audio_url)
+                
+                if audio_response.status_code == 200:
+                    temp_filename = os.path.join(output_dir, f"{video_id}.{extension}")
+                    with open(temp_filename, 'wb') as audio_file:
+                        audio_file.write(audio_response.content)
+                    
+                    # Convert to MP3 and downsample to 16000 Hz
+                    audio = AudioSegment.from_file(temp_filename, format=extension)
+                    audio = audio.set_frame_rate(16000)
+                    audio.export(output_filename, format="mp3", parameters=["-ar", "16000"])
+                    
+                    os.remove(temp_filename)  # Remove the temporary file
+                    return output_filename  # Return the final MP3 filename
+        
+        return None  # Return None if no successful download occurs
+    else:
+        print("Error:", response.status_code, response.text)
+        return None  # Return None on failure
+
 def inference_via_llm_api(input_text, min_new_tokens=2, max_new_tokens=64):
     print(input_text)
     one_vllm_input = f"<|im_start|>system\nYou are a translation expert.<|im_end|>\n<|im_start|>user\n{input_text}<|im_end|>\n<|im_start|>assistant"
@@ -37,9 +106,14 @@ def inference_via_llm_api(input_text, min_new_tokens=2, max_new_tokens=64):
     else:
         return "The system got some error during vLLM generation. Please try it again."
     
-def transcribe_and_speak(audio, source_lang, target_lang):
+def transcribe_and_speak(audio, source_lang, target_lang, youtube_url=None):
+    if youtube_url:
+        audio = download_youtube_audio(youtube_url)
+        if not audio:
+            return "Failed to download YouTube audio.", None, None
+
     if not audio:
-        return "Please provide an audio input.", None, None
+        return "Please provide an audio input or a valid YouTube URL.", None, None
 
     # ASR
     file_id = str(uuid.uuid4())
@@ -79,18 +153,19 @@ def transcribe_and_speak(audio, source_lang, target_lang):
 def check_password(password):
     return password == DEVELOPER_PASSWORD
 
-def user_interface(audio, source_lang, target_lang):
-    _, _, audio_url = transcribe_and_speak(audio, source_lang, target_lang)
-    return audio_url
+def user_interface(audio, source_lang, target_lang, youtube_url):
+    transcription, translated_text, audio_url = transcribe_and_speak(audio, source_lang, target_lang, youtube_url)
+    return transcription, translated_text, audio_url
 
 with gr.Blocks() as demo:
     gr.Markdown("# Speech Translation")
     
     with gr.Tab("User Mode"):
-        gr.Markdown("Speak into the microphone or upload an audio file. The app will translate and speak it back to you.")
+        gr.Markdown("Speak into the microphone, upload an audio file, or provide a YouTube URL. The app will translate and speak it back to you.")
         
         with gr.Row():
             user_audio_input = gr.Audio(sources=["microphone", "upload"], type="filepath")
+            user_youtube_url = gr.Textbox(label="YouTube URL (optional)")
             user_source_lang = gr.Dropdown(choices=["en", "ma", "ta", "zh"], label="Source Language", value="en")
             user_target_lang = gr.Dropdown(choices=["en", "ma", "ta", "zh"], label="Target Language", value="zh")
         
@@ -98,12 +173,14 @@ with gr.Blocks() as demo:
             user_button = gr.Button("Translate and Speak")
         
         with gr.Row():
+            user_transcription_output = gr.Textbox(label="Transcription")
+            user_translation_output = gr.Textbox(label="Translation")
             user_audio_output = gr.Audio(label="Translated Speech")
         
         user_button.click(
             fn=user_interface,
-            inputs=[user_audio_input, user_source_lang, user_target_lang],
-            outputs=[user_audio_output]
+            inputs=[user_audio_input, user_source_lang, user_target_lang, user_youtube_url],
+            outputs=[user_transcription_output, user_translation_output, user_audio_output]
         )
     
     with gr.Tab("Developer Mode"):