import gradio as gr
from rembg import remove
from PIL import Image
import io

def remove_bg(input_image):
    try:
        if input_image is None:
            return None

        # Convert input image to bytes
        input_bytes = io.BytesIO()
        input_image.save(input_bytes, format='PNG')
        input_bytes = input_bytes.getvalue()

        # Run background removal
        output_bytes = remove(input_bytes)

        # Return RGBA image (with transparent background)
        output_image = Image.open(io.BytesIO(output_bytes)).convert("RGBA")
        return output_image

    except Exception as e:
        print("Error:", e)
        return None

demo = gr.Interface(fn=remove_bg, inputs=gr.Image(type="pil"), outputs="image")
demo.launch()