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

def replace_bg(input_img, bg_img):
    # Convert to RGBA
    input_img = input_img.convert("RGBA")
    # Remove background
    fg = remove(input_img)

    if bg_img is None:
        return fg

    bg_img = bg_img.convert("RGBA").resize(fg.size)
    combined = Image.alpha_composite(bg_img, fg)
    return combined

demo = gr.Interface(
    fn=replace_bg,
    inputs=[gr.Image(type="pil", label="Upload Image"), 
            gr.Image(type="pil", label="Upload Background (Optional)")],
    outputs=gr.Image(type="pil", label="Result"),
    title="Background Replace Tool",
    description="Upload a photo and replace its background with your own image."
)

demo.launch()