import spaces
import gradio as gr
from transparent_background import Remover
from PIL import Image
import numpy as np

@spaces.GPU(duration=20)
def remove_background(image):
    remover = Remover()
    if isinstance(image, Image.Image):
        output = remover.process(image)
    elif isinstance(image, np.ndarray):
        image_pil = Image.fromarray(image)
        output = remover.process(image_pil)
    else:
        raise TypeError("Unsupported image type")
    return output

iface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(label="Upload Image"),
    outputs=gr.Image(label="Output Image")
)

if __name__ == "__main__":
    iface.launch()