codetechdevx commited on
Commit
59fabac
·
1 Parent(s): 6b8cbc8

change to cpu

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -10,7 +10,7 @@ torch.set_float32_matmul_precision(["high", "highest"][0])
10
  birefnet = AutoModelForImageSegmentation.from_pretrained(
11
  "ZhengPeng7/BiRefNet", trust_remote_code=True
12
  )
13
- birefnet.to("cuda")
14
 
15
  transform_image = transforms.Compose(
16
  [
@@ -27,10 +27,10 @@ def fn(image):
27
  processed_image = process(im)
28
  return (processed_image, origin)
29
 
30
- @spaces.GPU
31
  def process(image):
32
  image_size = image.size
33
- input_images = transform_image(image).unsqueeze(0).to("cuda")
34
  # Prediction
35
  with torch.no_grad():
36
  preds = birefnet(input_images)[-1].sigmoid().cpu()
 
10
  birefnet = AutoModelForImageSegmentation.from_pretrained(
11
  "ZhengPeng7/BiRefNet", trust_remote_code=True
12
  )
13
+ birefnet.to("cpu")
14
 
15
  transform_image = transforms.Compose(
16
  [
 
27
  processed_image = process(im)
28
  return (processed_image, origin)
29
 
30
+ #@spaces.GPU
31
  def process(image):
32
  image_size = image.size
33
+ input_images = transform_image(image).unsqueeze(0).to("cpu")
34
  # Prediction
35
  with torch.no_grad():
36
  preds = birefnet(input_images)[-1].sigmoid().cpu()