File size: 2,806 Bytes
e706c52 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
from typing import Dict, Any
from transformers import QwenImageProcessor, QwenTokenizer, QwenForMultiModalConditionalGeneration
import torch
from PIL import Image
import io
import json
import base64
import requests
class EndpointHandler():
def __init__(self, path=""):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = QwenForMultiModalConditionalGeneration.from_pretrained(
path,
torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32
).to(self.device)
self.image_processor = QwenImageProcessor.from_pretrained(path)
self.tokenizer = QwenTokenizer.from_pretrained(path)
self.model.generation_config.use_cache = False
def __call__(self, data: Any) -> Dict[str, Any]:
"""
Args:
data (Any): The input data, which can be:
- Binary image data in the request body.
- A dictionary with 'image' and 'text' keys:
- 'image': Base64-encoded image string or image URL.
- 'text': The text prompt.
Returns:
Dict[str, Any]: The generated text output from the model.
"""
if isinstance(data, (bytes, bytearray)):
image = Image.open(io.BytesIO(data)).convert('RGB')
text_input = "<|im_start|>user\nDescribe this image.\n<|im_end|><|im_start|>assistant\n"
elif isinstance(data, dict):
image_input = data.get('image', None)
text_input = data.get('text', '')
if image_input is None:
return {"error": "No image provided."}
if image_input.startswith('http'):
response = requests.get(image_input)
image = Image.open(io.BytesIO(response.content)).convert('RGB')
else:
image_data = base64.b64decode(image_input)
image = Image.open(io.BytesIO(image_data)).convert('RGB')
else:
return {"error": "Invalid input data. Expected binary image data or a dictionary with 'image' key."}
image_inputs = self.image_processor(images=image, return_tensors="pt").to(self.device)
if not text_input:
text_input = "<|im_start|>user\nDescribe this image.\n<|im_end|><|im_start|>assistant\n"
input_ids = self.tokenizer(text_input, return_tensors="pt").input_ids.to(self.device)
generated_ids = self.model.generate(
**image_inputs,
input_ids=input_ids,
max_new_tokens=256,
do_sample=True,
top_p=0.9,
temperature=0.7,
)
output_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
return {"generated_text": output_text}
|