Update app.py

app.py

@@ -32,4 +32,42 @@ def process(image, progress=gr.Progress()):
     im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
 
     if torch.cuda.is_available():
-        im_tensor
+        im_tensor = im_tensor.cuda()
+
+    # Inference
+    result = net(im_tensor)
+    progress(0.7)  # 70% progress during inference
+    
+    # Post-process
+    result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
+    result = (result - result.min()) / (result.max() - result.min())
+
+    # Convert to PIL image with alpha mask
+    result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
+    pil_mask = Image.fromarray(result_array)
+    new_im = orig_image.copy()
+    new_im.putalpha(pil_mask)
+
+    # Save the image for download
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+    new_im.save(temp_file.name, format='PNG')
+
+    progress(1.0)  # 100% complete
+    return new_im, temp_file.name
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## BRIA RMBG 1.4 - Background Remover")
+    gr.HTML("<p>Upload your image to remove the background.</p>")
+    
+    with gr.Row():
+        input_image = gr.Image(type="numpy", label="Upload Image")
+        output_image = gr.Image(type="pil", label="Processed Image")
+    
+    download_button = gr.File(label="Download Processed Image")
+    
+    input_image.change(fn=process, inputs=input_image, outputs=[output_image, download_button])
+
+# Run the app on port 7860
+demo.launch(server_name="0.0.0.0", server_port=7860)
+