app.py

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

@spaces.GPU
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")
    
    # Convert to PNG to preserve transparency
    if output.mode != 'RGBA':
        output = output.convert('RGBA')
    
    return output

iface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(label="Upload Image"),
    outputs=gr.Image(label="Output Image", type="pil", format="png"),
    title="✂️ Image Background Remover ✂️",
    description="⚠️ Sorry for the inconvenience. The model is currently running on the CPU, which might affect performance. We appreciate your understanding.",
    theme="Yntec/HaleyCH_Theme_Orange"
)

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