Testing

app.py

@@ -31,6 +31,7 @@ torch.cuda.empty_cache()
 nest_asyncio.apply()
 co = cohere.Client(st.secrets["COHERE_API_KEY"])
 
+device='cuda' if torch.cuda.is_available() else 'cpu'
 
 st.set_page_config(page_title="DermBOT", page_icon="🧬", layout="centered")
 
@@ -52,7 +53,6 @@ collection_name = "ks_collection_1.5BE"
 
 
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_safe_embedding_model():
     model_name = "Alibaba-NLP/gte-Qwen2-1.5B-instruct"
@@ -140,7 +140,8 @@ elif "Gemini" in selected_model:
     import google.generativeai as genai
     genai.configure(api_key=st.secrets["GEMINI_API_KEY"])  # Store in `.streamlit/secrets.toml`
 
-    gemini_model = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+    # gemini_model = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+    gemini_model = genai.GenerativeModel("gemini-2.5-pro-preview-05-06")
     def get_gemini_response(prompt):
         response = gemini_model.generate_content(prompt)
         return response.text
@@ -156,7 +157,7 @@ elif "All" in selected_model:
         openai_resp = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=st.secrets["OPENAI_API_KEY"]).invoke(
             [{"role": "system", "content": prompt}]).content
 
-        gemini = genai.GenerativeModel("gemini-2.5-pro-exp-03-25")
+        gemini = genai.GenerativeModel("gemini-2.5-pro-preview-05-06")
         gemini_resp = gemini.generate_content(prompt).text
 
         llama = Groq(api_key=st.secrets["GROQ_API_KEY"])
@@ -270,23 +271,23 @@ multilabel_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filenam
 multiclass_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filename="best_dermnet_vit_sd.pth")
 
 def load_model_with_fallback(model_class, weight_path, num_classes, model_name):
-    try:
-        print(f"🔍 Loading {model_name} on GPU...")
-        model = model_class(num_classes)
-        model.load_state_dict(torch.load(weight_path, map_location="cuda"))
-        model.to("cuda")
-        print(f"✅ {model_name} loaded on GPU.")
-        return model
-    except RuntimeError as e:
-        if "CUDA out of memory" in str(e):
-            print(f"⚠️ {model_name} OOM. Falling back to CPU.")
-        else:
-            print(f"❌ Error loading {model_name} on CUDA: {e}")
-        print(f"🔄 Loading {model_name} on CPU...")
-        model = model_class(num_classes)
-        model.load_state_dict(torch.load(weight_path, map_location="cpu"))
-        model.to("cpu")
-        return model
+    # try:
+    print(f"🔍 Loading {model_name} on GPU...")
+    model = model_class(num_classes)
+    model.load_state_dict(torch.load(weight_path, map_location=device))
+    model.to(device)
+    print(f"✅ {model_name} loaded on GPU.")
+    return model
+    # except RuntimeError as e:
+    #     if "CUDA out of memory" in str(e):
+    #         print(f"⚠️ {model_name} OOM. Falling back to CPU.")
+    #     else:
+    #         print(f"❌ Error loading {model_name} on CUDA: {e}")
+    #     print(f"🔄 Loading {model_name} on CPU...")
+    #     model = model_class(num_classes)
+    #     model.load_state_dict(torch.load(weight_path, map_location="cpu"))
+    #     model.to("cpu")
+    #     return model
 
 # Load both models with fallback
 multilabel_model = load_model_with_fallback(SkinViT, multilabel_model_path, len(multilabel_class_names), "SkinViT")