added youtube song import

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import tempfile
 import os 
 hf_token = os.environ.get('HF_TOKEN')
 
@@ -36,6 +37,67 @@ compel = Compel(
 
 from pydub import AudioSegment
 
+import yt_dlp as youtube_dl
+from moviepy.editor import VideoFileClip
+
+YT_LENGTH_LIMIT_S = 480  # limit to 1 hour YouTube files
+
+def download_yt_audio(yt_url, filename):
+    info_loader = youtube_dl.YoutubeDL()
+    
+    try:
+        info = info_loader.extract_info(yt_url, download=False)
+    except youtube_dl.utils.DownloadError as err:
+        raise gr.Error(str(err))
+    
+    file_length = info["duration_string"]
+    file_h_m_s = file_length.split(":")
+    file_h_m_s = [int(sub_length) for sub_length in file_h_m_s]
+    
+    if len(file_h_m_s) == 1:
+        file_h_m_s.insert(0, 0)
+    if len(file_h_m_s) == 2:
+        file_h_m_s.insert(0, 0)
+    file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2]
+    
+    if file_length_s > YT_LENGTH_LIMIT_S:
+        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
+        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
+        raise gr.Error(f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.")
+    
+    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
+    
+    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+        try:
+            ydl.download([yt_url])
+        except youtube_dl.utils.ExtractorError as err:
+            raise gr.Error(str(err))
+
+
+def convert_to_mp3(input_path, output_path):
+    try:
+        video_clip = VideoFileClip(input_path)
+        audio_clip = video_clip.audio
+        print("Converting to MP3...")
+        audio_clip.write_audiofile(output_path)
+    except Exception as e:
+        print("Error:", e)
+
+def load_youtube_audio(yt_link):
+    
+    gr.Info("Loading your YouTube link ... ")
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        filepath = os.path.join(tmpdirname, "video.mp4")
+        download_yt_audio(yt_link, filepath)
+
+        mp3_output_path = "video_sound.mp3"
+        convert_to_mp3(filepath, mp3_output_path)
+        print("Conversion complete. MP3 saved at:", mp3_output_path)
+    
+    
+        
+    return mp3_output_path
+
 def cut_audio(input_path, output_path, max_duration):
     audio = AudioSegment.from_file(input_path)
 
@@ -274,6 +336,9 @@ with gr.Blocks(css=css) as demo:
             </div>""")
         
         audio_input = gr.Audio(label="Music input", type="filepath", source="upload")
+        with gr.Row():
+            youtube_link = gr.Textbox(show_label=False, placeholder="you can also paste YT link and load it")
+            yt_load_btn = gr.Button("Load YT song")
         
         with gr.Row():
             has_lyrics = gr.Radio(label="Does your audio has lyrics ?", choices=["Yes", "No"], value="No", info="If yes, the image should reflect the lyrics, but be aware that because we add a step (getting lyrics), inference will take more time.")
@@ -333,6 +398,7 @@ with gr.Blocks(css=css) as demo:
         """)
 
     #infer_btn.click(fn=infer, inputs=[audio_input], outputs=[lpmc_cap, llama_trans_cap, img_result])
+    yt_load_btn.click(fn=load_youtube_audio, inputs=[youtube_link], outputs=[audio_input], queue=False)
     infer_btn.click(fn=infer, inputs=[audio_input, has_lyrics], outputs=[processed_audio, img_result, llama_trans_cap, tryagain_btn, share_group])
     share_button.click(None, [], [], _js=share_js)
     tryagain_btn.click(fn=solo_xd, inputs=[llama_trans_cap], outputs=[img_result])