import runpod
import requests
from PIL import Image
import io
import base64

GRADIO_SERVER_URL = "http://127.0.0.1:7860"  # Default Gradio port

def handler(event):
	try:
		# Extract data from the Runpod event
		data = event["input"]
		
		# Convert base64 image to PIL Image
		image_data = base64.b64decode(data["image"])
		image = Image.open(io.BytesIO(image_data))
		
		# Prepare the payload for Gradio endpoint
		payload = {
			"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", "")
		}

		# Forward to Gradio's stream_chat endpoint
		response = requests.post(
			f"{GRADIO_SERVER_URL}/stream_chat",
			json=payload
		)

		if response.status_code == 200:
			return {
				"status": "success",
				"data": response.json()
			}
		else:
			return {
				"status": "error",
				"message": f"Gradio request failed with status {response.status_code}"
			}

	except Exception as e:
		return {
			"status": "error",
			"message": str(e)
		}

runpod.serverless.start({"handler": handler}, test_input=None)