Update handler.py
Browse files- handler.py +3 -1
handler.py
CHANGED
|
@@ -5,9 +5,11 @@ import base64
|
|
| 5 |
from io import BytesIO
|
| 6 |
from transformers import CLIPProcessor, CLIPModel
|
| 7 |
|
|
|
|
|
|
|
| 8 |
class EndpointHandler():
|
| 9 |
def __init__(self, path=""):
|
| 10 |
-
self.model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14-336")
|
| 11 |
self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14-336")
|
| 12 |
|
| 13 |
def __call__(self, data: Any) -> List[float]:
|
|
|
|
| 5 |
from io import BytesIO
|
| 6 |
from transformers import CLIPProcessor, CLIPModel
|
| 7 |
|
| 8 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
| 9 |
+
|
| 10 |
class EndpointHandler():
|
| 11 |
def __init__(self, path=""):
|
| 12 |
+
self.model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14-336").to(device)
|
| 13 |
self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14-336")
|
| 14 |
|
| 15 |
def __call__(self, data: Any) -> List[float]:
|