davanstrien HF Staff commited on
Commit
8eeaa6b
·
1 Parent(s): fdde8cd

use better transformers

Browse files
Files changed (2) hide show
  1. app.py +5 -1
  2. requirements.txt +2 -1
app.py CHANGED
@@ -1,9 +1,12 @@
1
  from pathlib import Path
2
  import gradio as gr
3
- from transformers import pipeline
 
4
  from transformers import AutoModelForImageClassification
5
  import shutil
 
6
 
 
7
 
8
  chk_point = "davanstrien/autotrain-encyclopaedia-illustrations-blog-post-3327992158"
9
 
@@ -12,6 +15,7 @@ model = AutoModelForImageClassification.from_pretrained(chk_point)
12
  pipe = pipeline(
13
  "image-classification",
14
  "davanstrien/autotrain-encyclopaedia-illustrations-blog-post-3327992158",
 
15
  )
16
 
17
 
 
1
  from pathlib import Path
2
  import gradio as gr
3
+ import torch
4
+ # from transformers import pipeline
5
  from transformers import AutoModelForImageClassification
6
  import shutil
7
+ from optimum.pipelines import pipeline
8
 
9
+ device = 1 if torch.cuda.is_available() else "cpu"
10
 
11
  chk_point = "davanstrien/autotrain-encyclopaedia-illustrations-blog-post-3327992158"
12
 
 
15
  pipe = pipeline(
16
  "image-classification",
17
  "davanstrien/autotrain-encyclopaedia-illustrations-blog-post-3327992158",
18
+ accelerator="bettertransformer",device=device
19
  )
20
 
21
 
requirements.txt CHANGED
@@ -1 +1,2 @@
1
- transformers[torch]
 
 
1
+ transformers[torch]
2
+ optimum