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",
            trust_remote_code=True,
            device="cpu"
        )
        
        # 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 theme with valid colors
theme = gr.themes.Soft(
    primary_hue="orange",    # Changed from "gold" to "orange"
    secondary_hue="blue",
    neutral_hue="gray"
).set(
    body_background_fill="linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%)",
    body_text_color="#ffffff",
    button_primary_background_fill="linear-gradient(45deg, #FFD700, #FFA500)",
    button_primary_text_color="#000000",
    border_color_primary="#FFD700"
)

# Create Gradio interface
with gr.Blocks(theme=theme) as demo:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
            <h1 style="font-size: 2.5rem; 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: #cccccc; font-size: 1.2rem; margin-bottom: 2rem;">
                Remove backgrounds instantly using RMBG V1.4 model
            </p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                sources=["upload", "clipboard"]
            )
            
        with gr.Column():
            output_image = gr.Image(
                label="Result",
                type="pil"
            )
    
    with gr.Row():
        clear_btn = gr.Button("Clear", variant="secondary")
        process_btn = gr.Button("Remove Background", variant="primary")

    # Status message
    status_msg = gr.Textbox(
        label="Status",
        placeholder="Ready to process your image...",
        interactive=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],
    )
    
    clear_btn.click(
        fn=lambda: (None, None, "Ready to process your image..."),
        outputs=[input_image, output_image, status_msg],
    )

# Launch the app
demo.launch()