Fix model card: align import/clone to CycleGRPO repo (projects.transformers.vq_sam2)

README.md

@@ -16,161 +16,70 @@ rewarded by how well the model can ground it back to the region it describes
 stage. It produces descriptions with **interleaved segmentation masks** for the
 corresponding parts of the answer, decoded through the SAMTok mask tokenizer.
 
+Code: [github.com/devinxzhang/CycleGRPO](https://github.com/devinxzhang/CycleGRPO)
+
 ## Quickstart
-Install `transformers` (with Qwen3-VL support) and get the SAMTok mask-decoder code,
-which provides the `projects.samtok.models` module (`VQ_SAM2`, `SAM2Config`,
-`DirectResize`) used to turn the generated mask tokens into segmentation masks:
+
+CycleGRPO-4B is a Qwen3-VL-4B that emits **SAMTok mask tokens** (`<|mt_...|>`). Plain
+text generation works with 🤗 Transformers directly; turning the mask tokens into
+segmentation masks needs the VQ-SAM2 decoder from the CycleGRPO repo
+(`projects.transformers.vq_sam2`), so clone and install it first:
 
 ```bash
 pip install "transformers>=4.57"
-# SAMTok mask tokenizer / VQ-SAM2 decoder code (provides projects.samtok.models):
-git clone https://github.com/bytedance/Sa2VA.git
-cd Sa2VA   # run your script from the repo root so `projects.samtok.models` is importable
+git clone https://github.com/devinxzhang/CycleGRPO.git
+cd CycleGRPO            # run from the repo root so `projects.transformers.vq_sam2` imports
+pip install -e .
 ```
 
-### Using 🤗 Transformers to Chat
+### Generate (text + mask tokens)
 
 ```python
-import re
 import torch
-import numpy as np
-from PIL import Image
-
-def extract_mt_token_ids_v1(text):
-    pattern = r"<\|mt_(\d{4})\|>"
-    return [int(x) for x in re.findall(pattern, text)]
-
-def extract_mt_token_ids_v2(text):
-    pattern = re.compile(r'<\|mt_start\|><\|mt_(\d{4})\|><\|mt_(\d{4})\|><\|mt_end\|>')
-    matches = pattern.findall(text)
-    ret_list = []
-    for num1, num2 in matches:
-        ret_list.append(int(num1))
-        ret_list.append(int(num2))
-    return ret_list
-
-def find_first_index(arr, value):
-    indices = np.where(arr == value)[0]
-    return indices[0] if len(indices) > 0 else -1
-
-def fix_mt_format_comprehensive(text):
-    pattern_too_many = r'(<\|mt_start\|>)(<\|mt_\d+\|>)(<\|mt_\d+\|>)(?:<\|mt_\d+\|>)+<\|mt_end\|>'
-    replacement_too_many = r'\1\2\3<|mt_end|>'
-    text = re.sub(pattern_too_many, replacement_too_many, text)
-
-    pattern_too_few_with_end = r'(<\|mt_start\|>)(<\|mt_\d+\|>)(<\|mt_end\|>)'
-    replacement_too_few = r'\1\2<|mt_9999|><|mt_end|>'
-    text = re.sub(pattern_too_few_with_end, replacement_too_few, text)
-
-    pattern_too_few_no_end = r'(<\|mt_start\|>)(<\|mt_\d+\|>)(?!<\|mt_)'
-    replacement_too_few_no_end = r'\1\2<|mt_9999|><|mt_end|>'
-    text = re.sub(pattern_too_few_no_end, replacement_too_few_no_end, text)
-    return text
-
 from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
-from projects.samtok.models import DirectResize, VQ_SAM2, VQ_SAM2Config, SAM2Config
-
-CYCLEGRPO_REPO = "XinNUS/CycleGRPO-4B"
-SAMTOK_REPO = "zhouyik/Qwen3-VL-4B-SAMTok"       # mask tokenizer / SAM2 decoder assets are reused from the base repo
 
-# build VLM (CycleGRPO weights)
+model_id = "XinNUS/CycleGRPO-4B"
 model = Qwen3VLForConditionalGeneration.from_pretrained(
-    CYCLEGRPO_REPO, torch_dtype="auto"
-).cuda().eval()
-processor = AutoProcessor.from_pretrained(CYCLEGRPO_REPO)
-
-# build SAMTok mask decoder (reused from the base model repo)
-CODEBOOK_SIZE = 256
-CODEBOOK_DEPTH = 2
-sam2_config = SAM2Config(
-    ckpt_path=f"{SAMTOK_REPO}/sam2.1_hiera_large.pt",
-)
-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(f"{SAMTOK_REPO}/mask_tokenizer_256x2.pth", map_location="cpu")
-vq_sam2.load_state_dict(state)
-sam2_image_processor = DirectResize(1024)
-
-# message
-image_path = "figs/totoro.jpg"
-question = "Could you please give me a detail description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer."
-image = Image.open(image_path).convert('RGB')
-ori_width, ori_height = image.size
-messages = [
-    {
-        "role": "user",
-        "content": [
-            {
-                "type": "image",
-                "image": image_path,
-            },
-            {"type": "text", "text": question},
-        ],
-    }
-]
-
-# VLM inference
+    model_id, dtype="auto", device_map="auto"
+).eval()
+processor = AutoProcessor.from_pretrained(model_id)
+
+messages = [{
+    "role": "user",
+    "content": [
+        {"type": "image", "image": "figs/totoro.jpg"},
+        {"type": "text", "text": "Describe the image with interleaved segmentation "
+                                 "masks for the corresponding parts of the answer."},
+    ],
+}]
 inputs = processor.apply_chat_template(
-    messages,
-    tokenize=True,
-    add_generation_prompt=True,
-    return_dict=True,
-    return_tensors="pt"
-)
-inputs = inputs.to(model.device)
-
-generated_ids = model.generate(
-    **inputs,
-    max_new_tokens=512,
-    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)
-]
-output_text = processor.batch_decode(
-    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-)
-
-# decode mask
-quant_ids = extract_mt_token_ids_v1(output_text[0])
-if len(quant_ids) % CODEBOOK_DEPTH != 0:
-    output_text = [fix_mt_format_comprehensive(output_text[0])]
-    quant_ids = extract_mt_token_ids_v2(output_text[0])
-
-batch_size = len(quant_ids) // CODEBOOK_DEPTH
-remap_quant_ids = []
-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"{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]
-    if not (code2 >= 0 and code2 < CODEBOOK_SIZE):
-        code2 = -1
-    remap_chunk_quant_ids_error_handle = [code1, code2]
-    remap_quant_ids.append(remap_chunk_quant_ids_error_handle)
-
-batch_size = len(remap_quant_ids)
-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)
-sam2_pixel_values = sam2_pixel_values.repeat(batch_size, 1, 1, 1)
-
-quant_ids = torch.LongTensor(remap_quant_ids).to(vq_sam2.device)
-
-with torch.no_grad():
-    _pred_masks = vq_sam2.forward_with_codes(sam2_pixel_values, quant_ids)
-_pred_masks = torch.nn.functional.interpolate(_pred_masks, size=(ori_height, ori_width), mode='bilinear')
-_pred_masks = _pred_masks > 0.5
-_pred_masks = _pred_masks[:, 0, :, :].cpu().numpy().astype(np.uint8)
-text_token_2d_mask_mapping = {tag: _pred_mask for tag, _pred_mask in zip(tags, _pred_masks)}
+    messages, tokenize=True, add_generation_prompt=True,
+    return_dict=True, return_tensors="pt",
+).to(model.device)
+
+out = model.generate(**inputs, max_new_tokens=512, do_sample=False)
+text = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
+print(text)   # answer text interleaved with <|mt_start|><|mt_XXXX|><|mt_YYYY|><|mt_end|> mask tokens
 ```
+
+### Decode mask tokens → segmentation masks
+
+The `<|mt_...|>` tokens are decoded to masks by the VQ-SAM2 mask tokenizer. Use the
+**reference implementation in the CycleGRPO repo** rather than re-deriving it — see
+`evaluation/groundingsuite/qwen3vl_groundingsuite_infer.py` (or
+`evaluation/dlc_bench/inference.py`), which build the decoder and run the decode loop:
+
+```python
+from projects.transformers.vq_sam2 import VQ_SAM2, VQ_SAM2Config, SAM2Config
+# Those scripts also contain the `DirectResize` preprocessor, the mt-token parsing
+# (extract_mt_token_ids / fix_mt_format), and the `VQ_SAM2.forward_with_codes(...)`
+# decode step (codebook size 256, depth 2). Reuse them directly.
+```
+
+The decoder weights — `mask_tokenizer_256x2.pth` and `sam2.1_hiera_large.pt` — come
+from the base model [Qwen3-VL-4B-SAMTok](https://huggingface.co/zhouyik/Qwen3-VL-4B-SAMTok).
+
+## License
+
+Released under Apache-2.0. Derived from Qwen3-VL-4B-SAMTok; use is also subject to the
+base model's license and terms.