import numpy as np
import gradio as gr
from PIL import Image
from io import BytesIO
from copy import deepcopy
from core import process_inpaint
from huggingface_hub import login
import os
import spaces
login(os.getenv("HF_TOKEN"))

@spaces.GPU()
def process_image(image, mask, progress=gr.Progress(track_tqdm=True)):
    if np.unique(mask["background"]).size == 1:
        mask = image["layers"][0]
        output = process_inpaint(image["background"], mask)
    else: 
        mask = mask["background"]
        output = process_inpaint(image["background"], mask)
    img_output = Image.fromarray(output).convert("RGB")
    
    return img_output, mask


with gr.Blocks() as demo:
    gr.Markdown("# OBJECT REMOVER")
    with gr.Row():
        with gr.Column():
            image = gr.ImageMask(type="numpy", layers=False, label="Upload Image")
            with gr.Accordion(label="Advanced", open=False):
                mask = gr.ImageMask(label="Mask", format="png", value=None, sources=["upload"])
                mask_out = gr.Image(format="png", label="Output Image")
            button = gr.Button("Remove")
        with gr.Column():
            output = gr.Image(format="png", label="Output Image")
            
        button.click(fn=process_image, inputs=[image, mask], outputs=[output, mask_out])

demo.launch(debug=True,show_error=True)