Update app.py

app.py

@@ -9,22 +9,19 @@ import matplotlib as mpl
 import numpy as np
 import uuid
 import imageio.v3 as iio
+import base64
+import io
+import re
 
 import torch
+import torchvision
 from torchvision.transforms.functional import to_pil_image
 from huggingface_hub import hf_hub_download
 
 import spaces
 import gradio as gr
 
-# GRADIO_TMP = os.path.join(os.path.dirname(__file__), ".gradio_tmp")
-# Path(GRADIO_TMP).mkdir(parents=True, exist_ok=True)
-
-# os.environ["GRADIO_TEMP_DIR"] = GRADIO_TMP
-# os.environ["TMPDIR"] = GRADIO_TMP
-# os.environ["TEMP"] = GRADIO_TMP
-# os.environ["TMP"] = GRADIO_TMP
-
+from transformers import SamModel, SamProcessor
 from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
 from sam2 import VQ_SAM2, VQ_SAM2Config, SAM2Config
 from visualizer import sample_color, draw_mask
@@ -78,15 +75,14 @@ MODEL = 'zhouyik/Qwen3-VL-8B-SAMTok'
 TITLE = 'SAMTok: Representing Any Mask with Two Words'
 
 HEADER = """
-<p align="center" style="margin: 1em 0 2em;"><img width="260" src="https://github.com/bytedance/Sa2VA/blob/main/projects/samtok/figs/logo.png"></p>
-<h3 align="center">SAMTok: Representing Any Mask with Two Words</h3>
+<h2 align="center">SAMTok: Representing Any Mask with Two Words</h3>
 <div style="display: flex; justify-content: center; gap: 5px;">
     <a href="https://github.com/bytedance/Sa2VA/tree/main/projects/samtok" target="_blank"><img src="https://img.shields.io/badge/arXiv-2509.18094-red"></a>
     <a href="https://github.com/bytedance/Sa2VA/tree/main/projects/samtok" target="_blank"><img src="https://img.shields.io/badge/Project-Page-brightgreen"></a>
     <a href="https://huggingface.co/collections/zhouyik/samtok" target="_blank"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-blue"></a>
     <a href="https://github.com/bytedance/Sa2VA" target="_blank"><img src="https://img.shields.io/github/stars/bytedance/Sa2VA"></a>
 </div>
-<p style="margin-top: 1em;">SAMTok provides a unified mask-token interface for MLLMs. (1) SAMTok compresses region masks into two discrete tokens and faithfully reconstructs them across diverse visual domains. (2) Injecting these mask tokens into MLLMs enables a wide range of region-level mask generation and understanding tasks. (3) The text-based representation of region masks allows a purely textual answer-matching reward for the GRPO of the mask generation task.</p>
+<p style="margin-top: 1em;">SAMTok provides a unified mask-token interface for MLLMs.</p>
 """
 
 JS = """
@@ -99,33 +95,63 @@ function init() {
 window.addEventListener('load', init);
 """
 
-device = torch.device('cuda')
-
-model = Qwen3VLForConditionalGeneration.from_pretrained(
-    MODEL, torch_dtype="auto"
-).cuda().eval()
-
-processor = AutoProcessor.from_pretrained(MODEL)
+MT_START_TOKEN = '<|mt_start|>'
+MT_END_TOKEN = '<|mt_end|>'
+MT_CONTEXT_TOKEN = '<|mt_{}|>'
 
 # build vq-sam2 model
+vq_sam2 = None
+sam2_image_processor = DirectResize(1024)
 sam2_ckpt_local = hf_hub_download(repo_id=MODEL, filename="sam2.1_hiera_large.pt")
 mask_tokenizer_local = hf_hub_download(repo_id=MODEL, filename="mask_tokenizer_256x2.pth")
 CODEBOOK_SIZE = 256
 CODEBOOK_DEPTH = 2
