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

# 1. Load the AI Model into memory ONCE at startup
# 'isnet-general-use' is highly accurate for products, people, and objects
session = new_session("isnet-general-use")

# 2. Define the core function
def process_image(input_img):
    if input_img is None:
        return None
    try:
        # Remove background using the pre-loaded session
        output = remove(input_img, session=session)
        return output
    except Exception as e:
        print(f"Error processing image: {e}")
        return None

# 3. Build the User Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # ✨ AI Background Remover
        Optimized for high-speed API access via CPU.
        """
    )
    
    with gr.Row():
        img_input = gr.Image(type="pil", label="Upload Image")
        img_output = gr.Image(type="pil", label="Transparent Result")
    
    btn = gr.Button("Remove Background", variant="primary")
    
    # 4. Connect the button and declare the API name for the blog
    btn.click(
        fn=process_image, 
        inputs=img_input, 
        outputs=img_output, 
        api_name="remove_bg"
    )

# 5. Launch the app with strict queuing to protect the CPU
if __name__ == "__main__":
    # Explicitly define the Hugging Face Docker network to prevent localhost errors
    demo.queue(default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=7860)