Update app.py

app.py

@@ -39,29 +39,27 @@ geolocator = Nominatim(user_agent="skin-dashboard", timeout = 10)
 
 @st.cache_resource
 def load_image_model(token: str):
-    # 1) Download the repo into a writable path
-    cache_dir = "/mnt/data/model_cache"
-    # snapshot_download will create this folder if it doesn't exist
-    local_dir = snapshot_download(
-        repo_id=MODEL_NAME,
-        use_auth_token=token,
-        cache_dir=cache_dir
+    # 1) load the feature extractor from the Hub as usual
+    extractor = AutoFeatureExtractor.from_pretrained(
+        MODEL_NAME,
+        use_auth_token=token
     )
 
-    # 2) Inject a valid config.json there
-    cfg = {
-      "architectures": ["ConvNextForImageClassification"],
-      "model_type": "convnext",
-      "num_labels": 2,
-      "id2label": { "0": "benign", "1": "malignant" },
-      "label2id": { "benign": 0, "malignant": 1 }
-    }
-    with open(os.path.join(local_dir, "config.json"), "w") as f:
-        json.dump(cfg, f)
+    # 2) manually create a ConvNextConfig with the right num_labels / id2label
+    config = ConvNextConfig(
+        num_labels=2,
+        id2label={0: "benign", 1: "malignant"},
+        label2id={"benign": 0, "malignant": 1}
+    )
+
+    # 3) load the weights with that config override
+    model = AutoModelForImageClassification.from_pretrained(
+        MODEL_NAME,
+        config=config,
+        use_auth_token=token
+    )
 
-    # 3) Load from the patched local snapshot
-    extractor = AutoFeatureExtractor.from_pretrained(local_dir)
-    model     = AutoModelForImageClassification.from_pretrained(local_dir)
+    # 4) build your pipeline
     return pipeline(
         "image-classification",
         model=model,