import gradio as gr
from transformers import pipeline
import torch
import numpy as np
from PIL import Image
import io

def remove_background(input_image):
    try:
        # Initialize the pipeline with trust_remote_code=True
        segmentor = pipeline("image-segmentation", 
                           model="briaai/RMBG-1.4",
                           device=-1,
                           trust_remote_code=True)
        
        # Process the image
        result = segmentor(input_image)
        return result['output_image']
    except Exception as e:
        raise gr.Error(f"Error processing image: {str(e)}")

# Custom CSS for mobile-friendly design
css = """
.gradio-container {
    font-family: 'Segoe UI', sans-serif;
    max-width: 100% !important;
    padding: 10px !important;
}
.container {
    display: flex;
    flex-direction: column;
    gap: 20px;
}
.image-container {
    width: 100%;
    max-width: 500px;
    margin: 0 auto;
}
.gr-button {
    background: linear-gradient(45deg, #FFD700, #FFA500);
    border: none !important;
    color: black !important;
    padding: 12px 20px !important;
    border-radius: 8px !important;
    font-weight: bold !important;
    margin: 10px 0 !important;
    width: 100% !important;
    max-width: 300px !important;
}
@media (max-width: 768px) {
    .gradio-container {
        padding: 5px !important;
    }
    .gr-button {
        padding: 10px 15px !important;
    }
}
"""

# Create Gradio interface
with gr.Blocks(css=css) as demo:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
            <h1 style="font-size: 2rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
                AI Background Remover
            </h1>
            <p style="color: #666; font-size: 1rem; margin-bottom: 2rem;">
                Powered by RMBG V1.4 model
            </p>
        </div>
        """
    )
    
    with gr.Row(equal_height=True):
        with gr.Column():
            # Removed 'tool' parameter and added mobile-friendly settings
            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                sources=["upload", "clipboard"],
                interactive=True
            )
            
        with gr.Column():
            output_image = gr.Image(
                label="Result",
                type="pil",
                interactive=False
            )
    
    with gr.Row():
        clear_btn = gr.Button("Clear", variant="secondary")
        process_btn = gr.Button("Remove Background", variant="primary")

    # Status message for feedback
    status_msg = gr.Textbox(label="Status", interactive=False, visible=False)

    # Event handlers
    def process_and_update(image):
        if image is None:
            return None, "Please upload an image first"
        try:
            result = remove_background(image)
            return result, "Background removed successfully!"
        except Exception as e:
            return None, f"Error: {str(e)}"

    process_btn.click(
        fn=process_and_update,
        inputs=[input_image],
        outputs=[output_image, status_msg],
        api_name="remove_background"
    )
    
    clear_btn.click(
        fn=lambda: (None, None, ""),
        outputs=[input_image, output_image, status_msg],
        api_name="clear"
    )

# Launch the app
demo.launch(
    share=True,
    enable_queue=True,
    show_error=True,
    server_port=7860,
    server_name="0.0.0.0"
)