##### Set API variables #####
import os

hf_api_key = os.environ['HF_API_KEY']
hf_api_itt_base = os.environ['HF_API_ITT_BASE']

##### Helper functions #####
import io, base64, requests, json

def image_to_base64_str(pil_image):
    byte_arr = io.BytesIO()
    pil_image.save(byte_arr, format='PNG')
    byte_arr = byte_arr.getvalue()
    return str(base64.b64encode(byte_arr).decode('utf-8'))

def get_completion(inputs, parameters = None, ENDPOINT_URL = hf_api_itt_base): 
    headers = {
      "Authorization": f"Bearer {hf_api_key}",
      "Content-Type": "application/json"
    }    
    data = { "inputs" : inputs }
    if parameters is not None:
        data.update({"parameters" : parameters})
    response = requests.post(ENDPOINT_URL,
                             headers = headers,
                             data = json.dumps(data))
    return json.loads(response.content.decode("utf-8"))

def captioner(pil_image):
    base64_image = image_to_base64_str(pil_image)
    result = get_completion(base64_image)
    return result[0]['generated_text']
    
##### Gradio App #####
import gradio as gr 

gr.close_all()

demo = gr.Interface(fn = captioner,
                    inputs = [gr.Image(label = "Upload image", type = "pil")],
                    outputs = [gr.Textbox(label = "Caption")],
                    title = "Image captioning with BLIP",
                    description = "Caption any image using the BLIP model",
                    allow_flagging = "never",
                    examples = ["tank-aerial-view.jpg", "f35-carrier-takeoff.jpg", "the_los_angelesclass_submarine.jpg"])

demo.launch()