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()