import gradio as gr
from gradio_client import Client, handle_file
import os
from PIL import Image
import io
import base64

hf_token = os.environ.get("HF_TKN")

def convert_base64_to_img(image_string):
    # Split the input string to separate the metadata header and the base64-encoded data
    header, encoded_data = image_string.split(",", 1)
    
    # Now, encoded_data contains the base64-encoded image data
    image_data = base64.b64decode(encoded_data)
    
    # Create a BytesIO object to store the image data
    image_file = io.BytesIO(image_data)
    
    # Open the image using the BytesIO object
    img = Image.open(image_file)
    
    # Save the image as a JPEG file
    img.save('output.png', 'PNG')
    
    return "output.png"

def infer(image_string, question):
    image_in = convert_base64_to_img(image_string)
    client = Client("fffiloni/moondream1", hf_token=hf_token)
    result = client.predict(
        image = handle_file(image_in),
        question = question,
    	api_name="/predict"
    )
    print(result)
    return result

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_string = gr.Textbox(interactive=True)
            question = gr.Textbox(interactive=True)
            submit_btn = gr.Button("Submit", interactive=True)
        with gr.Column():
            answer = gr.Textbox(interactive=False)
            
    submit_btn.click(
        fn=infer,
        inputs=[image_string, question],
        outputs=[answer]
    )

demo.launch()