Update app.py

app.py

@@ -3,11 +3,17 @@ from PIL import Image
 from transformers import AutoModel, AutoTokenizer
 import gradio as gr
 
+# Check if GPU is available and set the device
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 # Load the model and tokenizer
-model = AutoModel.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True,
-    attn_implementation='sdpa', torch_dtype=torch.bfloat16)
-model = model.eval().cuda()
-tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    'openbmb/MiniCPM-V',
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16 if device == 'cuda' else torch.float32
+)
+model = model.eval().to(device)
+tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V', trust_remote_code=True)
 
 # Function to process the image and question
 def predict(image, question):