p1atdev commited on
Commit
ebf9292
1 Parent(s): f561c25

chore: zero gpu test

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -5,7 +5,7 @@ import gc
5
  from PIL import Image
6
  from transformers import AutoModel, AutoImageProcessor
7
  from anime2sketch.model import Anime2Sketch
8
- # import spaces
9
 
10
  setup()
11
 
@@ -17,7 +17,7 @@ class MangaLineExtractor:
17
  model = AutoModel.from_pretrained(MLE_MODEL_REPO, trust_remote_code=True)
18
  processor = AutoImageProcessor.from_pretrained(MLE_MODEL_REPO, trust_remote_code=True)
19
 
20
- # @spaces.GPU
21
  @torch.no_grad()
22
  def __call__(self, image: Image.Image) -> Image.Image:
23
  inputs = self.processor(image, return_tensors="pt")
 
5
  from PIL import Image
6
  from transformers import AutoModel, AutoImageProcessor
7
  from anime2sketch.model import Anime2Sketch
8
+ import spaces
9
 
10
  setup()
11
 
 
17
  model = AutoModel.from_pretrained(MLE_MODEL_REPO, trust_remote_code=True)
18
  processor = AutoImageProcessor.from_pretrained(MLE_MODEL_REPO, trust_remote_code=True)
19
 
20
+ @spaces.GPU
21
  @torch.no_grad()
22
  def __call__(self, image: Image.Image) -> Image.Image:
23
  inputs = self.processor(image, return_tensors="pt")