| import re | |
| import torch | |
| from torchvision import transforms | |
| from transformers import BlipForConditionalGeneration, BlipProcessor | |
| from internals.util.commons import download_image | |
| from internals.util.config import get_hf_cache_dir | |
| class Image2Text: | |
| __loaded = False | |
| def load(self): | |
| if self.__loaded: | |
| return | |
| self.processor = BlipProcessor.from_pretrained( | |
| "Salesforce/blip-image-captioning-large", | |
| cache_dir=get_hf_cache_dir(), | |
| ) | |
| self.model = BlipForConditionalGeneration.from_pretrained( | |
| "Salesforce/blip-image-captioning-large", | |
| torch_dtype=torch.float16, | |
| cache_dir=get_hf_cache_dir(), | |
| ).to("cuda") | |
| self.__loaded = True | |
| def process(self, imageUrl: str) -> str: | |
| self.load() | |
| image = download_image(imageUrl).resize((512, 512)) | |
| inputs = self.processor.__call__(image, return_tensors="pt").to( | |
| "cuda", torch.float16 | |
| ) | |
| output_ids = self.model.generate( | |
| **inputs, do_sample=False, top_p=0.9, max_length=128 | |
| ) | |
| output_text = self.processor.batch_decode(output_ids) | |
| print(output_text) | |
| output_text = output_text[0] | |
| output_text = re.sub("</.>|\\n|\[SEP\]", "", output_text) | |
| return output_text | |