|
|
|
from transformers import AutoProcessor, AutoModelForVision2Seq |
|
import torch |
|
from PIL import Image |
|
class VisionLLM: |
|
def __init__(self, device="cuda", model_id="microsoft/kosmos-2"): |
|
self.device = device |
|
self.processor = AutoProcessor.from_pretrained(model_id) |
|
self.model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16).to(self.device) |
|
|
|
|
|
def describe_images(self, images, prompt="", max_length=128): |
|
if isinstance(images, list): |
|
images = [img.convert("RGB") for img in images] |
|
else: |
|
images = images.convert("RGB") |
|
inputs = self.processor(images=images, text=prompt, return_tensors="pt").to(self.device) |
|
with torch.no_grad(), torch.autocast("cuda"): |
|
outputs = self.model.generate(**inputs, max_length=max_length) |
|
descriptions = self.processor.batch_decode(outputs, skip_special_tokens=True) |
|
return descriptions |
|
|
|
if __name__ == "__main__": |
|
vllm = VisionLLM() |
|
images = [Image.open("multi_view_output.png")] |
|
prompt = "Describe the objects in the image" |
|
descriptions = vllm.describe_images(images, prompt) |
|
print(descriptions) |
