from handler import EndpointHandler
from PIL import Image
import requests
import base64
from io import BytesIO
import time

# init handler
my_handler = EndpointHandler(path=".")



# prompt = "<grounding>An image of"
prompt = "<grounding>Describe this image in detail"

url = "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.png"
image = Image.open(requests.get(url, stream=True).raw)

# The original Kosmos-2 demo saves the image first then reload it. For some images, this will give slightly different image input and change the generation outputs.
image.save("new_image.jpg")
image = Image.open("img1.jpg")
buffered = BytesIO()
image.save(buffered, format='JPEG')
image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

# prepare sample payload
non_holiday_payload = {"prompt": prompt, "image_base64": image_base64}
# holiday_payload = {"inputs": "Today is a though day", "date": "2022-07-04"}
init_t = time.time()
# test the handler
non_holiday_pred=my_handler(non_holiday_payload)
# holiday_payload=my_handler(holiday_payload)

# show results
print("image_description", non_holiday_pred)
print(time.time() - init_t)
# print("holiday_payload", holiday_payload)

# non_holiday_pred [{'label': 'joy', 'score': 0.9985942244529724}]
# holiday_payload [{'label': 'happy', 'score': 1}]