Update quick start with single-image inference example

README.md

@@ -55,30 +55,79 @@ The model repository includes `configuration_uhr_bat.py` and `modeling_uhr_bat.p
 ## Quick Start
 
 ```python
-from transformers import AutoModelForCausalLM
+import importlib
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "FelixKAI/UHR-BAT"
+image_path = "your_remote_sensing_image.jpg"
+question = "Describe this remote-sensing image."
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
+image_processor = AutoImageProcessor.from_pretrained(model_id)
 
 model = AutoModelForCausalLM.from_pretrained(
-    "FelixKAI/UHR-BAT",
+    model_id,
     trust_remote_code=True,
     torch_dtype="auto",
     device_map="auto",
-)
-```
-
-Tokenizer:
-
-```python
-from transformers import AutoTokenizer
-
-tokenizer = AutoTokenizer.from_pretrained("FelixKAI/UHR-BAT")
-```
+).eval()
 
-Image processor:
+# Reuse the preprocessing helpers shipped with the model's remote code.
+uhrbat = importlib.import_module(model.__class__.__module__)
+image = Image.open(image_path).convert("RGB")
 
-```python
-from transformers import AutoImageProcessor
-
-image_processor = AutoImageProcessor.from_pretrained("FelixKAI/UHR-BAT")
+prompt = (
+    "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+    f"<|im_start|>user\n<image>\n{question}<|im_end|>\n"
+    "<|im_start|>assistant\n"
+)
+image_token_id = getattr(model.config, "image_token_index", -200)
+input_ids = uhrbat.tokenizer_image_token(
+    prompt,
+    tokenizer,
+    image_token_id,
+    return_tensors="pt",
+).unsqueeze(0).to(model.device)
+attention_mask = torch.ones_like(input_ids)
+
+target_sizes = [672, 1344, 2688, 4032]
+multiscale_pixels = [
+    uhrbat.split_image_to_multiscale_tiles(
+        image,
+        image_processor,
+        target_sizes=target_sizes,
+        tile_size=336,
+    )
+]
+
+with torch.inference_mode():
+    output = model.generate(
+        inputs=input_ids,
+        attention_mask=attention_mask,
+        image_sizes=[image.size],
+        modalities=["image"],
+        multiscale_pixels=multiscale_pixels,
+        multiscale_masks=[{}],
+        multiscale_topk=[80, 320, 600, 2000],
+        multiscale_target_sizes=target_sizes,
+        do_sample=False,
+        max_new_tokens=256,
+        return_dict_in_generate=True,
+        output_scores=True,
+        pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+prompt_len = output.sequences.shape[1] - len(output.scores)
+answer_ids = output.sequences[:, prompt_len:].clone()
+answer_ids[answer_ids < 0] = tokenizer.pad_token_id or tokenizer.eos_token_id
+answer = tokenizer.decode(answer_ids[0], skip_special_tokens=True).strip()
+print(answer)
 ```
 
 ## Citation