Use chat template for tiny model prompts

app.py

@@ -308,6 +308,21 @@ def score_binary_logits(logits: torch.Tensor, tokenizer) -> tuple[float, str]:
     return prob_fake, label
 
 
+def build_prompt(processor_obj, user_text: str) -> str:
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": user_text},
+            ],
+        }
+    ]
+    if hasattr(processor_obj, "apply_chat_template"):
+        return processor_obj.apply_chat_template(messages, add_generation_prompt=True)
+    return f"<image>\n{user_text}"
+
+
 def infer_image(image: Image.Image):
     try:
         if image is None:
@@ -316,7 +331,7 @@ def infer_image(image: Image.Image):
         model_obj, processor_obj, device = load_model_and_processor()
 
         prepared_image = preprocess_image(image)
-        prompt_text = TINY_PROMPT if MODEL_MODE == "tiny" else TMOS_PROMPT
+        prompt_text = build_prompt(processor_obj, TINY_PROMPT if MODEL_MODE == "tiny" else TMOS_PROMPT)
         inputs = processor_obj(text=prompt_text, images=prepared_image, return_tensors="pt", padding=True)
         inputs = {name: tensor.to(device) for name, tensor in inputs.items()}