samson-s
/

motiv-receipt-donut

vision-encoder-decoder

Inference Endpoints

Model card Files Files and versions Community

samson-s commited on Mar 13

Commit

da85024

•

1 Parent(s): 44cd819

Add gpu specification

Files changed (1) hide show

handler.py +2 -1

handler.py CHANGED Viewed

@@ -1,6 +1,6 @@
 from transformers import pipeline, AutoFeatureExtractor
 from PIL import Image
-import io
 class EndpointHandler:
@@ -9,6 +9,7 @@ class EndpointHandler:
             "image-to-text",
             model=path,
             feature_extractor=AutoFeatureExtractor,
         )
     def __call__(self, data) -> str:

 from transformers import pipeline, AutoFeatureExtractor
 from PIL import Image
+import torch
 class EndpointHandler:
             "image-to-text",
             model=path,
             feature_extractor=AutoFeatureExtractor,
+            device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
         )
     def __call__(self, data) -> str: