From bac2cf7655c4a3fd93052f60c42a72d36d2a7ef0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Tue, 3 Sep 2024 11:36:11 +0000
Subject: [PATCH] Idefics

---
 router/src/validation.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/router/src/validation.rs b/router/src/validation.rs
index 879e9bf38..09e1375c0 100644
--- a/router/src/validation.rs
+++ b/router/src/validation.rs
@@ -602,9 +602,10 @@ fn image_tokens(
     }
 }
 
-fn image_id(config: &Config) -> u32 {
+fn image_id(config: &Config, tokenizer: &Tokenizer) -> u32 {
     use Config::*;
     match config {
+        Idefics => tokenizer.token_to_id("<image>").unwrap(),
         Idefics2(idefics) => idefics.image_token_id,
         LlavaNext(llava) => llava.image_token_index,
         Paligemma(paligemma) => paligemma.image_token_index,
@@ -700,7 +701,7 @@ fn prepare_input(
             (
                 tokenizer_query,
                 input_chunks,
-                image_id(&config),
+                image_id(&config, &tokenizer),
                 image_hashes,
                 image_lens,
             )