app.py

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:
        print("\n Did not Receive Mask\n")
        output = process_inpaint(image["background"], image["layers"][0])
    else: 
        print("Processing Received Mask")
        print("Mask Background shape", mask["background"].shape)
        print("Mask : ", mask)
        print("Mask : ", mask["background"])
        output = process_inpaint(image["background"], mask["background"])
    img_output = Image.fromarray(output).convert("RGB")
    
    return img_output


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"])
            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])

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