Browse Source
fix the gpu usage.
Signed-off-by: wxywb <xy.wang@zilliz.com>
main
wxywb
2 years ago
1 changed files with
4 additions and
1 deletions
-
taiyi.py
|
|
|
@ -41,6 +41,9 @@ class Taiyi(NNOperator): |
|
|
|
self.clip_model = CLIPModel.from_pretrained(config['clip_model']) |
|
|
|
self.processor = CLIPProcessor.from_pretrained(config['processor']) |
|
|
|
|
|
|
|
self.text_encoder.to(self.device) |
|
|
|
self.clip_model.to(self.device) |
|
|
|
|
|
|
|
def inference_single_data(self, data): |
|
|
|
if self.modality == 'image': |
|
|
|
vec = self._inference_from_image(data) |
|
|
|
@ -72,7 +75,7 @@ class Taiyi(NNOperator): |
|
|
|
@arg(1, to_image_color('RGB')) |
|
|
|
def _inference_from_image(self, img): |
|
|
|
image = to_pil(img) |
|
|
|
image = self.processor(images=image, return_tensors="pt") |
|
|
|
image = self.processor(images=image, return_tensors="pt").to(self.device) |
|
|
|
image_features = self.clip_model.get_image_features(**image) |
|
|
|
return image_features |
|
|
|
|
|
|
|
|
