Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # MUST come before torch / transformers | |
| import re | |
| import time | |
| import gradio as gr | |
| import torch | |
| from PIL import Image, ImageDraw, ImageFont | |
| from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration | |
| MODEL_ID = "xyzzzh/Hi-Token" | |
| GROUNDING_PROMPT = ( | |
| "Please provide the bounding box coordinate of the region this sentence " | |
| "describes: {expr}." | |
| ) | |
| # A Hi-Token box is 12 special tokens: for each of the 4 coordinates | |
| # (x_min, y_min, x_max, y_max) an axis-specific hundreds/tens/ones triplet, | |
| # e.g. <x_hundreds_4><x_tens_2><x_ones_7> -> 427 in the normalized [0, 999] | |
| # coordinate space. | |
| COORD_RE = re.compile(r"<(x|y)_(hundreds|tens|ones)_(\d)>") | |
| PLACES = ("hundreds", "tens", "ones") | |
| def parse_hitoken_box(response: str): | |
| """Reconstruct [x_min, y_min, x_max, y_max] (0-999 ints) from the model's | |
| hierarchical coordinate tokens. | |
| The tokens appear in sequence: x_min triplet, y_min triplet, x_max triplet, | |
| y_max triplet. Each integer is 100*hundreds + 10*tens + ones. | |
| Returns None if the response does not contain a well-formed box. | |
| """ | |
| tokens = COORD_RE.findall(response) | |
| if len(tokens) < 12: | |
| return None | |
| tokens = tokens[:12] | |
| values = [] | |
| for i in range(0, 12, 3): | |
| trip = tokens[i : i + 3] | |
| places = [p for _, p, _ in trip] | |
| if places != list(PLACES): | |
| return None | |
| axes = {a for a, _, _ in trip} | |
| if len(axes) != 1: | |
| return None | |
| values.append((axes.pop(), 100 * int(trip[0][2]) + 10 * int(trip[1][2]) + int(trip[2][2]))) | |
| if [axis for axis, _ in values] != ["x", "y", "x", "y"]: | |
| return None | |
| return [v for _, v in values] | |
| processor = AutoProcessor.from_pretrained(MODEL_ID, max_pixels=768 * 28 * 28) | |
| model = ( | |
| Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID, torch_dtype=torch.bfloat16, attn_implementation="sdpa" | |
| ) | |
| .to("cuda") | |
| .eval() | |
| ) | |
| def ground( | |
| image, | |
| expression: str = "the main subject", | |
| show_raw_tokens: bool = False, | |
| ): | |
| """Locate a region in an image from a text description and draw its box. | |
| Args: | |
| image: input image to search. | |
| expression: referring expression describing the target region, | |
| e.g. "the smallest rubber duck". | |
| show_raw_tokens: also return the model's raw generated token string. | |
| Returns: | |
| Tuple of (annotated image, decoded coordinate string). | |
| """ | |
| t0 = time.time() | |
| image = image.convert("RGB") | |
| width, height = image.size | |
| query = GROUNDING_PROMPT.format(expr=expression.strip()) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image"}, | |
| {"type": "text", "text": query}, | |
| ], | |
| } | |
| ] | |
| prompt = processor.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| inputs = processor( | |
| text=[prompt], images=[image], padding=True, return_tensors="pt" | |
| ).to(model.device) | |
| with torch.inference_mode(): | |
| generated = model.generate( | |
| **inputs, | |
| max_new_tokens=64, | |
| do_sample=False, | |
| use_cache=True, | |
| ) | |
| response_ids = generated[:, inputs.input_ids.shape[1] :] | |
| # Coordinate tokens ARE special tokens, so they must survive decoding. | |
| response = processor.batch_decode( | |
| response_ids, | |
| skip_special_tokens=False, | |
| clean_up_tokenization_spaces=False, | |
| )[0] | |
| box = parse_hitoken_box(response) | |
| elapsed = time.time() - t0 | |
| if box is None: | |
| annotated = image.copy() | |
| summary = "No bounding box found in the model output." | |
| if show_raw_tokens: | |
| summary += f"\n\nRaw output:\n{response.strip()}" | |
| return annotated, summary | |
| x_min, y_min, x_max, y_max = box | |
| note = "" | |
| if x_max < x_min or y_max < y_min: | |
| # Normalize an inverted prediction so it can still be drawn. | |
| x_min, x_max = min(x_min, x_max), max(x_min, x_max) | |
| y_min, y_max = min(y_min, y_max), max(y_min, y_max) | |
| note = " (corners normalized)" | |
| px = [ | |
| x_min * width / 999.0, | |
| y_min * height / 999.0, | |
| x_max * width / 999.0, | |
| y_max * height / 999.0, | |
| ] | |
| annotated = image.copy() | |
| draw = ImageDraw.Draw(annotated) | |
| line = max(2, int(0.004 * max(width, height)) + 1) | |
| # White halo under an orange core so the box reads on any background. | |
| draw.rectangle(px, outline=(255, 255, 255), width=line * 3) | |
| draw.rectangle(px, outline=(255, 64, 0), width=line) | |
| label = expression.strip() | |
| try: | |
| font = ImageFont.load_default(size=max(14, int(0.02 * width) + 4)) | |
| except TypeError: | |
| font = ImageFont.load_default() | |
| tb = draw.textbbox((0, 0), label, font=font) | |
| tw, th = tb[2] - tb[0], tb[3] - tb[1] | |
| lx = max(0, min(px[0], width - tw - 8)) | |
| ly = px[1] - th - 12 | |
| if ly < 0: | |
| ly = px[1] + 6 | |
| draw.rectangle([lx - 5, ly - 5, lx + tw + 5, ly + th + 5], fill=(255, 64, 0)) | |
| draw.text((lx, ly), label, fill=(255, 255, 255), font=font) | |
| summary = ( | |
| f"Bounding box (0-999){note}: [{x_min}, {y_min}, {x_max}, {y_max}] " | |
| f"pixels: [{px[0]:.0f}, {px[1]:.0f}, {px[2]:.0f}, {px[3]:.0f}] " | |
| f"inference: {elapsed:.1f}s" | |
| ) | |
| if show_raw_tokens: | |
| summary += f"\n\nRaw output:\n{response.strip()}" | |
| return annotated, summary | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """ | |
| # Hi-Token: Hierarchical Coordinate Tokenization for Visual Grounding | |
| Point this generative grounding model at an image and describe a | |
| region in plain language — it emits each bounding-box coordinate as | |
| an axis-specific hundreds/tens/ones token triplet (12 tokens per | |
| box) instead of plain numbers, which sharpens boundary precision. | |
| Model: [xyzzzh/Hi-Token](https://huggingface.co/xyzzzh/Hi-Token) | |
| (Qwen2.5-VL-3B fine-tune, Hi-R1) · | |
| [Paper](https://huggingface.co/papers/2608.03471) | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_in = gr.Image( | |
| type="pil", label="Image", sources=["upload", "clipboard"] | |
| ) | |
| expression_in = gr.Textbox( | |
| label="Referring expression", | |
| placeholder="e.g. the smallest rubber duck", | |
| lines=1, | |
| ) | |
| show_raw = gr.Checkbox( | |
| label="Show raw coordinate tokens", value=False | |
| ) | |
| run_btn = gr.Button("Locate", variant="primary") | |
| with gr.Column(): | |
| image_out = gr.Image(type="pil", label="Grounded region") | |
| text_out = gr.Textbox(label="Decoded coordinates", lines=4) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/rubber_ducks.jpg", "the smallest rubber duck"], | |
| ["examples/rubber_ducks.jpg", "the largest rubber duck"], | |
| ["examples/vintage_camera.jpg", "the black instant camera"], | |
| ["examples/pizza_board.jpg", "the slice being lifted"], | |
| ["examples/motorcycle_vintage.jpg", "the green motorcycle"], | |
| ], | |
| inputs=[image_in, expression_in], | |
| outputs=[image_out, text_out], | |
| fn=ground, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| run_btn.click( | |
| ground, | |
| inputs=[image_in, expression_in, show_raw], | |
| outputs=[image_out, text_out], | |
| api_name="ground", | |
| ) | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |