import gradio as gr
from PIL import Image
import rembg

# Function to handle image uploads, optional background removal, and collage generation
def process_images(images, remove_bg):
    final_collage = Image.new("RGBA", (2480, 3508), (255, 255, 255, 0))  # A4 size transparent canvas
    
    # Positioning each image onto the A4 canvas
    x_offset = 0
    y_offset = 0
    for img in images:
        img = Image.open(img)
        
        # Optionally remove background
        if remove_bg:
            img = Image.open(rembg.remove(img))
        
        # Resize the image if necessary
        img.thumbnail((500, 500))  # Limit the size of each sticker for now
        
        # Paste the image onto the canvas
        final_collage.paste(img, (x_offset, y_offset), img)
        
        # Update position for next image (for simplicity stacking vertically)
        y_offset += img.height
        if y_offset + img.height > final_collage.height:
            y_offset = 0
            x_offset += img.width
            if x_offset + img.width > final_collage.width:
                break  # Stop adding images if there's no more space

    return final_collage

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Sticker Collage Creator (A4)")
    
    with gr.Row():
        images = gr.Files(label="Upload your images", file_types=["image"])  # Use gr.Files for multiple image uploads
        remove_bg = gr.Checkbox(label="Remove background?", value=True)
    
    output = gr.Image(label="Collage", type="pil")
    
    submit = gr.Button("Create Collage")
    submit.click(process_images, inputs=[images, remove_bg], outputs=output)

demo.launch()