Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-22 08:07:56 +01:00
Code Issues Packages Projects Releases Wiki Activity

Move BLIP to the CPU

Browse Source 
It's just as fast
This commit is contained in:
oobabooga 2023-02-15 00:03:19 -03:00
parent d4d90a8000
commit 1622059179
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
modules/bot_picture.py
View File

@ -5,10 +5,9 @@ from transformers import BlipForConditionalGeneration
from transformers import BlipProcessor from transformers import BlipProcessor
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=torch.float16).to("cuda") model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=torch.float32).to("cpu")
# raw_image = Image.open('/tmp/istockphoto-470604022-612x612.jpg').convert('RGB')
def caption_image(raw_image): def caption_image(raw_image):
inputs = processor(raw_image, return_tensors="pt").to("cuda", torch.float16) inputs = processor(raw_image.convert('RGB'), return_tensors="pt").to("cpu", torch.float32)
out = model.generate(**inputs, max_new_tokens=100) out = model.generate(**inputs, max_new_tokens=100)
return processor.decode(out[0], skip_special_tokens=True) return processor.decode(out[0], skip_special_tokens=True)