Update visualization

app_dialogue.py

@@ -32,7 +32,7 @@ EOS_TOKENS = "</s>;User"
 import logging
 
 from accelerate.utils import get_max_memory
-from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer
+from transformers import AutoConfig, IdeficsForVisionText2Text, AutoProcessor
 
 
 TOKENIZER_FAST = True
@@ -48,12 +48,13 @@ def load_processor_tokenizer_model(model_name):
         use_auth_token=os.getenv("HF_AUTH_TOKEN", True),
         truncation_side="left",
     )
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        use_fast=TOKENIZER_FAST,
-        use_auth_token=os.getenv("HF_AUTH_TOKEN", True),
-        truncation_side="left",
-    )
+    tokenizer = processor.tokenizer
+    # tokenizer = AutoTokenizer.from_pretrained(
+    #     model_name,
+    #     use_fast=TOKENIZER_FAST,
+    #     use_auth_token=os.getenv("HF_AUTH_TOKEN", True),
+    #     truncation_side="left",
+    # )
     # tokenizer.padding_side = "left" -> we don't need that, do we?
 
     config = AutoConfig.from_pretrained(model_name, use_auth_token=os.getenv("HF_AUTH_TOKEN", True))
@@ -66,7 +67,7 @@ def load_processor_tokenizer_model(model_name):
             # Decrease 2 for Pytorch overhead and 2 for the forward to be safe
             max_memory_map[key] = f"{max_memory_map[key] - 4} GiB"
 
-    model = AutoModelForCausalLM.from_pretrained(
+    model = IdeficsForVisionText2Text.from_pretrained(
         model_name,
         use_auth_token=os.getenv("HF_AUTH_TOKEN", True),
         device_map="auto",