Spaces:
Build error
Build error
Upload folder using huggingface_hub
Browse files
app.py
CHANGED
|
@@ -12,6 +12,13 @@ net = ORMBG()
|
|
| 12 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 13 |
net.to(device)
|
| 14 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 15 |
|
| 16 |
def resize_image(image):
|
| 17 |
image = image.convert("RGB")
|
|
@@ -59,7 +66,7 @@ gr.HTML(
|
|
| 59 |
</p>
|
| 60 |
"""
|
| 61 |
)
|
| 62 |
-
title = "Background
|
| 63 |
description = r"""
|
| 64 |
This model is a fully open-source background remover optimized for images with humans.
|
| 65 |
|
|
|
|
| 12 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 13 |
net.to(device)
|
| 14 |
|
| 15 |
+
if torch.cuda.is_available():
|
| 16 |
+
net.load_state_dict(torch.load(model_path))
|
| 17 |
+
net = net.cuda()
|
| 18 |
+
else:
|
| 19 |
+
net.load_state_dict(torch.load(model_path, map_location="cpu"))
|
| 20 |
+
net.eval()
|
| 21 |
+
|
| 22 |
|
| 23 |
def resize_image(image):
|
| 24 |
image = image.convert("RGB")
|
|
|
|
| 66 |
</p>
|
| 67 |
"""
|
| 68 |
)
|
| 69 |
+
title = "Open Remove Background Model (ormbg)"
|
| 70 |
description = r"""
|
| 71 |
This model is a fully open-source background remover optimized for images with humans.
|
| 72 |
|