Update app.py

app.py

@@ -4,6 +4,7 @@ from llava_llama3.model.builder import load_pretrained_model
 from PIL import Image
 import torch
 
+# Model configuration
 model_path = "TheFinAI/FinLLaVA"
 device = "cuda"
 conv_mode = "llama_3"
@@ -12,6 +13,7 @@ max_new_tokens = 512
 load_8bit = False
 load_4bit = False
 
+# Load the pretrained model
 tokenizer, llava_model, image_processor, context_len = load_pretrained_model(
     model_path, 
     None, 
@@ -21,7 +23,8 @@ tokenizer, llava_model, image_processor, context_len = load_pretrained_model(
     device=device
 )
 
-def predict(image, text):
+# Define the prediction function
+def predict(image, text, history):
     output = chat_llava(
         args=None,
         image_file=image,
@@ -31,23 +34,26 @@ def predict(image, text):
         image_processor=image_processor,
         context_len=context_len
     )
-    return output
+    history.append((text, output))
+    return history, gr.update(value="")
 
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot(label="FinLLaVA Chatbot")
+    image_input = gr.Image(type="filepath", label="Upload Image")
+    text_input = gr.Textbox(label="Enter your message")
+    submit_btn = gr.Button("Submit")
 
-chatbot = gr.Chatbot(scale=1)
-chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
+    # Define interaction: when submit is clicked, call predict and update the chatbot
+    submit_btn.click(fn=predict, inputs=[image_input, text_input, chatbot], outputs=[chatbot, text_input])
 
-with gr.Blocks(fill_height=True) as demo:
-    gr.ChatInterface(
-        fn=predict,
-        title="FinLLaVA",
-        examples=[{"text": "What is on the flower?", "files": ["./bee.jpg"]},
-                  {"text": "How to make this pastry?", "files": ["./baklava.png"]}],
-        stop_btn="Stop Generation",
-        multimodal=True,
-        textbox=chat_input,
-        chatbot=chatbot,
+    # Add example inputs
+    gr.Examples(
+        examples=[["./bee.jpg", "What is on the flower?"],
+                  ["./baklava.png", "How to make this pastry?"]],
+        inputs=[image_input, text_input]
     )
 
+# Launch the Gradio app
 demo.queue(api_open=False)
 demo.launch(show_api=False, share=False)