-sam2_config = SAM2Config(
-    ckpt_path=sam2_ckpt_local,
-)
-vq_sam2_config = VQ_SAM2Config(
-    sam2_config=sam2_config,
-    codebook_size=CODEBOOK_SIZE,
-    codebook_depth=CODEBOOK_DEPTH,
-    shared_codebook=False,
-    latent_dim=256,
-)
-vq_sam2 = VQ_SAM2(vq_sam2_config).cuda().eval()
-state = torch.load(mask_tokenizer_local, map_location="cpu")
-vq_sam2.load_state_dict(state)
-sam2_image_processor = DirectResize(1024)
+def load_vq_sam2():
+    global vq_sam2
+
+    if vq_sam2 is not None:
+        return vq_sam2
+    
+    if hasattr(torch, "set_default_device"):
+        torch.set_default_device("cpu")
+    
+    sam2_config = SAM2Config(
+        ckpt_path=sam2_ckpt_local,
+    )
+    vq_sam2_config = VQ_SAM2Config(
+        sam2_config=sam2_config,
+        codebook_size=CODEBOOK_SIZE,
+        codebook_depth=CODEBOOK_DEPTH,
+        shared_codebook=False,
+        latent_dim=256,
+    )
+
+    vq_sam2 = VQ_SAM2(vq_sam2_config)
+    state = torch.load(mask_tokenizer_local, map_location="cpu")
+    vq_sam2.load_state_dict(state)
+
+    vq_sam2 = vq_sam2.cuda().eval()
+    return vq_sam2
+
+processor = AutoProcessor.from_pretrained(MODEL)
+sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+
+_qwen = None
+_sam = None
+
+def get_qwen():
+    """Must be called only inside @spaces.GPU function."""
+    global _qwen
+    if _qwen is None:
+        _qwen = Qwen3VLForConditionalGeneration.from_pretrained(MODEL, torch_dtype="auto").to("cuda").eval()
+    return _qwen
+
+def get_sam():
+    """Must be called only inside @spaces.GPU function."""
+    global _sam
+    if _sam is None:
+        _sam = SamModel.from_pretrained("facebook/sam-vit-huge").to("cuda").eval()
+    return _sam
 
 
 colors = sample_color()
@@ -150,10 +176,312 @@ def reset_seg():
 def reset_reg():
     return 1, gr.update(interactive=False)
 
