Update app.py

app.py

@@ -5,17 +5,22 @@ from PIL import Image
 import torch
 from transformers import DetrImageProcessor, DetrForObjectDetection
 import os
+import spaces
 
 # Load DETR model for object detection
-# We use try-except to handle potential model loading issues in constrained environments
 try:
     processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
     model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+    # Warm up the model
+    dummy_input = processor(images=np.zeros((100, 100, 3)), return_tensors="pt")
+    with torch.no_grad():
+        _ = model(**dummy_input)
 except Exception as e:
     print(f"Error loading model: {e}")
     processor = None
     model = None
 
+@spaces.GPU
 def detect_objects(image):
     """
     Detect objects in an image using DETR model
@@ -144,3 +149,4 @@ with gr.Blocks(title="Object Scanner", theme=gr.themes.Soft()) as demo:
 
 if __name__ == "__main__":
     demo.launch()
+