Update app.py

app.py

@@ -1,9 +1,13 @@
 #!/usr/bin/env python3
 import os
 import json
+import base64
 import requests
 import gradio as gr
+from PIL import Image
+from io import BytesIO
 
+# Get environment variables from HF Spaces secrets
 ENDPOINT = os.environ.get("VLLM_ENDPOINT")
 MODEL = os.environ.get("VLLM_MODEL")
 
@@ -11,6 +15,13 @@ if not ENDPOINT or not MODEL:
     raise ValueError("VLLM_ENDPOINT and VLLM_MODEL environment variables must be set. Please add them as secrets in your Space settings.")
 
 
+def image_to_base64(image):
+    """Convert PIL Image to base64 string."""
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
 def respond(
     message,
     history: list[dict[str, str]],
@@ -20,11 +31,40 @@ def respond(
     top_p,
 ):
     """
-    Send messages to vLLM endpoint and stream the response.
+    Send messages (with optional images) to vLLM endpoint and stream the response.
     """
     messages = [{"role": "system", "content": system_message}]
-    messages.extend(history)
-    messages.append({"role": "user", "content": message})
+    
+    # Add conversation history
+    for msg in history:
+        messages.append(msg)
+    
+    # Process the current message - check if it contains an image
+    if message and "files" in message and message["files"]:
+        # Message has image(s)
+        content = []
+        
+        # Add text if present
+        if message.get("text", "").strip():
+            content.append({"type": "text", "text": message["text"]})
+        
+        # Add all images
+        for file_info in message["files"]:
+            try:
+                image = Image.open(file_info)
+                b64_image = image_to_base64(image)
+                content.append({
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/png;base64,{b64_image}"}
+                })
+            except Exception as e:
+                print(f"Error processing image: {e}")
+        
+        messages.append({"role": "user", "content": content})
+    else:
+        # Text-only message
+        text_content = message if isinstance(message, str) else message.get("text", "")
+        messages.append({"role": "user", "content": text_content})
 
     payload = {
         "model": MODEL,
@@ -70,34 +110,48 @@ def respond(
         yield f"Error: {str(e)}"
 
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-with gr.Blocks(title="vLLM Chatbot") as demo:
-    gr.Markdown("# 💬 Chat Interface")
+# Build the Gradio Interface
+with gr.Blocks(title="💬 Vision Chat", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 💬 Vision-Enabled Chat Interface
+        **💡 How to use:**
+        1. Type your message in the chat box
+        2. Optionally upload images by clicking the 📎 icon
+        3. Adjust parameters in the accordion below if needed
+        4. Press Enter or click Send
+        
+        The model can understand both text and images!
+        """
+    )
+    
+    chatbot = gr.ChatInterface(
+        respond,
+        type="messages",
+        multimodal=True,
+        additional_inputs=[
+            gr.Textbox(
+                value="You are a helpful AI assistant with vision capabilities. You can understand and analyze images.",
+                label="System message"
+            ),
+            gr.Slider(minimum=1, maximum=4096, value=2048, step=1, label="Max new tokens"),
+            gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
+            gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.95,
+                step=0.05,
+                label="Top-p (nucleus sampling)",
+            ),
+        ],
+    )
+    
+    chatbot.render()
+    
     gr.Markdown("""
-    Configure the endpoint via environment variables:
-    - `VLLM_ENDPOINT`: vLLM server URL
-    - `VLLM_MODEL`: Model name
+    ---
+    **Note:** Configure endpoint via `VLLM_ENDPOINT` and `VLLM_MODEL` environment variables.
     """)
-    chatbot.render()
 
 
 if __name__ == "__main__":