Update app.py

app.py

@@ -13,6 +13,7 @@ import networkx as nx
 from collections import Counter
 import json
 from datetime import datetime
+from transformers import pipeline
 
 @dataclass
 class ChatMessage:
@@ -32,6 +33,8 @@ class XylariaChat:
             model="Qwen/QwQ-32B-Preview",
             token=self.hf_token
         )
+        
+        self.image_captioning_pipe = pipeline("image-text-to-text", model="llava-onevision-qwen2-0.5b-si-hf")
 
         self.image_api_url = "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-large"
         self.image_api_headers = {"Authorization": f"Bearer {self.hf_token}"}
@@ -401,6 +404,29 @@ class XylariaChat:
             print(f"Error resetting API client: {e}")
 
         return None
+    
+    def caption_image_llava(self, image_path, prompt):
+        try:
+            with open(image_path, "rb") as img_file:
+                image_data = base64.b64encode(img_file.read()).decode("utf-8")
+
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "url": image_path},
+                        {"type": "text", "text": prompt},
+                    ],
+                },
+            ]
+            
+            caption_result = self.image_captioning_pipe(text=messages, max_new_tokens=50)
+            
+            caption = caption_result[0]['generated_text'] if caption_result else "No caption generated"
+            
+            return caption
+        except Exception as e:
+            return f"Error captioning image with llava: {str(e)}"
 
     def caption_image(self, image):
         try:
@@ -484,8 +510,9 @@ class XylariaChat:
                 messages.append(msg)
 
             if image:
-                image_caption = self.caption_image(image)
-                user_input = f"description of an image: {image_caption}\n\nUser's message about it: {user_input}"
+                image_caption_prompt = "Can you describe this image for me?"
+                caption = self.caption_image_llava(image, image_caption_prompt)
+                user_input = f"Here is a description of an image: {caption}\n\n{user_input}"
 
             messages.append(ChatMessage(
                 role="user",