Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -6,12 +6,14 @@ from transformers import AutoModelForImageSegmentation
|
|
6 |
import torch
|
7 |
from torchvision import transforms
|
8 |
|
9 |
-
|
|
|
|
|
10 |
|
11 |
birefnet = AutoModelForImageSegmentation.from_pretrained(
|
12 |
"ZhengPeng7/BiRefNet", trust_remote_code=True
|
13 |
)
|
14 |
-
birefnet.to("
|
15 |
|
16 |
transform_image = transforms.Compose(
|
17 |
[
|
@@ -28,10 +30,12 @@ def fn(image):
|
|
28 |
processed_image = process(im)
|
29 |
return (processed_image, origin)
|
30 |
|
31 |
-
@spaces.GPU
|
|
|
|
|
32 |
def process(image):
|
33 |
image_size = image.size
|
34 |
-
input_images = transform_image(image).unsqueeze(0).to("
|
35 |
# Prediction
|
36 |
with torch.no_grad():
|
37 |
preds = birefnet(input_images)[-1].sigmoid().cpu()
|
|
|
6 |
import torch
|
7 |
from torchvision import transforms
|
8 |
|
9 |
+
# GPU ์ค์ ์ CPU๋ก ๋ณ๊ฒฝ
|
10 |
+
# GPU ์ค์ ์ ์ญ์ ํ๊ฑฐ๋ "cuda"๋ฅผ "cpu"๋ก ๋ณ๊ฒฝ
|
11 |
+
# torch.set_float32_matmul_precision("high")๋ CPU์์ ํ์ ์์.
|
12 |
|
13 |
birefnet = AutoModelForImageSegmentation.from_pretrained(
|
14 |
"ZhengPeng7/BiRefNet", trust_remote_code=True
|
15 |
)
|
16 |
+
birefnet.to("cpu") # GPU -> CPU๋ก ๋ณ๊ฒฝ
|
17 |
|
18 |
transform_image = transforms.Compose(
|
19 |
[
|
|
|
30 |
processed_image = process(im)
|
31 |
return (processed_image, origin)
|
32 |
|
33 |
+
# @spaces.GPU ๋ฐ์ฝ๋ ์ดํฐ ์ ๊ฑฐ
|
34 |
+
# CPU ํ๊ฒฝ์์ ๋์ํ๋๋ก ์ค์
|
35 |
+
|
36 |
def process(image):
|
37 |
image_size = image.size
|
38 |
+
input_images = transform_image(image).unsqueeze(0).to("cpu") # GPU -> CPU๋ก ๋ณ๊ฒฝ
|
39 |
# Prediction
|
40 |
with torch.no_grad():
|
41 |
preds = birefnet(input_images)[-1].sigmoid().cpu()
|