Update app.py

app.py

@@ -55,13 +55,37 @@ collection_name = "ks_collection_1.5BE"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-local_embedding = HuggingFaceEmbeddings(
-    model_name="Alibaba-NLP/gte-Qwen2-1.5B-instruct",
-    model_kwargs={
-        "trust_remote_code": True,
-        "device": device
-    }
-)
+def get_safe_embedding_model():
+    model_name = "Alibaba-NLP/gte-Qwen2-1.5B-instruct"
+    
+    try:
+        print("Trying to load embedding model on CUDA...")
+        embedding = HuggingFaceEmbeddings(
+            model_name=model_name,
+            model_kwargs={
+                "trust_remote_code": True,
+                "device": "cuda"
+            }
+        )
+        print("Loaded embedding model on GPU.")
+        return embedding
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            print("CUDA OOM. Falling back to CPU.")
+        else:
+            print(" Error loading model on CUDA:", str(e))
+        print("Loading embedding model on CPU...")
+        return HuggingFaceEmbeddings(
+            model_name=model_name,
+            model_kwargs={
+                "trust_remote_code": True,
+                "device": "cpu"
+            }
+        )
+
+# Replace your old local_embedding line with this
+local_embedding = get_safe_embedding_model()
+
 print(" Qwen2-1.5B local embedding model loaded.")
 
 
@@ -178,17 +202,33 @@ class DermNetViT(nn.Module):
 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")
 
-multilabel_model = SkinViT(num_classes=len(multilabel_class_names))
-multiclass_model = DermNetViT(num_classes=len(multiclass_class_names))
-
-#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-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"))
-
+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
+
+# Load both models with fallback
+multilabel_model = load_model_with_fallback(SkinViT, multilabel_model_path, len(multilabel_class_names), "SkinViT")
+multiclass_model = load_model_with_fallback(DermNetViT, multiclass_model_path, len(multiclass_class_names), "DermNetViT")
 
 multilabel_model.eval()
 multiclass_model.eval()
 
+
 # === Session Init ===
 if "messages" not in st.session_state:
     st.session_state.messages = []