Update app.py

app.py

@@ -43,32 +43,17 @@ geolocator = Nominatim(user_agent="skin-dashboard", timeout = 10)
 
 @st.cache_resource
 def load_image_model(token: str):
-    # 1) Grab a working ConvNeXt feature extractor from the official base model
-    extractor = AutoFeatureExtractor.from_pretrained(
-        "microsoft/convnext-base-224"
-    )
-
-    # 2) Build a minimal ConvNextConfig that matches your fine-tuned head
-    config = ConvNextConfig(
-        num_labels=2,
-        id2label={0: "benign", 1: "malignant"},
-        label2id={"benign": 0, "malignant": 1},
-    )
+    """
+    Load the pre-trained skin cancer classification model using PyTorch.
+    """
+    try:
+        extractor = AutoFeatureExtractor.from_pretrained("Anwarkh1/Skin_Cancer-Image_Classification")
+        model = AutoModelForImageClassification.from_pretrained("Anwarkh1/Skin_Cancer-Image_Classification")
+        return pipeline("image-classification", model=model, feature_extractor=extractor, framework="pt")
+    except Exception as e:
+        st.error(f"Error loading the model: {e}")
+        return None
 
-    # 3) Load *just the weights* from your repo, using the config override
-    model = AutoModelForImageClassification.from_pretrained(
-        MODEL_NAME,
-        config=config,
-        use_auth_token=token,
-    )
-
-    # 4) Return the pipeline
-    return pipeline(
-        "image-classification",
-        model=model,
-        feature_extractor=extractor,
-        device=0,  # or -1 for CPU
-    )
 
 @st.cache_resource
 def load_llm(token: str):