1import torch
2from PIL import Image
3from transformers import ChameleonProcessor, ChameleonForConditionalGeneration
4
5# 1) Import & load
6model_id = "ModalityDance/Omni-R1" # or "ModalityDance/Omni-R1-Zero"
7processor = ChameleonProcessor.from_pretrained(model_id)
8model = ChameleonForConditionalGeneration.from_pretrained(
9 model_id,
10 torch_dtype=torch.bfloat16,
11 device_map="auto",
12)
13model.eval()
14
15# 2) Prepare a single input (prompt contains <image>)
16prompt = "What is the smiling man in the image wearing? <image>"
17image = Image.open("image.png").convert("RGB")
18
19inputs = processor(
20 prompt,
21 images=[image],
22 padding=False,
23 return_for_text_completion=True,
24 return_tensors="pt",
25).to(model.device)
26
27# --- minimal image token preprocessing: replace <image> placeholder with image tokens ---
28input_ids = inputs["input_ids"].long()
29pixel_values = inputs["pixel_values"]
30
31placeholder_id = processor.tokenizer.encode("<image>", add_special_tokens=False)[0]
32image_tokens = model.get_image_tokens(pixel_values) # shape: [1, N] (or compatible)
33
34mask = (input_ids == placeholder_id)
35input_ids = input_ids.clone()
36input_ids[mask] = image_tokens.reshape(-1).to(dtype=torch.long, device=input_ids.device)
37
38# 3) Call the model
39outputs = model.generate(
40 input_ids=input_ids,
41 max_length=4096,
42 do_sample=True,
43 temperature=0.5,
44 top_p=0.9,
45 pad_token_id=1,
46 multimodal_generation_mode="unrestricted",
47)
48
49# 4) Get results
50text = processor.batch_decode(outputs, skip_special_tokens=False)[0]
51print(text)
For full scripts (batch JSONL inference, interleaved decoding, and vLLM-based evaluation), please refer to the official GitHub repository:
https://github.com/ModalityDance/Omni-R1
1@misc{cheng2026omnir1unifiedgenerativeparadigm,
2 title={Omni-R1: Towards the Unified Generative Paradigm for Multimodal Reasoning},
3 author={Dongjie Cheng and Yongqi Li and Zhixin Ma and Hongru Cai and Yupeng Hu and Wenjie Wang and Liqiang Nie and Wenjie Li},
4 year={2026},
5 eprint={2601.09536},
6 archivePrefix={arXiv},
7 primaryClass={cs.AI},
8 url={https://arxiv.org/abs/2601.09536},
9}