Update loader.py

loader.py

@@ -1,90 +1,28 @@
 import torch
 import gc
 import open_clip
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor, BitsAndBytesConfig
-import os
 
 class ModelLoader:
     def __init__(self):
-        self.medgemma_model = None
-        self.medgemma_tokenizer = None
-        
-        self.biomedclip_model = None
-        self.biomedclip_preprocess = None
-        self.biomedclip_tokenizer = None
-        
-        self.maira2_model = None
-        self.maira2_processor = None
-        
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.hf_token = os.getenv("HF_TOKEN", "")
-        
-    def init_startup_models(self):
-        """Loads MedGemma and BiomedCLIP into VRAM."""
-        print("Pre-loading MedGemma 1.5 4B...")
-        try:
-            self.medgemma_tokenizer = AutoTokenizer.from_pretrained(
-                "google/medgemma-1.5-4b-it", token=self.hf_token
-            )
-            bnb_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.float16
-            )
-            self.medgemma_model = AutoModelForCausalLM.from_pretrained(
-                "google/medgemma-1.5-4b-it", 
-                token=self.hf_token,
-                quantization_config=bnb_config,
-                device_map="auto"
-            )
-        except Exception as e:
-            print(f"Failed to load MedGemma: {e}")
-            
-        print("Pre-loading BiomedCLIP...")
-        try:
-            model, preprocess, _ = open_clip.create_model_and_transforms('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
-            self.biomedclip_tokenizer = open_clip.get_tokenizer('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
-            self.biomedclip_model = model.to(self.device).eval()
-            self.biomedclip_preprocess = preprocess
-        except Exception as e:
-            print(f"Failed to load BiomedCLIP: {e}")
+        self.biomed_model = None
+        self.preprocess = None
 
-    def load_maira2_lazy(self):
-        """Lazy loads MAIRA-2 to GPU."""
-        if self.maira2_model is not None:
-            return
-            
-        print("Lazy Loading MAIRA-2...")
-        try:
-            self.maira2_processor = AutoProcessor.from_pretrained("microsoft/maira-2", token=self.hf_token, trust_remote_code=True)
-            bnb_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.float16
+    def load_biomed_clip(self):
+        """Universal Zero-Shot Auditor (BiomedCLIP)"""
+        if self.biomed_model is None:
+            print("🔄 Loading BiomedCLIP Universal Auditor...")
+            model, _, preprocess = open_clip.create_model_and_transforms(
+                'hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224'
             )
-            self.maira2_model = AutoModelForCausalLM.from_pretrained(
-                "microsoft/maira-2",
-                token=self.hf_token,
-                quantization_config=bnb_config,
-                device_map="auto",
-                trust_remote_code=True
-            )
-        except Exception as e:
-            print(f"Failed to load MAIRA-2: {e}")
+            self.biomed_model = model.to("cuda").eval()
+            self.preprocess = preprocess
+        return self.biomed_model, self.preprocess
 
     def clear_vram(self):
-        """Strictly moves MAIRA-2 out of VRAM/memory to ensure T4 capacity."""
-        if self.maira2_model is not None:
-            print("Offloading MAIRA-2 to free VRAM...")
-            del self.maira2_model
-            self.maira2_model = None
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            
-    def get_medgemma(self):
-        return self.medgemma_model, self.medgemma_tokenizer
-        
-    def get_biomedclip(self):
-        return self.biomedclip_model, self.biomedclip_preprocess, self.biomedclip_tokenizer
-        
-    def get_maira2(self):
-        return self.maira2_model, self.maira2_processor
+        """Safety flush to ensure Council stability on T4."""
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+# 👇 THIS IS THE CRUCIAL LINE THAT WAS MISSING 👇
+loader = ModelLoader()