Update app.py

app.py

@@ -1,4 +1,6 @@
 import os
+import base64
+import requests
 import gradio as gr
 from huggingface_hub import InferenceClient
 
@@ -11,10 +13,14 @@ class XylariaChat:
         
         # Initialize the inference client
         self.client = InferenceClient(
-            model= os.getenv("MODEL_NAME"), 
+            model=os.getenv("MODEL_NAME"), 
             api_key=self.hf_token
         )
         
+        # Image captioning API setup
+        self.image_api_url = "https://api-inference.huggingface.co/models/microsoft/git-large-coco"
+        self.image_api_headers = {"Authorization": f"Bearer {self.hf_token}"}
+        
         # Initialize conversation history and persistent memory
         self.conversation_history = []
         self.persistent_memory = {}
@@ -22,6 +28,7 @@ class XylariaChat:
         # System prompt with more detailed instructions
         self.system_prompt = """You are a helpful and harmless assistant. You are Xylaria developed by Sk Md Saad Amin(india, 12 year old). You should think step-by-step. 
 """
+    
     def store_information(self, key, value):
         """Store important information in persistent memory"""
         self.persistent_memory[key] = value
@@ -51,12 +58,53 @@ class XylariaChat:
         
         return None  # To clear the chatbot interface
 
-    def get_response(self, user_input):
+    def caption_image(self, image):
+        """
+        Caption an uploaded image using Hugging Face API
+        
+        Args:
+            image (str): Base64 encoded image or file path
+        
+        Returns:
+            str: Image caption or error message
+        """
+        try:
+            # If image is a file path, read and encode
+            if isinstance(image, str) and os.path.isfile(image):
+                with open(image, "rb") as f:
+                    data = f.read()
+            # If image is already base64 encoded
+            elif isinstance(image, str):
+                # Remove data URI prefix if present
+                if image.startswith('data:image'):
+                    image = image.split(',')[1]
+                data = base64.b64decode(image)
+            # If image is a file-like object
+            else:
+                data = image.read()
+            
+            # Send request to Hugging Face API
+            response = requests.post(
+                self.image_api_url, 
+                headers=self.image_api_headers, 
+                data=data
+            )
+            
+            # Check response
+            if response.status_code == 200:
+                caption = response.json()[0].get('generated_text', 'No caption generated')
+                return caption
+            else:
+                return f"Error captioning image: {response.text}"
+        
+        except Exception as e:
+            return f"Error processing image: {str(e)}"
+
+    def get_response(self, user_input, image=None):
         # Prepare messages with conversation context and persistent memory
         messages = [
             {"role": "system", "content": self.system_prompt},
             *self.conversation_history,
-            {"role": "user", "content": user_input}
         ]
         
         # Add persistent memory context if available
@@ -66,6 +114,14 @@ class XylariaChat:
             )
             messages.insert(1, {"role": "system", "content": memory_context})
         
+        # Process image if uploaded
+        if image:
+            image_caption = self.caption_image(image)
+            user_input = f"Image description: {image_caption}\n\nUser's message: {user_input}"
+        
+        # Add user input
+        messages.append({"role": "user", "content": user_input})
+        
         # Generate response with streaming
         try:
             stream = self.client.chat.completions.create(
@@ -82,13 +138,13 @@ class XylariaChat:
             return f"Error generating response: {str(e)}"
 
     def create_interface(self):
-        def streaming_response(message, chat_history):
+        def streaming_response(message, chat_history, image):
             # Clear input textbox
-            response_stream = self.get_response(message)
+            response_stream = self.get_response(message, image)
             
             # If it's an error, return immediately
             if isinstance(response_stream, str):
-                return "", chat_history + [[message, response_stream]]
+                return "", chat_history + [[message, response_stream]], None
             
             # Prepare for streaming response
             full_response = ""
@@ -102,7 +158,7 @@ class XylariaChat:
                     
                     # Update the last message in chat history with partial response
                     updated_history[-1][1] = full_response
-                    yield "", updated_history
+                    yield "", updated_history, None
             
             # Update conversation history
             self.conversation_history.append(
@@ -144,30 +200,37 @@ class XylariaChat:
                     show_copy_button=True
                 )
                 
-                # Input row with improved layout
+                # Input row with improved layout and image upload
                 with gr.Row():
-                    txt = gr.Textbox(
-                        show_label=False, 
-                        placeholder="Type your message...", 
-                        container=False,
-                        scale=4
-                    )
+                    with gr.Column(scale=4):
+                        txt = gr.Textbox(
+                            show_label=False, 
+                            placeholder="Type your message...", 
+                            container=False
+                        )
+                        img = gr.Image(
+                            sources=["upload", "camera"], 
+                            type="filepath", 
+                            label="Upload or Capture Image"
+                        )
+                    
                     btn = gr.Button("Send", scale=1)
                 
                 # Clear history and memory buttons
-                clear = gr.Button("Clear Conversation")
-                clear_memory = gr.Button("Clear Memory")
+                with gr.Row():
+                    clear = gr.Button("Clear Conversation")
+                    clear_memory = gr.Button("Clear Memory")
             
-            # Submit functionality with streaming
+            # Submit functionality with streaming and image support
             btn.click(
                 fn=streaming_response, 
-                inputs=[txt, chatbot], 
-                outputs=[txt, chatbot]
+                inputs=[txt, chatbot, img], 
+                outputs=[txt, chatbot, img]
             )
             txt.submit(
                 fn=streaming_response, 
-                inputs=[txt, chatbot], 
-                outputs=[txt, chatbot]
+                inputs=[txt, chatbot, img], 
+                outputs=[txt, chatbot, img]
             )
             
             # Clear conversation history