Update app.py

app.py

@@ -11,6 +11,7 @@ import io
 import base64
 import logging
 import runpod
+import requests
 
 # Add logging configuration right after
 logging.basicConfig(level=logging.INFO)
@@ -301,32 +302,48 @@ def base64_to_pil(base64_str):
 # Simple RunPod handler that forwards to stream_chat
 def handler(event):
 	try:
-		logger.info("Received RunPod request")
-		input_data = event.get("input", {})
+		# Extract data from the Runpod event
+		data = event["input"]
 		
+		# Check if input is a dictionary (which seems to be the case from the error)
+		if isinstance(data["input_image"], dict):
+			return {
+				"status": "error",
+				"message": "Invalid image format. Expected base64 string or file data."
+			}
+			
 		# Convert base64 image to PIL Image
-		input_image = base64_to_pil(input_data.get("input_image"))
+		if isinstance(data["input_image"], str):
+			# Remove data URL prefix if present
+			if 'base64,' in data["input_image"]:
+				data["input_image"] = data["input_image"].split('base64,')[1]
+			image_data = base64.b64decode(data["input_image"])
+			image = Image.open(io.BytesIO(image_data))
+		else:
+			return {
+				"status": "error",
+				"message": "Invalid image format"
+			}
 		
-		# Forward request to stream_chat
-		prompt, caption = stream_chat(
-			input_image=input_image,
-			caption_type=input_data.get("caption_type", "Descriptive"),
-			caption_length=input_data.get("caption_length", "long"),
-			extra_options=input_data.get("extra_options", []),
-			name_input=input_data.get("name_input", ""),
-			custom_prompt=input_data.get("custom_prompt", "")
+		# Now we have a valid PIL Image, proceed with the stream_chat call
+		result = stream_chat(
+			input_image=image,
+			caption_type=data.get("caption_type", "Descriptive"),
+			caption_length=data.get("caption_length", "any"),
+			extra_options=data.get("extra_options", []),
+			name_input=data.get("name_input", ""),
+			custom_prompt=data.get("custom_prompt", "")
 		)
 		
 		return {
-			"output": {
-				"prompt": prompt,
-				"caption": caption
-			}
+			"output": result
 		}
-		
+
 	except Exception as e:
-		logger.error(f"Handler failed: {str(e)}")
-		return {"error": str(e)}
+		return {
+			"status": "error",
+			"message": str(e)
+		}
 
 if __name__ == "__main__":
 	# Start RunPod serverless