Update app.py

app.py

@@ -3,7 +3,7 @@ from PIL import Image
 import torch
 import torch.nn as nn
 from torchvision import transforms
-from torchvision.models import vit_b_16, ViT_B_16_Weights
+from torchvision.models import vit_b_16, vit_l_16, ViT_B_16_Weights, ViT_L_16_Weights
 import pandas as pd
 from huggingface_hub import hf_hub_download
 from langchain_huggingface import HuggingFaceEmbeddings
@@ -31,7 +31,6 @@ nest_asyncio.apply()
 
 st.set_page_config(page_title="DermBOT", page_icon="🧬", layout="centered")
 
-#os.environ["PGVECTOR_CONNECTION_STRING"] = "postgresql+psycopg2://postgres:postgres@localhost:5432/VectorDB"
 
 # === Model Selection ===
 available_models = ["OpenAI GPT-4o", "LLaMA 3", "Gemini Pro"]
@@ -49,10 +48,6 @@ collection_name = "ks_collection_1.5BE"
 #local_embedding = SentenceTransformerEmbeddings(model=embedding_model)
 
 
-#local_embedding = HuggingFaceEmbeddings(
- #   model_name="Alibaba-NLP/gte-Qwen2-1.5B-instruct",
-  #  model_kwargs={"trust_remote_code": True, "device":"cpu"}
-#)
 
 local_embedding = HuggingFaceEmbeddings(
     model_name="Alibaba-NLP/gte-Qwen2-1.5B-instruct",
@@ -143,25 +138,33 @@ class SkinViT(nn.Module):
     def __init__(self, num_classes):
         super(SkinViT, self).__init__()
         self.model = vit_b_16(weights=ViT_B_16_Weights.DEFAULT)
-        in_features = self.model.heads[0].in_features
-        self.model.heads[0] = nn.Linear(in_features, num_classes)
+        in_features = self.model.heads.head.in_features
+        self.model.heads.head = nn.Linear(in_features, num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+class DermNetViT(nn.Module):
+    def __init__(self, num_classes):
+        super(DermNetViT, self).__init__()
+        self.model = vit_l_16(weights=ViT_L_16_Weights.DEFAULT)
+        in_features = self.model.heads.head.in_features
+        self.model.heads.head = nn.Linear(in_features, num_classes)
+
     def forward(self, x):
         return self.model(x)
 
 #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')
 
-# Download state_dicts from Hugging Face Hub
+# === Load Model State Dicts ===
 multilabel_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filename="skin_vit_fold10_sd.pth")
 multiclass_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filename="best_dermnet_vit_sd.pth")
 
-# Load models
 multilabel_model = SkinViT(num_classes=len(multilabel_class_names))
-multiclass_model = SkinViT(num_classes=len(multiclass_class_names))
-
+multiclass_model = DermNetViT(num_classes=len(multiclass_class_names))
 multilabel_model.load_state_dict(torch.load(multilabel_model_path, map_location="cpu"))
 multiclass_model.load_state_dict(torch.load(multiclass_model_path, map_location="cpu"))
-
 multilabel_model.eval()
 multiclass_model.eval()