Update app.py

app.py

@@ -5,6 +5,8 @@ import torch.nn as nn
 from torchvision import transforms
 from torchvision.models import vit_b_16, ViT_B_16_Weights
 import pandas as pd
+from huggingface_hub import hf_hub_download
+from langchain_huggingface import HuggingFaceEmbeddings
 import io
 import os
 import base64
@@ -149,18 +151,13 @@ class SkinViT(nn.Module):
 #multilabel_model = torch.load("D:/DR/RAG/BestModels2703/skin_vit_fold10.pth", map_location='cpu')
 #multiclass_model = torch.load("D:/DR/RAG/BestModels2703/best_dermnet_vit.pth", map_location='cpu')
 
-multilabel_model = SkinViT(num_classes=len(multilabel_class_names))
-multiclass_model = SkinViT(num_classes=len(multiclass_class_names))
-
-multilabel_model.load_state_dict(torch.hub.load_state_dict_from_url(
-    "https://huggingface.co/santhoshraghu/DermBOT/resolve/main/skin_vit_fold10.pth",
-    map_location="cpu"
-))
-multiclass_model.load_state_dict(torch.hub.load_state_dict_from_url(
-    "https://huggingface.co/santhoshraghu/DermBOT/resolve/main/best_dermnet_vit.pth",
-    map_location="cpu"
-))
+# Download the full model from your Hugging Face repo
+multilabel_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filename="skin_vit_fold10.pth")
+multiclass_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filename="best_dermnet_vit.pth")
 
+# Load the models from local path
+multilabel_model = torch.load(multilabel_model_path, map_location="cpu")
+multiclass_model = torch.load(multiclass_model_path, map_location="cpu")
 
 multilabel_model.eval()
 multiclass_model.eval()