Update loader.py

loader.py

@@ -1,36 +1,26 @@
 import torch
-import torchxrayvision
 import gc
+import open_clip
 
 class ModelLoader:
     def __init__(self):
-        self.active_model = None
+        self.biomed_model = None
+        self.preprocess = None
 
-    def clear_memory(self):
-        """Clears the currently loaded model from memory."""
-        if self.active_model is not None:
-            del self.active_model
-            self.active_model = None
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
+    def load_biomed_clip(self):
+        """Universal Zero-Shot Auditor (BiomedCLIP)"""
+        if self.biomed_model is None:
+            print("🔄 Loading BiomedCLIP Auditor...")
+            model, _, preprocess = open_clip.create_model_and_transforms(
+                'hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224'
+            )
+            self.biomed_model = model.to("cuda").eval()
+            self.preprocess = preprocess
+        return self.biomed_model, self.preprocess
 
-    def load_model(self):
-        """Loads the X-Ray DenseNet model into memory (CPU)."""
-        if self.active_model is not None:
-            return self.active_model
+    def clear_vram(self):
+        """Safety flush to prevent OOM during live demo."""
+        gc.collect()
+        torch.cuda.empty_cache()
 
-        self.clear_memory()
-        print("Loading X-Ray DenseNet model into memory (CPU)...")
-        
-        try:
-            # We strictly use X-Ray DenseNet
-            model = torchxrayvision.models.DenseNet(weights="densenet121-res224-all")
-            model.eval()
-            self.active_model = model
-            print("X-Ray model loaded successfully.")
-            return self.active_model
-            
-        except Exception as e:
-            print(f"CRITICAL ERROR loading X-Ray model: {e}")
-            return None
+loader = ModelLoader()