Update app.py

app.py

@@ -7,6 +7,8 @@ print(torch.cuda.is_available())
 
 print(os.system('python -m bitsandbytes'))
 
+import os
+import torch
 import warnings
 warnings.filterwarnings('ignore')
 
@@ -18,18 +20,32 @@ from llava.model.builder import load_pretrained_model
 from llava.mm_utils import get_model_name_from_path
 from llava.eval.run_llava import eval_model
 
+# Check CUDA availability with error handling
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
 # Define the model path
 model_path = "Veda0718/llava-med-v1.5-mistral-7b-finetuned"
 
 # Load the model
-tokenizer, model, image_processor, context_len = load_pretrained_model(
-    model_path=model_path,
-    model_base=None,
-    model_name=get_model_name_from_path(model_path)
-)
+try:
+    tokenizer, model, image_processor, context_len = load_pretrained_model(
+        model_path=model_path,
+        model_base=None,
+        model_name=get_model_name_from_path(model_path)
+    )
+    
+    # Move model to appropriate device
+    model = model.to(device)
+except Exception as e:
+    print(f"Error loading model: {e}")
+    tokenizer, model, image_processor, context_len = None, None, None, None
 
 # Define the inference function
 def run_inference(image, question):
+    if model is None:
+        return "Model failed to load. Please check the logs."
+    
     args = type('Args', (), {
         "model_path": model_path,
         "model_base": None,