import reimport torchfrom transformers import AutoModelForImageTextToText, AutoProcessorfrom qwen_vl_utils import process_vision_infoย MODEL_DIR = "densereward/densereward-3frame-thinking" # or a local path to this checkpointย # The exact thinking system prompt used in training ships with the model:SYSTEM_PROMPT = open(f"{MODEL_DIR}/system_prompt.txt").read().strip()ย model = AutoModelForImageTextToText.from_pretrained( MODEL_DIR, torch_dtype=torch.bfloat16, device_map="auto")processor = AutoProcessor.from_pretrained(MODEL_DIR)ย # 3 chronological frames: oldest -> currentframes = [ "file:///path/to/frame_t-2.png", "file:///path/to/frame_t-1.png", "file:///path/to/frame_t.png",]messages = [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": [ {"type": "image", "image": frames[0]}, {"type": "image", "image": frames[1]}, {"type": "image", "image": frames[2]}, {"type": "text", "text": "put the black bowl on the plate"}, ], },]ย text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True)image_inputs, video_inputs = process_vision_info(messages)inputs = processor( text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt",).to(model.device)ย with torch.no_grad(): out = model.generate(**inputs, max_new_tokens=32, do_sample=False) # greedyย gen = out[:, inputs.input_ids.shape[1]:]raw = processor.batch_decode(gen, skip_special_tokens=True)[0].strip()print(raw)ย word = re.search(r"<think>\s*(\w+)\s*</think>", raw)reward = re.search(r"</think>\s*([01](?:\.\d+)?)", raw)print("reason:", word.group(1) if word else None, "| reward:", float(reward.group(1)) if reward else None)