Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 import transformers
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from PIL import Image
 import warnings
 
@@ -16,10 +16,17 @@ print(f"Using device: {device}")
 
 model_name = 'cognitivecomputations/dolphin-vision-72b'
 
-# create model and load it to the specified device
+# Configure 8-bit quantization
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_threshold=6.0,
+    llm_int8_has_fp16_weight=False
+)
+
+# create model and load it to the specified device with 8-bit quantization
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.float16,
+    quantization_config=quantization_config,
     device_map="auto",  # This will automatically use the GPU if available
     trust_remote_code=True
 )
@@ -50,7 +57,7 @@ def inference(prompt, image):
     print(f"Device of image_tensor: {image_tensor.device}")
 
     # generate
-    with torch.cuda.amp.autocast():
+    with torch.inference_mode():
         output_ids = model.generate(
             input_ids,
             images=image_tensor,