+def new_mu_state():
+    return {
+        "image_path": None,
+        "ori_size": None,                 # (w, h)
+        "original_sizes": None,           # e.g. [h, w]
+        "reshaped_input_sizes": None,     # e.g. [h', w']
+        "image_embeddings": None,         # numpy array on CPU
+        "points": [],
+        "labels": [],
+        "cur_mask": None,                 # np.uint8 (H,W)
+        "regions": {},
+        "next_region_id": 1,
+    }
+
+@spaces.GPU
+def mu_on_upload_image(media_path, mu_state):
+    if not media_path:
+        return new_mu_state(), None, None
+
+    sam_model = get_sam()  # GPU-side
+
+    img = Image.open(media_path).convert("RGB")
+    w, h = img.size
+
+    inputs = sam_processor(img, return_tensors="pt").to("cuda")
+    with torch.no_grad():
+        emb = sam_model.get_image_embeddings(inputs["pixel_values"])  # CUDA tensor
+
+    st = new_mu_state()
+    st["image_path"] = media_path
+    st["ori_size"] = (w, h)
+
+    # store sizes as python lists (not tensors)
+    st["original_sizes"] = inputs["original_sizes"][0].detach().cpu().tolist()
+    st["reshaped_input_sizes"] = inputs["reshaped_input_sizes"][0].detach().cpu().tolist()
+
+    # store embeddings as CPU numpy (picklable)
+    st["image_embeddings"] = emb[0].detach().cpu().to(torch.float16).numpy()  # (256,64,64)
+
+    return st, media_path, None
+
+def mu_predict_mask_from_state(mu_state):
+    if mu_state["image_embeddings"] is None or mu_state["image_path"] is None:
+        return None
+    if len(mu_state["points"]) == 0:
+        return None
+
+    sam_model = get_sam()
+
+    img = Image.open(mu_state["image_path"]).convert("RGB")
+
+    prompt_inputs = sam_processor(
+        img,
+        input_points=[mu_state["points"]],
+        input_labels=[mu_state["labels"]],
+        return_tensors="pt",
+    ).to("cuda")
+
+    # restore embedding to CUDA tensor, shape (1,256,64,64)
+    emb = torch.from_numpy(mu_state["image_embeddings"]).to("cuda")
+    emb = emb.unsqueeze(0)
+
+    with torch.no_grad():
+        outputs = sam_model(
+            image_embeddings=emb,
+            input_points=prompt_inputs["input_points"],
+            input_labels=prompt_inputs["input_labels"],
+            multimask_output=False,
+        )
+
+    # postprocess needs lists/tensors on CPU
+    original_sizes = torch.tensor([mu_state["original_sizes"]], dtype=torch.long)
+    reshaped_sizes = torch.tensor([mu_state["reshaped_input_sizes"]], dtype=torch.long)
+
+    masks = sam_processor.post_process_masks(
+        outputs.pred_masks.detach().cpu(),
+        original_sizes,
+        reshaped_sizes,
+    )
+    mask = masks[0][0].numpy()
+    mask = (mask > 0).astype(np.uint8)
+    return mask
+
+@spaces.GPU
+def mu_add_point(evt: gr.SelectData, mu_state, is_positive: bool):
+    if mu_state["image_path"] is None:
+        return mu_state, None
+
+    x, y = evt.index
+    mu_state["points"].append([float(x), float(y)])
+    mu_state["labels"].append(1 if is_positive else 0)
+
+    mask = mu_predict_mask_from_state(mu_state)
+    mu_state["cur_mask"] = mask
+    return mu_state, mask
+
+def mu_clear_prompts(mu_state):
+    mu_state["points"] = []
+    mu_state["labels"] = []
+    mu_state["cur_mask"] = None
+    return mu_state, None
+
+
+@spaces.GPU
+def mu_save_region(mu_state):
+    if mu_state["cur_mask"] is None:
+        return mu_state, gr.update(choices=[], value=None)
+
+    rid = f"region{mu_state['next_region_id']}"
+    mu_state["next_region_id"] += 1
+
+    reg = {"mask": mu_state["cur_mask"], "token_str": None, "zoom_in_token_str": None, "zoom_in_image": None}
+
+    vq_sam2 = load_vq_sam2()
+
+    image = Image.open(mu_state["image_path"]).convert('RGB')
+    ori_width, ori_height = image.size
+
+    sam2_image = np.array(image)
+    sam2_image = sam2_image_processor.apply_image(sam2_image)
+    sam2_pixel_values = torch.from_numpy(sam2_image).permute(2, 0, 1).contiguous()
+    sam2_pixel_values = sam2_pixel_values.unsqueeze(0).to(vq_sam2.dtype).to(vq_sam2.device)
+
+    masks = torch.stack([torch.from_numpy(np.ascontiguousarray(mu_state["cur_mask"].copy()))])
+
+    boxes = torchvision.ops.masks_to_boxes(masks)
+    x1, y1, x2, y2 = boxes.squeeze().cpu().numpy().tolist()
+    boxes_w = x2 - x1
+    boxes_h = y2 - y1
+    boxes_area = boxes_h * boxes_w
+    image_area = ori_height * ori_width
+    boxes_occupied_ratio = boxes_area / image_area
+
+    whwh = torch.as_tensor([[ori_width, ori_height, ori_width, ori_height]])
+    boxes = boxes / whwh
+    boxes = boxes.to(vq_sam2.device)
+    masks = [m.unsqueeze(0).to(vq_sam2.device) for m in masks]
+    
+    with torch.no_grad():
+        vq_sam2_output = vq_sam2(
+            sam2_pixel_values,
+            masks,
+            boxes,
+            reconstruct_mask=False,
+        )
+
+    quant_codes = vq_sam2_output.quant_codes.squeeze().cpu().numpy().astype(np.int32).tolist()
+    remap_quant_codes = [depth_idx*CODEBOOK_SIZE+quant_code for depth_idx, quant_code in enumerate(quant_codes)]
+    quant_codes = remap_quant_codes
+    global_mask_tokens_str = MT_START_TOKEN + ''.join([MT_CONTEXT_TOKEN.format(str(code).zfill(4)) for code in quant_codes]) + MT_END_TOKEN
+
+    reg["token_str"] = global_mask_tokens_str
+
+    if boxes_occupied_ratio < 0.3:
+        bbox_w = x2 - x1
+        bbox_h = y2 - y1
+        if bbox_w < 140:
+            x1 = x1 - (140 - bbox_w) // 2
+            x2 = x2 + (140 - bbox_w) // 2
+        if bbox_h < 140:
+            y1 = y1 - (140 - bbox_h) // 2
+            y2 = y2 + (140 - bbox_h) // 2
+        x1 = int(max(0, x1))
+        x2 = int(min(ori_width, x2))
+        y1 = int(max(0, y1))
+        y2 = int(min(ori_height, y2))
+
+        cropped_image = image.crop((x1, y1, x2, y2))
+        crop_width, crop_height = cropped_image.size
+        
+        if crop_width > crop_height and crop_width < 280:
+            ratio = 280 / crop_height
+            new_height = 280
+            new_width = int(crop_width * ratio)
+            resized_crop_image = cropped_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        elif crop_height > crop_width and crop_height < 280:
+            ratio = 280 / crop_width
+            new_width = 280
+            new_height = int(crop_height * ratio)
+            resized_crop_image = cropped_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        elif crop_height == crop_width and crop_width < 280:
+            ratio = 280 / crop_height
+            new_height = 280
+            new_width = int(crop_width * ratio)
+            resized_crop_image = cropped_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        else:
+            new_height = new_width = None
+            resized_crop_image = None
+
+        if resized_crop_image is None:
+            cropped_sam2_image = np.array(cropped_image)
+            cropped_sam2_image = sam2_image_processor.apply_image(cropped_sam2_image)
+            cropped_sam2_pixel_values = torch.from_numpy(cropped_sam2_image).permute(2, 0, 1).contiguous()
+            cropped_sam2_pixel_values = cropped_sam2_pixel_values.unsqueeze(0).to(vq_sam2.dtype).to(vq_sam2.device)
+        else:
+            cropped_sam2_image = np.array(resized_crop_image)
+            cropped_sam2_image = sam2_image_processor.apply_image(cropped_sam2_image)
+            cropped_sam2_pixel_values = torch.from_numpy(cropped_sam2_image).permute(2, 0, 1).contiguous()
+            cropped_sam2_pixel_values = cropped_sam2_pixel_values.unsqueeze(0).to(vq_sam2.dtype).to(vq_sam2.device)
+
+        cropped_masks = torch.stack([torch.from_numpy(np.ascontiguousarray(mu_state["cur_mask"].copy()[y1:y2, x1:x2]))])
+        assert cropped_masks.shape[-2] == crop_height and cropped_masks.shape[-1] == crop_width
+
+        if resized_crop_image is not None:
+            resized_crop_masks = torch.nn.functional.interpolate(cropped_masks.unsqueeze(0), size=(new_height, new_width), mode='bilinear')
+            resized_crop_masks = resized_crop_masks[0] > 0.5
+            cropped_masks = resized_crop_masks
+        crop_height, crop_width = cropped_masks.shape[-2:]
+        cropped_boxes = torchvision.ops.masks_to_boxes(cropped_masks)
+        crop_whwh = torch.as_tensor([[crop_width, crop_height, crop_width, crop_height]])
+        cropped_boxes = cropped_boxes / crop_whwh
+        cropped_boxes = cropped_boxes.to(vq_sam2.device)
+        cropped_masks = [m.unsqueeze(0).to(vq_sam2.device) for m in cropped_masks]
+
+        with torch.no_grad():
+            cropped_vq_sam2_output = vq_sam2(
+                cropped_sam2_pixel_values,
+                cropped_masks,
+                cropped_boxes,
+                reconstruct_mask=True,
+            )
+        
+        crop_quant_codes = cropped_vq_sam2_output.quant_codes.squeeze().detach().cpu().numpy().astype(np.int32).tolist()
+        remap_crop_quant_codes = [depth_idx*CODEBOOK_SIZE+quant_code for depth_idx, quant_code in enumerate(crop_quant_codes)]
+        crop_quant_codes = remap_crop_quant_codes
+        zoom_in_mask_tokens_str = MT_START_TOKEN + ''.join([MT_CONTEXT_TOKEN.format(str(code).zfill(4)) for code in crop_quant_codes]) + MT_END_TOKEN
+
+        buffer = io.BytesIO()
+        if resized_crop_image is None:
+            cropped_image.save(buffer, format='JPEG')
+        else:
+            resized_crop_image.save(buffer, format='JPEG')
+        buffer.seek(0)
+        b64 = base64.b64encode(buffer.read()).decode("utf-8")
+
+        reg["zoom_in_token_str"] = zoom_in_mask_tokens_str
+        reg["zoom_in_image"] = b64
+
+    mu_state["regions"][rid] = reg
+    choices = list(mu_state["regions"].keys())
+    return mu_state, gr.update(choices=choices, value=rid)
+
+def replace_region_all(text: str, rid: str, token_str: str) -> str:
+    pattern = re.compile(rf"(?<![A-Za-z0-9_]){re.escape(rid)}(?![A-Za-z0-9_])")
+    return pattern.sub(f"{rid} {token_str}", text)
+
+def short_tag_from_codes(code_a: int, code_b: int) -> str:
+    return f"<{code_a:04d}-{code_b:04d}>"
+
+@spaces.GPU
+def infer_understanding(mu_media, mu_query, mu_state):
+    model = get_qwen()
+
+    if not mu_media:
+        gr.Warning("Please upload an image")
+        return ""
+    if not mu_query:
+        gr.Warning("Please provide a text prompt.")
+        return ""
+    
+    raw_query = mu_query
+
+    # 1) find which regions are referenced in the ORIGINAL query
+    used = []
+    for rid in mu_state["regions"].keys():
+        if re.search(rf"(?<![A-Za-z0-9_]){re.escape(rid)}(?![A-Za-z0-9_])", raw_query):
+            used.append(rid)
+
+    # 2) replace ALL occurrences for each used rid
+    for rid in used:
+        reg = mu_state["regions"][rid]
+        token_str = reg.get("token_str")
+        if token_str:
+            mu_query = replace_region_all(mu_query, rid, token_str)
+
+    content = [{"type": "image", "image": mu_media}]
+    content.append({"type": "text", "text": mu_query})
+
+    # 3) zoom-in blocks only for used regions
+    for rid in used:
+        reg = mu_state["regions"][rid]
+        zoom_in_image = reg.get("zoom_in_image")
+        zoom_in_token_str = reg.get("zoom_in_token_str")
+        if zoom_in_image and zoom_in_token_str:
+            content.append({"type": "text", "text": f" Zoom in {rid}: "})
+            content.append({"type": "image", "image": f"data:image/jpeg;base64,{zoom_in_image}"})
+            content.append({"type": "text", "text": f", {zoom_in_token_str}."})
+    
+    messages = [{"role": "user", "content": content}]
+    inputs = processor.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+        return_dict=True, return_tensors="pt"
+    ).to(device)
+
+    generated_ids = model.generate(**inputs, max_new_tokens=1024)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    return processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+
 @spaces.GPU
 def infer_seg(media, query):
