Update app.py

app.py

@@ -42,34 +42,25 @@ def cleanup_memory():
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        # torch.cuda.synchronize()
 
 
 def load_model(name):
-    """Load model with proper memory management and dtype handling"""
+    """Load model with CORRECT dtype"""
     global processor, model
-
-    # Clean up existing model
+    
     cleanup_memory()
-
     model_id = MODEL_MAP[name]
-
-    # Load processor
+    
     processor = AutoImageProcessor.from_pretrained(model_id)
-
-    model = (
-        AutoModel.from_pretrained(
-            model_id,
-            torch_dtype="auto",
-        )
-        .to(DEVICE)
-        .eval()
-    )
-
-    # Get model info
+    
+    model = AutoModel.from_pretrained(
+        model_id,
+        torch_dtype="auto",  
+    ).eval()
+    
     param_count = sum(p.numel() for p in model.parameters()) / 1e9
-
-    return f"Loaded: {name} | {param_count:.1f}B params | {DEVICE.upper()}"
+    return f"Loaded: {name} | {param_count:.1f}B params | Ready"
 
 
 # Initialize default model
@@ -79,22 +70,31 @@ load_model(DEFAULT_NAME)
 @spaces.GPU(duration=60)
 def _extract_grid(img):
     """Extract feature grid from image"""
+    global model
+    
     with torch.inference_mode():
-        pv = processor(images=img, return_tensors="pt").pixel_values
-
-        if DEVICE == "cuda":
-            pv = pv.to(DEVICE)
-
+        # Move model to GPU for this call
+        model = model.to('cuda')
+        
+        # Process image and move to GPU
+        pv = processor(images=img, return_tensors="pt").pixel_values.to(model.device)
+        
+        # Run inference
         out = model(pixel_values=pv)
         last = out.last_hidden_state[0].to(torch.float32)
-
+        
+        # Extract features
         num_reg = getattr(model.config, "num_register_tokens", 0)
         p = model.config.patch_size
         _, _, Ht, Wt = pv.shape
         gh, gw = Ht // p, Wt // p
-
-        feats = last[1 + num_reg :, :].reshape(gh, gw, -1).cpu()
-
+        
+        feats = last[1 + num_reg:, :].reshape(gh, gw, -1).cpu()
+        
+        # Move model back to CPU before function exits
+        model = model.cpu()
+        torch.cuda.empty_cache()
+    
     return feats, gh, gw