Update README.md

README.md

@@ -10,6 +10,7 @@ datasets:
 - ModalityDance/Omni-Bench
 base_model:
 - GAIR/Anole-7b-v0.1
+pipeline_tag: any-to-any
 ---
 
 # Omni-R1
@@ -30,7 +31,7 @@ from PIL import Image
 from transformers import ChameleonProcessor, ChameleonForConditionalGeneration
 
 # 1) Import & load
-model_id = "ModalityDance/Omni-R1"  # or a local checkpoint path
+model_id = "ModalityDance/Omni-R1"  # or "ModalityDance/Omni-R1-Zero"
 processor = ChameleonProcessor.from_pretrained(model_id)
 model = ChameleonForConditionalGeneration.from_pretrained(
     model_id,
@@ -39,7 +40,7 @@ model = ChameleonForConditionalGeneration.from_pretrained(
 )
 model.eval()
 
-# 2) Prepare a single input
+# 2) Prepare a single input (prompt contains <image>)
 prompt = "What is the smiling man in the image wearing? <image>"
 image = Image.open("image.png").convert("RGB")
 
@@ -51,9 +52,20 @@ inputs = processor(
     return_tensors="pt",
 ).to(model.device)
 
+# --- minimal image token preprocessing: replace <image> placeholder with image tokens ---
+input_ids = inputs["input_ids"].long()
+pixel_values = inputs["pixel_values"]
+
+placeholder_id = processor.tokenizer.encode("<image>", add_special_tokens=False)[0]
+image_tokens = model.get_image_tokens(pixel_values)  # shape: [1, N] (or compatible)
+
+mask = (input_ids == placeholder_id)
+input_ids = input_ids.clone()
+input_ids[mask] = image_tokens.reshape(-1).to(dtype=torch.long, device=input_ids.device)
+
 # 3) Call the model
 outputs = model.generate(
-    **inputs,
+    input_ids=input_ids,
     max_length=4096,
     do_sample=True,
     temperature=0.5,