Update app.py

app.py

@@ -10,15 +10,21 @@ from PIL import Image, ImageDraw
 import torch
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
-from transformers import AutoModel  # trust_remote_code=True
+#from transformers import AutoModel  # trust_remote_code=True
+from transformers import pipeline
+ 
+
 
-import gradio as gr
 
 # ============================
 # Config
 # ============================
-DEFAULT_MODEL_ID = "facebook/dinov3-vits16plus-pretrain-lvd1689m"
-ALT_MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
+#DEFAULT_MODEL_ID = "facebook/dinov3-vits16plus-pretrain-lvd1689m"
+#ALT_MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
+
+DEFAULT_MODEL_ID = "onnx-community/dinov3-vits16-pretrain-lvd1689m-ONNX"
+ALT_MODEL_ID = "onnx-community/dinov3-vith16-pretrain-lvd1689m-ONNX"
+
 AVAILABLE_MODELS = [DEFAULT_MODEL_ID, ALT_MODEL_ID]
 
 PATCH_SIZE = 16
@@ -45,7 +51,7 @@ _model_cache = {}
 _current_model_id = None
 model = None
 
-def load_model_from_hub(model_id: str):
+def load_model_from_hubold(model_id: str):
     print(f"Loading model '{model_id}' from HF Hub…")
     token = os.environ.get("HF_TOKEN")
     mdl = AutoModel.from_pretrained(model_id, token=token, trust_remote_code=True)
@@ -53,6 +59,24 @@ def load_model_from_hub(model_id: str):
     print(f"✅ Loaded '{model_id}' on {DEVICE}")
     return mdl
 
+    
+def load_model_from_hub(model_id: str):
+    print(f"Loading model '{model_id}' from HF Hub…")
+    token = os.environ.get("HF_TOKEN")
+    # Use pipeline instead of AutoModel
+    extractor = pipeline(
+        "image-feature-extraction",
+        model=model_id,
+        token=token,
+        trust_remote_code=True,
+        device=0 if DEVICE == "cuda" else -1,
+    )
+    print(f"✅ Loaded '{model_id}' on {DEVICE}")
+    return extractor
+
+
+
+
 def get_model(model_id: str):
     if model_id in _model_cache:
         return _model_cache[model_id]