add gradio UI

poc_app.py

@@ -1,25 +1,31 @@
 # external imports
 import time
+import gradio as gr
 
 # local imports
 from blip_image_caption_large import Blip_Image_Caption_Large
 from phi3_mini_4k_instruct import Phi3_Mini_4k_Instruct
 from musicgen_small import Musicgen_Small
 
-def main():
+#image_to_music function
+def image_to_music(image_path):
     # test image captioning
     image_caption_start_time = time.time()
     image_caption_model = Blip_Image_Caption_Large()
-    test_caption = image_caption_model.caption_image_local_pipeline("data/test3.jpg")
+    
+    test_caption = image_caption_model.caption_image_local_pipeline(image_path)
+    
     print(test_caption)
+    
     image_caption_end_time = time.time()
 
     # test text generation
     text_generation_start_time = time.time()
     text_generation_model = Phi3_Mini_4k_Instruct()
-
+    
     #TODO: move this to a config file
-    text_generation_model.local_pipeline.model.config.max_length = 200
+    text_generation_model.local_pipeline.model.config.max_new_tokens = 200
+
 
     #TODO: move system prompt somewhere else, allow for genre override
     messages = [
@@ -44,10 +50,31 @@ def main():
     music_generation_duration = music_generation_end_time - music_generation_start_time
     total_duration = music_generation_end_time - image_caption_start_time
 
-    # output durations
-    print(f"Image Captioning Duration: {image_caption_duration}")
-    print(f"Text Generation Duration: {text_generation_duration}")
-    print(f"Music Generation Duration: {music_generation_duration}")
-    print(f"Total Duration: {total_duration}")
+    # output generated_text, audio and duration to gradio
+    return (test_caption[0]["generated_text"], test_text[-1]['generated_text'][-1]['content'], "data/musicgen_out.wav",
+            f"Image Captioning Duration: {image_caption_duration} sec",
+            f"Text Generation Duration: {text_generation_duration} sec",
+            f"Music Generation Duration: {music_generation_duration} sec",
+            f"Total Duration: {total_duration} sec")
+
+# Gradio UI 
+def gradio():
+    # Define Gradio Interface, information from (https://www.gradio.app/docs/chatinterface)
+    with gr.Blocks() as demo:
+        gr.Markdown("<h1 style='text-align: center;'> ⛺ Image to Music Generator 🎼</h1>")
+        image_input = gr.Image(type="filepath", label="Upload Image")
+        with gr.Row():
+            caption_output = gr.Textbox(label="Image Caption")
+            music_description_output = gr.Textbox(label="Music Description")
+            durations = gr.Textbox(label="Processing Times", interactive=False, placeholder="Time statistics will appear here")
+
+        music_output = gr.Audio(label="Generated Music")
+        # Button to trigger the process
+        generate_button = gr.Button("Generate Music")
+
+        generate_button.click(fn=image_to_music, inputs=[image_input], outputs=[caption_output, music_description_output, music_output, durations])
+
+    # Launch Gradio app
+    demo.launch()
 
-main()
+gradio()