Try enabling GPU if available
Browse files- handler.py +4 -2
handler.py
CHANGED
@@ -5,10 +5,12 @@ from PIL import Image
|
|
5 |
import torch
|
6 |
from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection
|
7 |
|
|
|
|
|
8 |
class EndpointHandler():
|
9 |
def __init__(self, path=""):
|
10 |
-
self.text_model = CLIPTextModel.from_pretrained("rbanfield/clip-vit-large-patch14")
|
11 |
-
self.image_model = CLIPVisionModelWithProjection.from_pretrained("rbanfield/clip-vit-large-patch14")
|
12 |
self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")
|
13 |
|
14 |
def __call__(self, data):
|
|
|
5 |
import torch
|
6 |
from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection
|
7 |
|
8 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
9 |
+
|
10 |
class EndpointHandler():
|
11 |
def __init__(self, path=""):
|
12 |
+
self.text_model = CLIPTextModel.from_pretrained("rbanfield/clip-vit-large-patch14").to(device)
|
13 |
+
self.image_model = CLIPVisionModelWithProjection.from_pretrained("rbanfield/clip-vit-large-patch14").to(device)
|
14 |
self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")
|
15 |
|
16 |
def __call__(self, data):
|