bring back automodelforcausalLM

app_dialogue.py

@@ -36,7 +36,7 @@ from transformers import (
     AutoTokenizer,
     AutoProcessor,
     AutoConfig,
-    IdeficsForVisionText2Text,
+    AutoModelForCausalLM,
 )
 
 
@@ -73,7 +73,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 = IdeficsForVisionText2Text.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         model_name,
         use_auth_token=os.getenv("HF_AUTH_TOKEN", True),
         device_map="auto",