import torch
from PIL import Image
from qwen_vl_utils import process_vision_info
from src.lantern_generate.generate import generate as lantern_generate
from src.models import load_model
device = "cuda" if torch.cuda.is_available() else "cpu"
model, processor = load_model("AGViveiros/LanteRn-3B-Tetris", compute_dtype=torch.bfloat16, use_cache=True)
model.eval().to(device)
processor.tokenizer.padding_side = "left"
image = Image.open("path/to/image.jpg").convert("RGB")
question = "Your question here"
messages = [{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": question},
],
}]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, _ = process_vision_info(messages)
inputs = processor(text=[text], images=image_inputs, return_tensors="pt").to(device)
prompt_len = inputs["input_ids"].shape[1]
output = model.generate(
**inputs,
max_new_tokens=512,
do_sample=False,
custom_generate=lantern_generate,
use_cache=True,
return_dict_in_generate=True,
)
generated = output.sequences[0][prompt_len:]
print(processor.decode(generated, skip_special_tokens=False))