Update app.py

app.py

@@ -1,22 +1,53 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Depends
 from typing import Optional
 from fastapi.responses import StreamingResponse
 from huggingface_hub import InferenceClient
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 import os
-import uvicorn
-
+from base64 import b64encode
+from io import BytesIO
+from PIL import Image  # Add this import
+import logging
 
 app = FastAPI()
 
+# Configure logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+# Get HuggingFace token from environment variable
 
 # Default model
 DEFAULT_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"
 
 class QueryRequest(BaseModel):
+    model_config = ConfigDict(protected_namespaces=())
+    
     query: str
+    image_data: Optional[str] = None  # Base64 encoded image data
     stream: bool = False
-    model_name: Optional[str] = None  # If not provided, will use DEFAULT_MODEL
+    model_name: Optional[str] = None
+
+class ChatForm(BaseModel):
+    model_config = ConfigDict(protected_namespaces=())
+    
+    query: str
+    stream: bool = False
+    model_name: Optional[str] = None
+    
+    @classmethod
+    def as_form(
+        cls,
+        query: str = Form(...),
+        stream: bool = Form(False),
+        model_name: Optional[str] = Form(None),
+        image: Optional[UploadFile] = File(None)
+    ):
+        return cls(
+            query=query,
+            stream=stream,
+            model_name=model_name
+        ), image
 
 def get_client(model_name: Optional[str] = None):
     """Get inference client for specified model or default model"""
@@ -25,7 +56,7 @@ def get_client(model_name: Optional[str] = None):
         model_path = model_name if model_name and model_name.strip() else DEFAULT_MODEL
         
         return InferenceClient(
-            model_path
+            model=model_path
         )
     except Exception as e:
         raise HTTPException(
@@ -33,12 +64,26 @@ def get_client(model_name: Optional[str] = None):
             detail=f"Error initializing model {model_path}: {str(e)}"
         )
 
-def generate_response(query: str, model_name: Optional[str] = None):
+def generate_response(query: str, image_data: Optional[str] = None, model_name: Optional[str] = None):
     messages = []
-    messages.append({
-        "role": "user", 
-        "content": f"[SYSTEM] You are ASSISTANT who answer question asked by user in short and concise manner. [USER] {query}"
-    })
+    
+    # Create the system and user message
+    user_content = f"[SYSTEM] You are ASSISTANT who answer question asked by user in short and concise manner. [USER] {query}"
+    
+    # If there's an image, add it to the message
+    if image_data:
+        messages.append({
+            "role": "user",
+            "content": [
+                {"type": "text", "text": user_content},
+                {"type": "image_url", "image_url": {"url": f"data:image/*;base64,{image_data}"}}
+            ]
+        })
+    else:
+        messages.append({
+            "role": "user",
+            "content": user_content
+        })
 
     try:
         client = get_client(model_name)
@@ -57,17 +102,56 @@ async def root():
     return {"message": "Welcome to FastAPI server!"}
 
 @app.post("/chat")
-async def chat(request: QueryRequest):
+async def chat(form_data: tuple[ChatForm, Optional[UploadFile]] = Depends(ChatForm.as_form)):
+    form, image = form_data
     try:
-        if request.stream:
+        image_data = None
+        if image:
+            logger.debug("Image received")
+            # Read the image
+            contents = await image.read()
+            
+            # Convert image to appropriate format if needed
+            try:
+                logger.debug("Attempting to open image")
+                img = Image.open(BytesIO(contents))
+                logger.debug(f"Image format before conversion: {img.format}, mode: {img.mode}")
+                # Convert to RGB if needed
+                if img.mode != 'RGB':
+                    img = img.convert('RGB')
+                logger.debug(f"Image format after conversion: {img.format}, mode: {img.mode}")
+                
+                # Save as JPEG in memory
+                buffer = BytesIO()
+                img.save(buffer, format="JPEG")
+                image_data = b64encode(buffer.getvalue()).decode('utf-8')
+                logger.debug("Image processed and encoded to base64")
+            except Exception as img_error:
+                logger.error(f"Error processing image: {str(img_error)}")
+                raise HTTPException(
+                    status_code=422,
+                    detail=f"Error processing image: {str(img_error)}"
+                )
+        
+        if form.stream:
             return StreamingResponse(
-                generate_response(request.query, request.model_name),
+                generate_response(form.query, image_data, form.model_name),
                 media_type="text/event-stream"
             )
         else:
             response = ""
-            for chunk in generate_response(request.query, request.model_name):
+            for chunk in generate_response(form.query, image_data, form.model_name):
                 response += chunk
             return {"response": response}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Error in /chat endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(
+        "main:app", 
+        port=8000,
+        reload=True,  # Enable auto-reload
+        reload_dirs=["./"]  # Watch the current directory for changes
+    )