Update app.py

app.py

@@ -3,9 +3,16 @@ from transformers import AutoProcessor, Blip2ForConditionalGeneration
 import torch
 from PIL import Image
 
-# Load the BLIP-2 model and processor with 8-bit loading
+# Load the BLIP-2 model and processor
 processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
-model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True)
+# Load model in int8 using bitsandbytes, and pass device_map='auto'
+model = Blip2ForConditionalGeneration.from_pretrained(
+    "Salesforce/blip2-opt-2.7b", load_in_8bit=True, device_map='auto'
+)
+
+# Set device to GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
 # Set device to GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"