from transformers import pipeline
from PIL import Image
import gradio as gr
import spaces

@spaces.GPU
def remove_background_bria(input_image):
    print(f"Removing background using bria for image")
    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=0)
    input_image = Image.fromarray(input_image)
    output_image = pipe(input_image)
    return output_image

with gr.Blocks() as demo:
    gr.Markdown("# Background Removal with BRIA")
    with gr.Row():
        input_image = gr.Image(label="Input Image")
        output_image = gr.Image(label="Output Image")
    
    remove_button = gr.Button("Remove Background")
    remove_button.click(fn=remove_background_bria, inputs=input_image, outputs=output_image)

if __name__ == "__main__":
    demo.launch()