-    print("=======>>>enter infer seg")
-    global model
+    model = get_qwen()
+    vq_sam2 = load_vq_sam2()
 
     if not media:
         gr.Warning('Please upload an image')
@@ -185,15 +513,13 @@ def infer_seg(media, query):
         return_tensors="pt"
     )
 
-    model = model.to(device)
-
     inputs = inputs.to(model.device)
 
     generated_ids = model.generate(
         **inputs, 
         max_new_tokens=1024,
-        do_sample=False,
-        top_p=1.0,
+        # do_sample=False,
+        # top_p=1.0,
     )
     generated_ids_trimmed = [
         out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
@@ -201,21 +527,37 @@ def infer_seg(media, query):
     output_text = processor.batch_decode(
         generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )[0]
-
-    print("========>>>>output_text", output_text)
-    exit(0)
     
     quant_ids = extract_mt_token_ids_v1(output_text)
+    if len(quant_ids) == 0:
+        # only show model response; hide masks & download
+        answer = dict(text=output_text, entities=[])
+        return (
+            answer,
+            gr.update(value=None, visible=False),                      # hide AnnotatedImage
+            gr.update(value=None, interactive=False, visible=False),   # hide DownloadButton
+        )
+    
     if len(quant_ids) % CODEBOOK_DEPTH != 0:
         output_text = fix_mt_format_comprehensive(output_text)
         quant_ids = extract_mt_token_ids_v2(output_text)
 
+    if len(quant_ids) == 0 or (len(quant_ids) % CODEBOOK_DEPTH != 0):
+        answer = dict(text=output_text, entities=[])
+        return (
+            answer,
+            gr.update(value=None, visible=False),
+            gr.update(value=None, interactive=False, visible=False),
+        )
+
     batch_size = len(quant_ids) // CODEBOOK_DEPTH
     remap_quant_ids = []
     tags = []
+    short_tags = []
     for bs_id in range(batch_size):
         chunk_quant_ids = quant_ids[bs_id*CODEBOOK_DEPTH:(bs_id+1)*CODEBOOK_DEPTH]
         tags.append(f'<|mt_start|><|mt_{str(chunk_quant_ids[0]).zfill(4)}|><|mt_{str(chunk_quant_ids[1]).zfill(4)}|><|mt_end|>')
+        short_tags.append(short_tag_from_codes(chunk_quant_ids[0], chunk_quant_ids[1]))
         remap_chunk_quant_ids = [quant_id - book_id*CODEBOOK_SIZE for book_id, quant_id in enumerate(chunk_quant_ids)]
         code1 = remap_chunk_quant_ids[0]
         code2 = remap_chunk_quant_ids[1]
@@ -240,63 +582,194 @@ def infer_seg(media, query):
     # _pred_masks = _pred_masks[:, 0, :, :].cpu().numpy().astype(np.uint8)
     _pred_masks = _pred_masks.long().unsqueeze(2).cpu() # n, 1, 1, h, w 
 
+    tag_to_mask_idx = {}
+    tag_to_short = {}
+    for i, (tag, stag) in enumerate(zip(tags, short_tags)):
+        if tag not in tag_to_mask_idx:
+            tag_to_mask_idx[tag] = i
+            tag_to_short[tag] = stag
+    unique_tags = list(tag_to_mask_idx.keys())
+
     entities = []
-    unique_tags = list(set(tags))
-    entity_names = []
-    for i, tag in enumerate(unique_tags):
+    for tag in unique_tags:
         for m in re.finditer(re.escape(tag), output_text):
-            entities.append(dict(entity=f'Target {i + 1}', start=m.start(), end=m.end()))
-            entity_names.append(f'Target {i + 1}')
-    
+            entities.append(dict(entity=tag, start=m.start(), end=m.end()))
+
     answer = dict(text=output_text, entities=entities)
 
+    # entities = []
+    # unique_tags = list(set(tags))
+    # entity_names = []
+    # for i, tag in enumerate(unique_tags):
+    #     for m in re.finditer(re.escape(tag), output_text):
+    #         entities.append(dict(entity=f'Target {i + 1}', start=m.start(), end=m.end()))
+    #         entity_names.append(f'Target {i + 1}')
+    
+    # answer = dict(text=output_text, entities=entities)
+
     frames = torch.from_numpy(np.array(image)).unsqueeze(0)
     imgs = draw_mask(frames, _pred_masks, colors=colors)
 
     path = f"/tmp/{uuid.uuid4().hex}.png"
     iio.imwrite(path, imgs, duration=100, loop=0)
 
-    masks = media, [(m[0, 0].numpy(), entity_names[i]) for i, m in enumerate(_pred_masks)]
+    # masks_value = (media, [(m[0, 0].numpy(), entity_names[i]) for i, m in enumerate(_pred_masks)])
+    # masks_value = (
+    #     media,
+    #     [( _pred_masks[tag_to_mask_idx[tag]][0, 0].numpy(), tag ) for tag in unique_tags]
+    # )
+
+    entity_names = [f"Target {i+1}" for i in range(len(unique_tags))]
+    masks_value = (
+        media,
+        [(_pred_masks[tag_to_mask_idx[tag]][0, 0].numpy().astype(np.uint8) * 255, entity_names[i]) for i, tag in enumerate(unique_tags)]
+    )
+
+    lines = []
+    for i, tag in enumerate(unique_tags):
+        short_tag = tag_to_short[tag]
+        lines.append(f"- **{entity_names[i]}** → `{short_tag}`")
+    tag_map_text = "### Mask-Token Mapping\n" + "\n".join(lines)
 
-    return answer, masks, path
+    # dynamic color maps keyed by tag
+    dyn_color_map = {}
+    dyn_color_map_light = {}
+    for i, tag in enumerate(unique_tags):
+        c = colors[i % len(colors)]
+        dyn_color_map[tag] = f'#{int(c[0]):02x}{int(c[1]):02x}{int(c[2]):02x}'
+        dyn_color_map_light[tag] = f'#{int(c[0] * 127.5 + 127.5):02x}{int(c[1] * 127.5 + 127.5):02x}{int(c[2] * 127.5 + 127.5):02x}'
+
+    # return answer, masks, path
+    return (
+        gr.update(value=answer, color_map=dyn_color_map_light, visible=True),  # ans_1
+        gr.update(value=masks_value, visible=True),   # msk_1
+        gr.update(value=path, interactive=True, visible=True),                 # download
+        gr.update(value=tag_map_text, visible=True)
+    )
 
 
 def build_demo():
     with gr.Blocks(title=TITLE, js=JS, theme=gr.themes.Soft()) as demo:
         gr.HTML(HEADER)
 
-        # with gr.Tab('Mask Generation'):
-        download_btn_1 = gr.DownloadButton(label='📦 Download', interactive=False, render=False)
-        msk_1 = gr.AnnotatedImage(label='De-tokenized 2D masks', color_map=color_map, render=False)
-        ans_1 = gr.HighlightedText(
-            label='Model Response', color_map=color_map_light, show_inline_category=False, render=False)
-        with gr.Row():
-            with gr.Column():
-                media_1 = gr.Image(type='filepath')
-
-                sample_frames_1 = gr.Slider(1, 32, value=16, step=1, visible=False)
-
-                query_1 = gr.Textbox(label='Text Prompt', placeholder='Please segment the...', elem_id='query_1')
-
-                with gr.Row():
-                    random_btn_1 = gr.Button(value='🔮 Random', visible=False)
-
-                    reset_btn_1 = gr.ClearButton([media_1, query_1, msk_1, ans_1], value='🗑️ Reset')
-                    reset_btn_1.click(reset_seg, None, [sample_frames_1, download_btn_1])
-
-                    download_btn_1.render()
-
-                    submit_btn_1 = gr.Button(value='🚀 Submit', variant='primary', elem_id='submit_1')
-            
-            with gr.Column():
-                msk_1.render()
-                ans_1.render()
-
-        ctx_1 = submit_btn_1.click(disable_btns, None, [random_btn_1, reset_btn_1, download_btn_1, submit_btn_1])
-        ctx_1 = ctx_1.then(infer_seg, [media_1, query_1], [ans_1, msk_1, download_btn_1])
-        ctx_1.then(enable_btns, None, [random_btn_1, reset_btn_1, download_btn_1, submit_btn_1])
-        # with gr.Tab('Mask Understanding'):
-        #     pass
+        with gr.Tab('Mask Generation'):
+            download_btn_1 = gr.DownloadButton(label='📦 Download', interactive=False, render=False)
+            msk_1 = gr.AnnotatedImage(label='De-tokenized 2D masks', color_map=color_map, render=False)
+            ans_1 = gr.HighlightedText(
+                label='Model Response', color_map=color_map_light, show_inline_category=False, render=False)
+            tag_map_md = gr.Markdown(label="Mask-Token Mapping", value="", visible=False)
+            with gr.Row():
+                with gr.Column():
+                    media_1 = gr.Image(type='filepath')
+
+                    sample_frames_1 = gr.Slider(1, 32, value=16, step=1, visible=False)
+
+                    # query_1 = gr.Textbox(label='Text Prompt', placeholder='Please segment the...', elem_id='query_1')
+                    query_1 = gr.Textbox(
+                        label='Text Prompt',
+                        placeholder='Please segment the...',
+                        lines=3,
+                        max_lines=12,
+                        elem_id='query_1'
+                    )
+
+                    with gr.Row():
+                        random_btn_1 = gr.Button(value='🔮 Random', visible=False)
+
+                        reset_btn_1 = gr.ClearButton([media_1, query_1, msk_1, ans_1, tag_map_md], value='🗑️ Reset')
+                        reset_btn_1.click(reset_seg, None, [sample_frames_1, download_btn_1])
+
+                        download_btn_1.render()
+
+                        submit_btn_1 = gr.Button(value='🚀 Submit', variant='primary', elem_id='submit_1')
+                
+                with gr.Column():
+                    msk_1.render()
+                    tag_map_md
+                    ans_1.render()
+
+            ctx_1 = submit_btn_1.click(disable_btns, None, [random_btn_1, reset_btn_1, download_btn_1, submit_btn_1])
+            ctx_1 = ctx_1.then(infer_seg, [media_1, query_1], [ans_1, msk_1, download_btn_1, tag_map_md])
+            ctx_1.then(enable_btns, None, [random_btn_1, reset_btn_1, download_btn_1, submit_btn_1])
+
+            EXAMPLES = [
+                ["examples/example1.jpeg", "Locate the tissue box in this image and response with its segmentation mask."],
+                ["examples/example2.jpg", "Could you please give me a detail description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer."],
+                ["examples/example3.png", "Find all the people who are currently standing and response with segmentation masks."],
+                ["examples/example4.jpg", "Segment every instance that belongs to the following categories: person, bicycle, car, motorcycle, airplane, bus, train, truck, boat, traffic light, fire hydrant, stop sign, parking meter, bench, bird, cat, dog, horse, sheep, cow, elephant, bear, zebra, giraffe, backpack, umbrella, handbag, tie, suitcase, frisbee, skis, snowboard, sports ball, kite, baseball bat, baseball glove, skateboard, surfboard, tennis racket, bottle, wine glass, cup, fork, knife, spoon, bowl, banana, apple, sandwich, orange, broccoli, carrot, hot dog, pizza, donut, cake, chair, couch, potted plant, bed, dining table, toilet, tv, laptop, mouse, remote, keyboard, cell phone, microwave, oven, toaster, sink, refrigerator, book, clock, vase, scissors, teddy bear, hair drier, toothbrush, banner, blanket, bridge, cardboard, counter, curtain, door-stuff, floor-wood, flower, fruit, gravel, house, light, mirror-stuff, net, pillow, platform, playingfield, railroad, river, road, roof, sand, sea, shelf, snow, stairs, tent, towel, wall-brick, wall-stone, wall-tile, wall-wood, water-other, window-blind, window-other, tree-merged, fence-merged, ceiling-merged, sky-other-merged, cabinet-merged, table-merged, floor-other-merged, pavement-merged, mountain-merged, grass-merged, dirt-merged, paper-merged, food-other-merged, building-other-merged, rock-merged, wall-other-merged, rug-merged"],
+                ["examples/example5.jpg", "Generate a scene graph for this image. Identify the main objects and describe their relationships to each other."],
+                ["examples/example6.jpg", "Which person, wearing a shirt of a primary color, is positioned between the individual in athletic attire and the one in a uniform? A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>"]
+            ]
+            gr.Markdown("## Examples")
+            gr.Examples(
+                examples=EXAMPLES,
+                inputs=[media_1, query_1],
+                label="Click an example to load the image and prompt",
+            )
+        with gr.Tab("Mask Understanding"):
+            mu_state = gr.State(new_mu_state())
+            mu_point_is_pos = gr.State(True)
+
+            with gr.Row():
+                with gr.Column():
+                    mu_media = gr.Image(type="filepath", label="Upload Image")
+                    mu_click_img = gr.Image(label="Click to add points", interactive=True)
+
+                    with gr.Row():
+                        mu_pos_btn = gr.Button("Positive Point")
+                        mu_neg_btn = gr.Button("Negative Point")
+                        mu_clear_btn = gr.Button("Clear Prompts")
+                        mu_save_btn = gr.Button("Save Region")
+
+                    mu_region_dd = gr.Dropdown(label="Saved Regions", choices=[], interactive=True)
+
+                    mu_query = gr.Textbox(label="Text Prompt", lines=3, max_lines=12)
+                    mu_submit = gr.Button("Submit", variant="primary")
+
+                with gr.Column():
+                    mu_mask_preview = gr.Image(label="Current Mask")
+                    mu_answer = gr.Textbox(label="Model Response", lines=12)
+
+            mu_media.change(
+                fn=mu_on_upload_image,
+                inputs=[mu_media, mu_state],
+                outputs=[mu_state, mu_click_img, mu_mask_preview],
+            )
+
+            mu_pos_btn.click(lambda: True, None, mu_point_is_pos)
+            mu_neg_btn.click(lambda: False, None, mu_point_is_pos)
+
+            mu_click_img.select(
+                fn=mu_add_point,
+                inputs=[mu_state, mu_point_is_pos],
+                outputs=[mu_state, mu_mask_preview],
+            )
+
+            mu_clear_btn.click(mu_clear_prompts, [mu_state], [mu_state, mu_mask_preview])
+
+            mu_save_btn.click(mu_save_region, [mu_state], [mu_state, mu_region_dd])
+
+            mu_submit.click(
+                fn=infer_understanding,
+                inputs=[mu_media, mu_query, mu_state],
+                outputs=[mu_answer],
+            )
+
+            EXAMPLES_MU = [
+                ["examples/example1.jpeg"],
+                ["examples/example2.jpg"],
+                ["examples/example3.png"],
+                ["examples/example4.jpg"],
+                ["examples/example5.jpg"],
+                ["examples/example6.jpg"],
+            ]
+
+            gr.Markdown("## Examples")
+            gr.Examples(
+                examples=EXAMPLES_MU,
+                inputs=[mu_media],  # only load image
+                label="Click an example to load the image",
+            )
 
     return demo
 
@@ -304,4 +777,5 @@ if __name__ == '__main__':
     demo = build_demo()
 
     demo.queue()
-    demo.launch(server_name='0.0.0.0')
+    # demo.launch(server_name='0.0.0.0')
+    demo.launch()