use IdeficsForVisionText2Text

app_dialogue.py

@@ -32,7 +32,7 @@ EOS_TOKENS = "</s>;User"
 import logging
 
 from accelerate.utils import get_max_memory
-from transformers import AutoTokenizer, AutoProcessor, AutoConfig, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoProcessor, AutoConfig, IdeficsForVisionText2Text
 
 
 TOKENIZER_FAST = True
@@ -66,7 +66,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",