mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Move bot_picture.py inside the extension

Browse source
This commit is contained in:
oobabooga 2023-02-25 03:00:19 -03:00
parent 5ac24b019e
commit 91f5852245
3 changed files with 15 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

10
modules/bot_picture.py
View file

@ -1,10 +0,0 @@
import torch
from transformers import BlipForConditionalGeneration, BlipProcessor
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=torch.float32).to("cpu")
def caption_image(raw_image):
inputs = processor(raw_image.convert('RGB'), return_tensors="pt").to("cpu", torch.float32)
out = model.generate(**inputs, max_new_tokens=100)
return processor.decode(out[0], skip_special_tokens=True)

9
modules/extensions.py
View file

@ -33,10 +33,11 @@ def apply_extensions(text, typ):
def create_extensions_block():
# Updating the default values
for extension, name in iterator():
for param in extension.params:
_id = f"{name}-{param}"
if _id in shared.settings:
extension.params[param] = shared.settings[_id]
if hasattr(extension, 'params'):
for param in extension.params:
_id = f"{name}-{param}"
if _id in shared.settings:
extension.params[param] = shared.settings[_id]
# Creating the extension ui elements
for extension, name in iterator():