Update app.py

app.py

@@ -1,4 +1,3 @@
-import time
 from threading import Thread
 from llava_llama3.serve.cli import chat_llava
 from llava_llama3.model.builder import load_pretrained_model
@@ -8,6 +7,7 @@ from PIL import Image
 import argparse
 import spaces
 import os
+import time
 
 root_path = os.path.dirname(os.path.abspath(__file__))
 print(root_path)
@@ -29,7 +29,8 @@ tokenizer, llava_model, image_processor, context_len = load_pretrained_model(
     'llava_llama3', 
     args.load_8bit, 
     args.load_4bit, 
-    device=args.device)
+    device=args.device
+)
 
 @spaces.GPU
 def bot_streaming(message, history):
@@ -53,8 +54,9 @@ def bot_streaming(message, history):
     if image is None:
         raise gr.Error("You need to upload an image for LLaVA to work.")
     
-    # Load the image
-    image = Image.open(image)
+    # Load the image if it's a path, otherwise use the existing PIL image
+    if isinstance(image, str):
+        image = Image.open(image).convert('RGB')
     
     # Generate the prompt for the model
     prompt = message['text']
@@ -62,8 +64,6 @@ def bot_streaming(message, history):
     # Use a streamer to generate the output in a streaming fashion
     streamer = []
 
-    image_file = image if isinstance(image, str) else image.filename
-
     # Define a function to call chat_llava in a separate thread
     def generate_output():
         output = chat_llava(
@@ -98,20 +98,19 @@ def bot_streaming(message, history):
         yield buffer
 
 
-chatbot=gr.Chatbot(scale=1)
+chatbot = gr.Chatbot(scale=1)
 chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
-with gr.Blocks(fill_height=True, ) as demo:
+with gr.Blocks(fill_height=True) as demo:
     gr.ChatInterface(
-    fn=bot_streaming,
-    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,
+        fn=bot_streaming,
+        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,
     )
 
 demo.queue(api_open=False)
-demo.launch(show_api=False, share=False)
+demo.launch(show_api=False, share=False)