Update app.py

app.py

@@ -14,8 +14,8 @@ from transformers import AutoImageProcessor, AutoModel
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 MODEL_MAP = {
-    "DINOv3 ViT-L/16 Satellite (493M)": "facebook/dinov3-vitl16-pretrain-sat493m",
-    "DINOv3 ViT-L/16 LVD (1.7B)": "facebook/dinov3-vitl16-pretrain-lvd1689m",
+    "DINOv3 ViT-L/16 Satellite": "facebook/dinov3-vitl16-pretrain-sat493m",
+    "DINOv3 ViT-L/16 LVD (web data)": "facebook/dinov3-vitl16-pretrain-lvd1689m",
     "⚠️ DINOv3 ViT-7B/16 Satellite": "facebook/dinov3-vit7b16-pretrain-sat493m",
 }
 
@@ -49,33 +49,28 @@ def load_model(name):
     """Load model with proper memory management and dtype handling"""
     global processor, model
 
-    try:
-        # Clean up existing model
-        cleanup_memory()
+    # Clean up existing model
+    cleanup_memory()
 
-        model_id = MODEL_MAP[name]
+    model_id = MODEL_MAP[name]
 
-        # Load processor
-        processor = AutoImageProcessor.from_pretrained(model_id)
+    # Load processor
+    processor = AutoImageProcessor.from_pretrained(model_id)
 
-        # Load model with auto dtype for optimal performance
-        model = (
-            AutoModel.from_pretrained(
-                model_id,
-                torch_dtype="auto",
-            )
-            .to(DEVICE)
-            .eval()
+    model = (
+        AutoModel.from_pretrained(
+            model_id,
+            torch_dtype="auto",
         )
+        .to(DEVICE)
+        .eval()
+    )
 
-        # Get model info
-        param_count = sum(p.numel() for p in model.parameters()) / 1e9
+    # Get model info
+    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 | {DEVICE.upper()}"
 
-    except Exception as e:
-        cleanup_memory()
-        return f"Failed to load {name}: {str(e)}"
 
 
 # Initialize default model
@@ -226,4 +221,5 @@ with gr.Blocks() as demo:
             cache_examples=False,
         )
 
-demo.launch(share=False, debug=True)
+if __name__ == "__main__":
+    demo.launch(share=False, debug=True)