File size: 1,863 Bytes
79b6316 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | #!/usr/bin/env python3
"""Run GUI grounding with a local directory or Hugging Face model ID."""
import argparse
import re
import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor
POINT_PATTERN = re.compile(r"<loc_(\d+)>,<loc_(\d+)>")
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--model", default=".")
parser.add_argument("--image", required=True)
parser.add_argument("--prompt", required=True)
return parser.parse_args()
def main():
args = parse_args()
if not torch.cuda.is_available():
raise RuntimeError("This example requires a CUDA GPU")
model = AutoModelForCausalLM.from_pretrained(
args.model,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
attn_implementation="sdpa",
).cuda().eval()
processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
image = Image.open(args.image).convert("RGB")
inputs = processor(images=image, text=args.prompt, return_tensors="pt").to(
"cuda", dtype=torch.bfloat16
)
with torch.inference_mode():
output_ids = model.generate(
**inputs,
do_sample=False,
max_new_tokens=16,
)
text = processor.tokenizer.batch_decode(
output_ids, skip_special_tokens=False
)[0]
match = POINT_PATTERN.search(text)
if match is None:
print(text)
raise RuntimeError("the model output did not contain location tokens")
normalized = tuple(map(int, match.groups()))
pixels = (
normalized[0] / 999 * image.width,
normalized[1] / 999 * image.height,
)
print(f"raw_output={text}")
print(f"normalized_point={normalized}")
print(f"pixel_point=({pixels[0]:.2f}, {pixels[1]:.2f})")
if __name__ == "__main__":
main()
|