qwenvl/app/qwenvl.py

106 lines
3.0 KiB
Python
Raw Normal View History

2024-11-12 16:39:19 +08:00
import torch
from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
from appPublic.worker import awaitify
from appPublic.jsonConfig import getConfig
from ahserver.serverenv import ServerEnv
from ahserver.webapp import webapp
class Qwen2VLClass:
def __init__(self, modelname):
# default: Load the model on the available device(s)
self.model = Qwen2VLForConditionalGeneration.from_pretrained(
modelname,
torch_dtype=torch.bfloat16,
# attn_implementation="flash_attention_2",
device_map="auto"
)
self.min_pixels = 256 * 28 * 28
self.max_pixels = 1280 * 28 * 28
# We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
# model = Qwen2VLForConditionalGeneration.from_pretrained(
# "Qwen/Qwen2-VL-7B-Instruct",
# torch_dtype=torch.bfloat16,
# attn_implementation="flash_attention_2",
# device_map="auto",
# )
# default processer
self.processor = AutoProcessor.from_pretrained(modelname,
min_pixels=self.min_pixels,
max_pixels=self.max_pixels
)
# The default range for the number of visual tokens per image in the model is 4-16384.
# You can set min_pixels and max_pixels according to your needs, such as a token range of 256-1280, to balance performance and cost.
# min_pixels = 256*28*28
# max_pixels = 1280*28*28
# processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
2024-11-12 17:57:40 +08:00
def inference(self, prompt, image, videofile):
2024-11-12 16:39:19 +08:00
content = [
{
"type":"text",
"text":prompt
}
]
if image:
if not image.startswith('file:///') \
and not image.startswith('http://') \
and not image.startswith('https://'):
image = f'data:image;base64,{image}'
content.append({
"type":"image",
"image":image
})
if videofile:
if not videofile.startswith('file:///'):
return 'only local video file support'
content.append({
"type":"video",
"video":videofile
})
messages = [
{
"role": "user",
"content": content
}
]
# Preparation for inference
text = self.processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
image_inputs, video_inputs = process_vision_info(messages)
inputs = self.processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
)
inputs = inputs.to("cuda")
# Inference: Generation of the output
generated_ids = self.model.generate(**inputs, max_new_tokens=128)
generated_ids_trimmed = [
out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
2024-11-12 17:57:40 +08:00
output_text = self.processor.batch_decode(
2024-11-12 16:39:19 +08:00
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
2024-11-12 17:57:40 +08:00
return output_text[0]
2024-11-12 16:39:19 +08:00
def main():
config = getConfig()
modelname = config.modelname
m = Qwen2VLClass(modelname)
g = ServerEnv()
g.inference = awaitify(m.inference)
if __name__ == '__main__':
webapp(main)