add youtube feature

app.py

@@ -33,6 +33,8 @@ from huggingface_hub import snapshot_download
 
 from gradio import Markdown
 
+from pytube import YouTube
+
 all_key_names = ['C major', 'G major', 'D major', 'A major',
                  'E major', 'B major', 'F major', 'Bb major',
                  'Eb major', 'Ab major', 'Db major', 'Gb major',
@@ -709,6 +711,29 @@ def gradio_generate(input_video, input_primer, input_key):
     output_filename = video2music.generate(input_video, input_primer, input_key)
     return str(output_filename)
 
+def gradio_generate2(input_youtube, input_primer, input_key):
+
+    youtube_dir = Path("./youtube")
+    if youtube_dir.exists():
+        shutil.rmtree(str(youtube_dir))
+    
+    youtube_dir.mkdir(parents=True)
+
+    yObject = YouTube(input_youtube)
+    yObject_stream = yObject.streams.get_by_resolution("240p")
+    fname = yObject.video_id +".mp4"
+    if yObject_stream == None:
+        yObject_stream = yObject.streams.get_lowest_resolution()
+    try:
+        yObject_stream.download(output_path=youtube_dir, filename= fname )
+    except:
+        print("An error has occurred")
+    
+    input_video = youtube_dir / fname
+
+    output_filename = video2music.generate(input_video, input_primer, input_key)
+    return str(output_filename)
+
 
 title="Video2Music: Suitable Music Generation from Videos using an Affective Multimodal Transformer model"
 description_text = """
@@ -718,10 +743,15 @@ Generate background music using Video2Music by providing an input video.
 <a href="https://arxiv.org/abs/2311.00968">Read our paper.</a>
 <p/>
 """
-input_video = gr.Video(label="Input Video")
-input_primer = gr.Textbox(label="Input Primer", value="C Am F G")
-input_key = gr.Dropdown(choices=["C major", "A minor"], value="C major", label="Input Key")
-output_video = gr.Video(label="Output Video")
+
+
+# input_video = gr.Video(label="Input Video")
+# input_primer = gr.Textbox(label="Input Primer", value="C Am F G")
+# input_key = gr.Dropdown(choices=["C major", "A minor"], value="C major", label="Input Key")
+# output_video = gr.Video(label="Output Video")
+# input_youtube = gr.Textbox(label="YouTube URL")
+
+
 
 css = '''
 #duplicate-button {
@@ -732,21 +762,86 @@ border-radius: 100vh;
 }
 '''
 
-# Gradio interface
-gr_interface = gr.Interface(
-    fn=gradio_generate,
-    inputs=[input_video, input_primer, input_key ],
-    outputs=[output_video],
-    description=description_text,
-    allow_flagging='never',
-    cache_examples=True,
-)
 
 
+# Gradio interface
+# gr_interface = gr.Interface(
+#     fn=gradio_generate,
+#     inputs=[input_video, input_primer, input_key ],
+#     outputs=[output_video],
+#     description=description_text,
+#     allow_flagging='never',
+#     cache_examples=True,
+# )
+
+# gr_interface2 = gr.Interface(
+#     fn=gradio_generate2,
+#     inputs=[input_youtube, input_primer, input_key ],
+#     outputs=[output_video],
+#     description=description_text,
+#     allow_flagging='never',
+#     cache_examples=True,
+# )
+
+def filter(choice):
+    if choice == "Upload Video":
+        return [gr.update(visible=True), gr.update(visible=False)]
+    if choice == "YouTube URL":
+        return [gr.update(visible=False), gr.update(visible=True)]
+    
+
 # with gr.Blocks() as demo:
 with gr.Blocks(css=css) as demo:
     title=gr.HTML(f"<h1><center>{title}</center></h1>")
-    gr_interface.render()
+    gr.Markdown(
+            """
+            This is the demo for Video2Music: Suitable Music Generation from Videos using an Affective Multimodal Transformer model.
+            [Read our paper](https://arxiv.org/abs/2311.00968).
+            """
+    )
+    with gr.Row():
+        with gr.Column():
+            radio = gr.Radio(["Upload Video", "YouTube URL"], value="Upload Video", label = "Choose the input method")
+            with gr.Row(visible=True) as mainA:
+                with gr.Column(visible=True) as colA:
+                    with gr.Row(visible=True) as rowA:
+
+                        with gr.Column():
+                            input_video = gr.Video(label="Input Video", height=250)
+                        with gr.Column():
+                            with gr.Row():
+                                input_primer = gr.Textbox(label="Input Primer", value="C Am F G")
+                                input_key = gr.Dropdown(choices=["C major", "A minor"], value="C major", label="Input Key")
+                        with gr.Column():
+                            btn = gr.Button("Generate")
+
+                    with gr.Row(visible=False) as rowB:
+                        with gr.Column():
+                            input_video_yt = gr.Textbox(label="YouTube URL")
+                        with gr.Column():
+                            with gr.Row():
+                                input_primer_yt = gr.Textbox(label="Input Primer", value="C Am F G")
+                                input_key_yt = gr.Dropdown(choices=["C major", "A minor"], value="C major", label="Input Key")
+                        with gr.Column():
+                            btn_yt = gr.Button("Generate")
+
+        with gr.Column():
+            with gr.Row():
+                output_video = gr.Video(label="Output Video")
+
+    radio.change(filter, radio, [rowA, rowB])
+
+    btn.click(
+        fn=gradio_generate, 
+        inputs=[input_video,input_primer,input_key],
+        outputs=[output_video],
+    )
+
+    btn_yt.click(
+        fn=gradio_generate2, 
+        inputs=[input_video_yt,input_primer_yt,input_key_yt],
+        outputs=[output_video],
+    )
 
 #demo.queue()
 # demo.launch(debug=True)