import gradio as gr
import io
from PIL import Image
import numpy as np
from models import make_image_controlnet, make_inpainting
from preprocessing import preprocess_seg_mask, get_image, get_mask

def image_to_byte_array(image: Image) -> bytes:
    # BytesIO is a fake file stored in memory
    imgByteArr = io.BytesIO()
    # image.save expects a file as a argument, passing a bytes io ins
    image.save(imgByteArr, format='png')  # image.format
    # Turn the BytesIO object back into a bytes object
    imgByteArr = imgByteArr.getvalue()
    return imgByteArr

def predict(input_img1,input_img2):

    print("predict")

    input_img1 = input_img1.resize((512, 512))
    input_img2 = input_img2.resize((512, 512))

    canvas_mask = np.array(input_img2)
    mask = get_mask(canvas_mask)

    print(input_img1, mask)

    result_image = make_inpainting(positive_prompt='test1',
                                   image=input_img1,
                                   mask_image=mask,
                                   negative_prompt="xxx",
                                   )

    return result_image

gradio_app = gr.Interface(
    predict,
    inputs=[gr.Image(label="img", sources=['upload', 'webcam'], type="pil"),
            gr.Image(label="mask", sources=['upload', 'webcam'], type="pil")
            ],
    outputs= gr.Image(label="resp"),
    title="rem fur 1",
)

gradio_app.launch(share=True)