Jonny001 commited on
Commit
ac68fcd
·
verified ·
1 Parent(s): 95fb114

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
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("cuda")
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("cuda")
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()