Update app.py

app.py

@@ -34,13 +34,20 @@ def load_model():
     """Load and cache the model and processor"""
     try:
         model_name = "google/vit-base-patch16-224"
+        # Initialize the processor first
+        processor = ViTImageProcessor.from_pretrained(model_name)
+        
+        # Load model with specific device configuration
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         model = ViTForImageClassification.from_pretrained(
             model_name,
             num_labels=len(DAMAGE_TYPES),
             ignore_mismatched_sizes=True,
-            device_map="auto"
-        )
-        processor = ViTImageProcessor.from_pretrained(model_name)
+        ).to(device)
+        
+        # Ensure model is in evaluation mode
+        model.eval()
+        
         return model, processor
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
@@ -77,13 +84,16 @@ def preprocess_image(uploaded_file):
 def analyze_damage(image, model, processor):
     """Analyze structural damage in the image"""
     try:
+        device = next(model.parameters()).device
         with torch.no_grad():
             image = image.convert('RGB')
             inputs = processor(images=image, return_tensors="pt")
+            # Move inputs to the same device as model
+            inputs = {k: v.to(device) for k, v in inputs.items()}
             outputs = model(**inputs)
             probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
             cleanup_memory()
-            return probs
+            return probs.cpu()  # Move results back to CPU
     except RuntimeError as e:
         if "out of memory" in str(e):
             cleanup_memory()