import gradio as gr
import base64
import io
from PIL import Image
import spaces


@spaces.GPU
def base64_to_image(base64_str):
    print("Base to image")
    image_data = base64.b64decode(base64_str)
    image = Image.open(io.BytesIO(image_data))
    return image

@spaces.GPU
def image_to_base64(image):
    print("Image to base")
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

@spaces.GPU
def convert_to_grayscale(base64_str):
    image = base64_to_image(base64_str)
    grayscale_image = image.convert("L")
    return {
        "grayscale_image": image_to_base64(grayscale_image),
        "Normal": image_to_base64(grayscale_image)
    }


iface = gr.Interface(
    fn=convert_to_grayscale,
    inputs=gr.Textbox(label="Base64 Image"),
    outputs=gr.JSON(label="Base64 Grayscale Image"),
)


iface.launch()