Update app.py
Browse files
app.py
CHANGED
|
@@ -12,7 +12,7 @@ torch.set_float32_matmul_precision(["high", "highest"][0])
|
|
| 12 |
birefnet = AutoModelForImageSegmentation.from_pretrained(
|
| 13 |
"ZhengPeng7/BiRefNet", trust_remote_code=True
|
| 14 |
)
|
| 15 |
-
birefnet.to("
|
| 16 |
|
| 17 |
transform_image = transforms.Compose(
|
| 18 |
[
|
|
@@ -32,7 +32,7 @@ def fn(image: Union[Image.Image, str]) -> Tuple[Image.Image, Image.Image]:
|
|
| 32 |
@spaces.GPU
|
| 33 |
def process(image: Image.Image) -> Image.Image:
|
| 34 |
image_size = image.size
|
| 35 |
-
input_images = transform_image(image).unsqueeze(0).to("
|
| 36 |
# Prediction
|
| 37 |
with torch.no_grad():
|
| 38 |
preds = birefnet(input_images)[-1].sigmoid().cpu()
|
|
|
|
| 12 |
birefnet = AutoModelForImageSegmentation.from_pretrained(
|
| 13 |
"ZhengPeng7/BiRefNet", trust_remote_code=True
|
| 14 |
)
|
| 15 |
+
birefnet.to("cpu")
|
| 16 |
|
| 17 |
transform_image = transforms.Compose(
|
| 18 |
[
|
|
|
|
| 32 |
@spaces.GPU
|
| 33 |
def process(image: Image.Image) -> Image.Image:
|
| 34 |
image_size = image.size
|
| 35 |
+
input_images = transform_image(image).unsqueeze(0).to("cpu")
|
| 36 |
# Prediction
|
| 37 |
with torch.no_grad():
|
| 38 |
preds = birefnet(input_images)[-1].sigmoid().cpu()
|