File size: 1,732 Bytes
94c9578 | 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 | """Single-image inference with UHR-BAT2's native Qwen3-VL architecture."""
import argparse
from pathlib import Path
import torch
from PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", default="RL-MIND/UHR-BAT2")
parser.add_argument("--image", required=True)
parser.add_argument("--question", default="Describe this remote-sensing image.")
parser.add_argument("--max-new-tokens", type=int, default=128)
args = parser.parse_args()
processor = AutoProcessor.from_pretrained(args.model)
model = Qwen3VLForConditionalGeneration.from_pretrained(
args.model,
dtype="auto",
device_map="auto",
attn_implementation="sdpa",
).eval()
image = Image.open(Path(args.image)).convert("RGB")
messages = [{"role": "user", "content": [
{"type": "image"}, {"type": "text", "text": args.question},
]}]
prompt = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True,
)
inputs = processor(text=[prompt], images=[image], return_tensors="pt")
inputs.pop("token_type_ids", None)
inputs = inputs.to(model.device)
with torch.inference_mode():
generated = model.generate(
**inputs,
max_new_tokens=args.max_new_tokens,
do_sample=False,
temperature=None,
top_p=None,
top_k=None,
use_cache=True,
)
answer = generated[:, inputs["input_ids"].shape[1]:]
print(processor.batch_decode(answer, skip_special_tokens=True)[0].strip())
if __name__ == "__main__":
main()
|