Update app.py

app.py

@@ -19,6 +19,7 @@ import pytube as pt
 from pytube.exceptions import VideoUnavailable
 
 from inference.style_transfer import *
+from inference.mastering_transfer import *
 
 
 yt_video_dir = "./yt_dir/0"
@@ -56,9 +57,6 @@ def get_audio_from_yt_video_ref(yt_link: str, start_point_in_second=0, duration_
     return filename_ref, filename_ref
 
 def inference(file_uploaded_in, file_uploaded_ref):
-
-    output_wav_path = None
-
     # Perform music mixing style transfer
     args = set_up()
     
@@ -67,6 +65,15 @@ def inference(file_uploaded_in, file_uploaded_ref):
     
     return output_wav_path, output_wav_path
 
+def inference_mastering(file_uploaded_in, file_uploaded_ref):
+    # Perform music mastering style transfer
+    args = set_up()
+    
+    inference_mastering_style_transfer = Mastering_Style_Transfer_Inference(args)
+    output_wav_path_mastering = inference_mastering_style_transfer.inference(file_uploaded_in, file_uploaded_ref)
+    
+    return output_wav_path_mastering, output_wav_path_mastering
+
 
 with gr.Blocks() as demo:
     gr.HTML(
@@ -126,26 +133,40 @@ with gr.Blocks() as demo:
                         inputs=[yt_link_ref],
                         outputs=[yt_audio_path_ref, file_uploaded_ref],
                     )
-            with gr.Column():
-                inference_btn = gr.Button("Perform Style Transfer")
                 
     with gr.Group():
         gr.HTML(
             """
-            <div> <h3> <center> Mixing Style Transferred Output. </h3> </div>
+            <div> <h3> <center> Mixing Style Transfer. </h3> </div>
             """
         )
-        with gr.Row().style(mobile_collapse=False, equal_height=True):
-            output_mix_l = gr.Audio(label="Listen to output mix")
-            output_mix_p = gr.Audio(label="Download output mix")
+        with gr.Column():
+            inference_btn = gr.Button("Perform Style Transfer")
         with gr.Row():
-            output_mix = gr.File(label="Download style transferred music track")
+            output_mix = gr.Audio(label="mixing style transferred music track")
             inference_btn.click(
                 inference,
                 inputs=[file_uploaded_in, file_uploaded_ref],
-                outputs=[output_mix_l, output_mix_p],
+                outputs=[output_mix],
             )
 
+    
+    with gr.Group():
+        gr.HTML(
+            """
+            <div> <h3> <center> Mastering Style Transfer. </h3> </div>
+            """
+        )
+        with gr.Column():
+            inference_mastering_btn = gr.Button("Perform Mastering Style Transfer")
+        with gr.Row():
+            output_master = gr.Audio(label="mastering style transferred music track")
+            inference_mastering_btn.click(
+                inference_mastering,
+                inputs=[file_uploaded_in, file_uploaded_ref],
+                outputs=[output_master],
+            )
+            
 
     
 if __name__ == "__main__":