Update app.py

app.py

@@ -1,13 +1,18 @@
 import torch
 from transformers import IdeficsForVisionText2Text, AutoProcessor
-from PIL import Image
+from peft import PeftModel, PeftConfig
 import gradio as gr
 
-model_id = "mrm8488/idefics-9b-ft-describe-diffusion-bf16"
+peft_model_id = "mrm8488/idefics-9b-ft-describe-diffusion-bf16-adapter"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-model = IdeficsForVisionText2Text.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-processor = AutoProcessor.from_pretrained(model_id)
+
+config = PeftConfig.from_pretrained(peft_model_id)
+model = IdeficsForVisionText2Text.from_pretrained(config.base_model_name_or_path, torch_dtype=torch.bfloat16)
+model = PeftModel.from_pretrained(model, peft_model_id)
+processor = AutoProcessor.from_pretrained(config.base_model_name_or_path)
+model = model.to(device)
+model.eval()
 
 def predict(prompt, image_url, max_length):
     image = processor.image_processor.fetch_images(image_url)