modify app

app.py

@@ -40,6 +40,29 @@ def loudness_normalize(audio, sample_rate, target_loudness=-12.0):
     loudness_normalized_audio = pyln.normalize.loudness(audio, loudness, target_loudness)
     return loudness_normalized_audio
 
+def process_youtube_url(url):
+    try:
+        audio, sr = download_youtube_audio(url)
+        return (sr, audio)
+    except Exception as e:
+        return None, f"Error processing YouTube URL: {str(e)}"
+
+def process_audio_with_youtube(input_audio, input_youtube_url, reference_audio, reference_youtube_url):
+    if input_youtube_url:
+        input_audio, error = process_youtube_url(input_youtube_url)
+        if error:
+            return None, None, error
+    
+    if reference_youtube_url:
+        reference_audio, error = process_youtube_url(reference_youtube_url)
+        if error:
+            return None, None, error
+    
+    if input_audio is None or reference_audio is None:
+        return None, None, "Both input and reference audio are required."
+    
+    return process_audio(input_audio, reference_audio)
+
 def process_audio(input_audio, reference_audio):
     output_audio, predicted_params, sr = mastering_transfer.process_audio(
         input_audio, reference_audio, reference_audio
@@ -138,6 +161,34 @@ with gr.Blocks() as demo:
             outputs=[output_audio, param_output]
         )
 
+    with gr.Tab("YouTube Audio"):
+        with gr.Row():
+            input_audio_yt = gr.Audio(label="Input Audio (Optional)")
+            input_youtube_url = gr.Textbox(label="Input YouTube URL (Optional)")
+        with gr.Row():
+            reference_audio_yt = gr.Audio(label="Reference Audio (Optional)")
+            reference_youtube_url = gr.Textbox(label="Reference YouTube URL (Optional)")
+        
+        process_button_yt = gr.Button("Process Mastering Style Transfer")
+        
+        with gr.Row():
+            output_audio_yt = gr.Audio(label="Output Audio", type='numpy')
+            param_output_yt = gr.Textbox(label="Predicted Parameters", lines=5)
+        
+        error_message_yt = gr.Textbox(label="Error Message", visible=False)
+
+        def process_and_handle_errors(input_audio, input_youtube_url, reference_audio, reference_youtube_url):
+            result = process_audio_with_youtube(input_audio, input_youtube_url, reference_audio, reference_youtube_url)
+            if len(result) == 3 and isinstance(result[2], str):  # Error occurred
+                return None, None, gr.update(visible=True, value=result[2])
+            return result[0], result[1], gr.update(visible=False, value="")
+
+        process_button_yt.click(
+            process_and_handle_errors, 
+            inputs=[input_audio_yt, input_youtube_url, reference_audio_yt, reference_youtube_url], 
+            outputs=[output_audio_yt, param_output_yt, error_message_yt]
+        )
+
     gr.Markdown("## Inference Time Optimization (ITO)")
     
     with gr.Row():