stojnvla commited on
Commit
11fd0ad
·
1 Parent(s): 61e7520

add decorator

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -7,6 +7,7 @@ import torchvision.transforms as T
7
  import torch.nn.functional as F
8
  from lposs import lposs, lposs_plus
9
  import torch
 
10
 
11
  device = "cpu"
12
  if torch.cuda.is_available():
@@ -21,6 +22,7 @@ maskclip = MaskClip().to(device)
21
  dino = DINO().to(device)
22
  to_torch_tensor = T.Compose([T.Resize(size=448, max_size=2048), T.ToTensor()])
23
 
 
24
  def segment_image(img: PIL.Image.Image, classnames: str, use_lposs_plus: bool | None) -> tuple[np.ndarray | PIL.Image.Image | str, list[tuple[np.ndarray | tuple[int, int, int, int], str]]]:
25
  img_tensor = to_torch_tensor(PIL.Image.fromarray(img)).unsqueeze(0).to(device)
26
  classnames = [c.strip() for c in classnames.split(",")]
 
7
  import torch.nn.functional as F
8
  from lposs import lposs, lposs_plus
9
  import torch
10
+ import spaces
11
 
12
  device = "cpu"
13
  if torch.cuda.is_available():
 
22
  dino = DINO().to(device)
23
  to_torch_tensor = T.Compose([T.Resize(size=448, max_size=2048), T.ToTensor()])
24
 
25
+ @spaces.GPU
26
  def segment_image(img: PIL.Image.Image, classnames: str, use_lposs_plus: bool | None) -> tuple[np.ndarray | PIL.Image.Image | str, list[tuple[np.ndarray | tuple[int, int, int, int], str]]]:
27
  img_tensor = to_torch_tensor(PIL.Image.fromarray(img)).unsqueeze(0).to(device)
28
  classnames = [c.strip() for c in classnames.split(",")]