import os
import gradio as gr
from rembg import remove

# Set environment variable for Numba
os.environ['NUMBA_CACHE_DIR'] = '/tmp/numba_cache'

def remove_background(input_image):
    try:
        output = remove(input_image)
        return output
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        return None

interface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="filepath", label="Input Image"),
    outputs=gr.Image(type="filepath", label="Output Image"),
    title="Background Remover",
    description="Upload an image to remove the background",
    allow_flagging="never"
)

interface.launch(server_name="0.0.0.0", server_port=7860)