| """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() |
|
|
