Kims12 commited on
Commit
4e96c9f
ยท
verified ยท
1 Parent(s): a505bd1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -4
app.py CHANGED
@@ -6,12 +6,14 @@ from transformers import AutoModelForImageSegmentation
6
  import torch
7
  from torchvision import transforms
8
 
9
- torch.set_float32_matmul_precision("high")
 
 
10
 
11
  birefnet = AutoModelForImageSegmentation.from_pretrained(
12
  "ZhengPeng7/BiRefNet", trust_remote_code=True
13
  )
14
- birefnet.to("cuda")
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("cuda")
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()