Update app.py

app.py

@@ -9,7 +9,13 @@ import gradio as gr
 import PIL.Image
 import spaces
 import torch
-from transformers import AutoProcessor, Blip2ForConditionalGeneration
+from transformers import AutoProcessor, BitsAndBytesConfig, Blip2ForConditionalGeneration
+
+from transformers import BitsAndBytesConfig
+
+bnb_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+)
 
 DESCRIPTION = "# [BLIP-2](https://github.com/salesforce/LAVIS/tree/main/projects/blip2)"
 
@@ -27,7 +33,7 @@ assert MODEL_ID in [MODEL_ID_OPT_2_7B, MODEL_ID_OPT_6_7B, MODEL_ID_FLAN_T5_XL, M
 
 if torch.cuda.is_available():
     processor = AutoProcessor.from_pretrained(MODEL_ID)
-    model = Blip2ForConditionalGeneration.from_pretrained(MODEL_ID, device_map="auto", load_in_8bit=True)
+    model = Blip2ForConditionalGeneration.from_pretrained(MODEL_ID, device_map="auto", quantization_config=bnb_config)
 
 
 @spaces.GPU