Update app.py

app.py

@@ -5,31 +5,25 @@ import gradio as gr
 from ormbg import ORMBG
 from PIL import Image
 
-
 model_path = "ormbg.pth"
 
+# Load the model globally but don't send to device yet
 net = ORMBG()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-net.to(device)
-
-if torch.cuda.is_available():
-    net.load_state_dict(torch.load(model_path))
-    net = net.cuda()
-else:
-    net.load_state_dict(torch.load(model_path, map_location="cpu"))
+net.load_state_dict(torch.load(model_path, map_location="cpu"))
 net.eval()
 
-
 def resize_image(image):
     image = image.convert("RGB")
     model_input_size = (1024, 1024)
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
-
 def inference(image):
+    # Check for CUDA and set the device inside inference
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    net.to(device)
 
-    # prepare input
+    # Prepare input
     orig_image = Image.fromarray(image)
     w, h = orig_image.size
     image = resize_image(orig_image)
@@ -37,50 +31,41 @@ def inference(image):
     im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
     im_tensor = torch.unsqueeze(im_tensor, 0)
     im_tensor = torch.divide(im_tensor, 255.0)
+
     if torch.cuda.is_available():
-        im_tensor = im_tensor.cuda()
+        im_tensor = im_tensor.to(device)
 
-    # inference
+    # Inference
     result = net(im_tensor)
-    # post process
+    # Post process
     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)
-    # image to pil
+    # Image to PIL
     im_array = (result * 255).cpu().data.numpy().astype(np.uint8)
     pil_im = Image.fromarray(np.squeeze(im_array))
-    # paste the mask on the original image
+    # Paste the mask on the original image
     new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
     new_im.paste(orig_image, mask=pil_im)
 
     return new_im
 
-
-gr.Markdown("## Open Remove Background Model (ormbg)")
-gr.HTML(
-    """
-  <p style="margin-bottom: 10px; font-size: 94%">
-    This is a demo for Open Remove Background Model (ormbg) that using
-    <a href="https://huggingface.co/schirrmacher/ormbg" target="_blank">Open Remove Background Model (ormbg) model</a> as backbone.
-  </p>
-"""
-)
+# Gradio interface setup
 title = "Open Remove Background Model (ormbg)"
 description = r"""
 This model is a <strong>fully open-source background remover</strong> optimized for images with humans.
-
 It is based on [Highly Accurate Dichotomous Image Segmentation research](https://github.com/xuebinqin/DIS).
 The model was trained with the synthetic [Human Segmentation Dataset](https://huggingface.co/datasets/schirrmacher/humans).
 
 This is the first iteration of the model, so there will be improvements!
-If you identify cases were the model fails, <a href='https://huggingface.co/schirrmacher/ormbg/discussions' target='_blank'>upload your examples</a>!
+If you identify cases where the model fails, <a href='https://huggingface.co/schirrmacher/ormbg/discussions' target='_blank'>upload your examples</a>!
 
 - <a href='https://huggingface.co/schirrmacher/ormbg' target='_blank'>Model card</a>: find inference code, training information, tutorials
 - <a href='https://huggingface.co/schirrmacher/ormbg' target='_blank'>Dataset</a>: see training images, segmentation data, backgrounds
 - <a href='https://huggingface.co/schirrmacher/ormbg\#research' target='_blank'>Research</a>: see current approach for improvements
-
 """
+
 examples = ["./example1.png", "./example2.png", "./example3.png"]
 
 demo = gr.Interface(
@@ -89,7 +74,7 @@ demo = gr.Interface(
     outputs="image",
     examples=examples,
     title=title,
-    description=description,
+    description=description
 )
 
 if __name__ == "__main__":