import gradio as gr
from models import make_inpainting
import io
from PIL import Image
import numpy as np

from PIL import Image
from typing import Union
import random
import numpy as np
import os
import time

from models import make_image_controlnet, make_inpainting
from segmentation import segment_image
from config import HEIGHT, WIDTH, POS_PROMPT, NEG_PROMPT, COLOR_MAPPING, map_colors, map_colors_rgb
from palette import COLOR_MAPPING_CATEGORY
from preprocessing import preprocess_seg_mask, get_image, get_mask
from explanation import make_inpainting_explanation, make_regeneration_explanation, make_segmentation_explanation


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):

    # image = Image.open(requests.get("https://applydesignblobs-chh5aahjdzh0cnew.z01.azurefd.net/spaceimages/org_sqr_7fee0869-3187-4363-b5fb-5233e943649d.png", stream=True).raw)
    # mask = Image.open(requests.get("https://applydesign.blob.core.windows.net/spaceimages/mask_e85b1585-8.png", stream=True).raw)

    result_image = make_inpainting(positive_prompt='test1',
                                   image=image_to_byte_array(input_img1),
                                   mask_image=np.array(input_img2),
                                   negative_prompt="xxx",
                                   )


    # predictions = pipeline(input_img1)
    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)