Update app.py

app.py

@@ -75,20 +75,37 @@ def load_vision_model():
         print("Starting to load vision model...")
         model_id = "microsoft/Phi-3.5-vision-instruct"
         print(f"Loading model from {model_id}")
+        
+        # Check for CUDA availability
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {device}")
+        
+        # Load model with potential memory optimization
         model = AutoModelForCausalLM.from_pretrained(
             model_id, 
             trust_remote_code=True, 
             torch_dtype=torch.float16, 
-            use_flash_attention_2=False
+            use_flash_attention_2=True,  # Enable if supported
+            device_map="auto",  # Automatically manage model placement
+            low_cpu_mem_usage=True
         )
         print("Model loaded successfully")
+        
         print("Loading processor...")
         processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
         print("Processor loaded successfully")
+        
         return model, processor
+    except ImportError as e:
+        print(f"Error importing required modules: {str(e)}")
+        print("Please ensure all required dependencies are installed.")
+    except RuntimeError as e:
+        print(f"Runtime error (possibly CUDA out of memory): {str(e)}")
+        print("Consider using a smaller model or enabling GPU offloading.")
     except Exception as e:
-        print(f"Detailed error in loading vision model: {str(e)}")
-        return None, None
+        print(f"Unexpected error in loading vision model: {str(e)}")
+    
+    return None, None
 
 
 # Process audio input within a GPU-decorated function