Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoProcessor, Idefics2ForConditionalGeneration
+from transformers import AutoProcessor, Idefics2ForConditionalGeneration, AutoModelForPreTraining
 import subprocess
 import torch
 from peft import LoraConfig
@@ -16,7 +16,9 @@ bnb_config = BitsAndBytesConfig(
         bnb_4bit_quant_type="nf4",
         bnb_4bit_compute_dtype=torch.float16
     ) 
-model = Idefics2ForConditionalGeneration.from_pretrained("HuggingFaceM4/idefics2-8b",torch_dtype=torch.float16)#quantization_config=bnb_config)
+
+model = AutoModelForPreTraining.from_pretrained("HuggingFaceM4/idefics2-8b",quantization_config=bnb_config)
+# model = Idefics2ForConditionalGeneration.from_pretrained("HuggingFaceM4/idefics2-8b",quantization_config=bnb_config)
 
 # if USE_QLORA or USE_LORA:
 #     lora_config = LoraConfig(