Align config.use_frame_special_tokens with the processor (checkpoint mismatch)

app.py

@@ -37,7 +37,8 @@ SAMPLING_FPS = 2
 
 TRACK_STYLE = "video_point_track_per_frame"
 DEFAULT_PROMPT = "track all fish"
-IM_END_TOKEN_ID = 151937  # <im_end>; config.image_end_token_id
+IM_END_TOKEN_ID = 151937     # <im_end>    (config.image_end_token_id)
+FRAME_END_TOKEN_ID = 151944  # <frame_end> (config.frame_end_token_id)
 
 # html-v2 pointing format, exactly as in olmo/preprocessing/point_formatter.py
 COORD_RE = re.compile(r"<(?:points|tracks).*? coords=\"([0-9\t:;, .]+)\"/?>")
@@ -62,6 +63,17 @@ processor = AutoProcessor.from_pretrained(
 model = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID, trust_remote_code=True, dtype=torch.bfloat16
 ).to("cuda")
+
+# The released checkpoint ships a mismatch: processor_config.json has
+# use_frame_special_tokens=true (so the prompt gets <frame_start>/<frame_end>
+# around every frame, matching training — see olmo/models/molmo2/molmo2.py,
+# which asserts it), but config.json has it false, so the model counts
+# <im_end> instead and asserts "Expected 0 videos, but got 1". Align them.
+if processor.use_frame_special_tokens and not model.config.use_frame_special_tokens:
+    print("[molmo2fish] aligning config.use_frame_special_tokens -> True", flush=True)
+    model.config.use_frame_special_tokens = True
+    model.model.config.use_frame_special_tokens = True
+
 model.eval()
 print("Model ready.", flush=True)
 
@@ -193,6 +205,7 @@ def run_model(video_path: str, turns: list, max_new_tokens: int) -> str:
     print(
         f"[molmo2fish] input_ids={tuple(ids.shape)} "
         f"im_end={int((ids == IM_END_TOKEN_ID).sum())} "
+        f"frame_end={int((ids == FRAME_END_TOKEN_ID).sum())} "
         f"keys={sorted(inputs.keys())}",
         flush=True,
     )