rbanfield commited on
Commit
e62633b
·
1 Parent(s): 3c0b02a

Try enabling GPU if available

Browse files
Files changed (1) hide show
  1. 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):