Update app.py

app.py

@@ -9,12 +9,11 @@ from PIL import Image
 import tempfile
 import os
 
-
 # Load the pre-trained model
 net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 net.to(device)
-net.eval()    
+net.eval()
 
 # Resize the input image for model compatibility
 def resize_image(image):
@@ -34,48 +33,42 @@ def process(image):
     im_tensor = torch.unsqueeze(im_tensor, 0)
     im_tensor = torch.divide(im_tensor, 255.0)
     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.cuda()
-
+    
     # Inference with the model
-    result = net(im_tensor)
+    with torch.no_grad():
+        result = net(im_tensor)
     
     # Post-process the result
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
     ma = torch.max(result)
     mi = torch.min(result)
-    result = (result - mi) / (ma - mi)  # Normalize result
-
+    result = (result - mi) / (ma - mi)
+    
     # Convert the result to an image
     result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
     pil_mask = Image.fromarray(np.squeeze(result_array))
-
+    
     # Add the mask as alpha channel to the original image
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
-
+    
     # Save the processed image to a temporary file
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
-    new_im.save(temp_file, format='PNG')
-    temp_file.close()  # Ensure the file is closed before Gradio uses it
-
-    return temp_file.name  # Return the path to the temporary file for downloading
+    new_im.save(temp_file.name, format='PNG')
+    temp_file.close()
+    
+    # Return the path to the temporary file for downloading
+    return temp_file.name
 
 # Gradio interface setup
-gr.Markdown("## BRIA RMBG 1.4")
-gr.HTML('''<p style="margin-bottom: 10px; font-size: 94%">
-  This is a demo for BRIA RMBG 1.4 that uses 
-  <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as a backbone.
-</p>''')
-
 title = "Background Removal"
-description = r"""Background removal model developed by <a href='https://BRIA.AI' target='_blank'><b>BRIA.AI</b></a>, trained on a carefully selected dataset and is available as an open-source model for non-commercial use.<br> 
-For testing, upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>. To purchase a commercial license, simply click <a href='https://go.bria.ai/3ZCBTLH' target='_blank'><b>Here</b></a>. <br>"""
+description = r"""Background removal model developed by <a href='https://BRIA.AI' target='_blank'><b>BRIA.AI</b></a>, trained on a carefully selected dataset and is available as an open-source model for non-commercial use.<br> For testing, upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>. To purchase a commercial license, simply click <a href='https://go.bria.ai/3ZCBTLH' target='_blank'><b>Here</b></a>. <br>"""
 
-examples = [['./input.jpg'],]
+examples = [['./input.jpg']]
 
-# Modify the interface to use live updates and file download
+# Create the Gradio interface
 demo = gr.Interface(
     fn=process,  # The function to process the image
     inputs=gr.Image(type="numpy"),  # Input type (image)
@@ -87,4 +80,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch(share=False)